1 // Copyright (c) 2017-2023, Lawrence Livermore National Security, LLC and other CEED contributors. 2 // All Rights Reserved. See the top-level LICENSE and NOTICE files for details. 3 // 4 // SPDX-License-Identifier: BSD-2-Clause 5 // 6 // This file is part of CEED: http://github.com/ceed 7 8 #include "../include/petsc_ops.h" 9 10 #include <ceed.h> 11 #include <petscdm.h> 12 13 #include "../navierstokes.h" 14 15 // @brief Get information about a DM's local vector 16 PetscErrorCode DMGetLocalVectorInfo(DM dm, PetscInt *local_size, PetscInt *global_size, VecType *vec_type) { 17 Vec V_loc; 18 19 PetscFunctionBeginUser; 20 PetscCall(DMGetLocalVector(dm, &V_loc)); 21 if (local_size) PetscCall(VecGetLocalSize(V_loc, local_size)); 22 if (global_size) PetscCall(VecGetSize(V_loc, global_size)); 23 if (vec_type) PetscCall(VecGetType(V_loc, vec_type)); 24 PetscCall(DMRestoreLocalVector(dm, &V_loc)); 25 PetscFunctionReturn(PETSC_SUCCESS); 26 } 27 28 // @brief Get information about a DM's global vector 29 PetscErrorCode DMGetGlobalVectorInfo(DM dm, PetscInt *local_size, PetscInt *global_size, VecType *vec_type) { 30 Vec V; 31 32 PetscFunctionBeginUser; 33 PetscCall(DMGetGlobalVector(dm, &V)); 34 if (local_size) PetscCall(VecGetLocalSize(V, local_size)); 35 if (global_size) PetscCall(VecGetSize(V, global_size)); 36 if (vec_type) PetscCall(VecGetType(V, vec_type)); 37 PetscCall(DMRestoreGlobalVector(dm, &V)); 38 PetscFunctionReturn(PETSC_SUCCESS); 39 } 40 41 /** 42 * @brief Create OperatorApplyContext struct for applying FEM operator in a PETSc context 43 * 44 * All passed in objects are reference copied and may be destroyed if desired (with the exception of `CEED_VECTOR_NONE`). 45 * Resulting context should be destroyed with `OperatorApplyContextDestroy()`. 46 * 47 * @param[in] dm_x `DM` associated with the operator active input. May be `NULL` 48 * @param[in] dm_y `DM` associated with the operator active output. May be `NULL` 49 * @param[in] ceed `Ceed` object 50 * @param[in] op_apply `CeedOperator` representing the local action of the FEM operator 51 * @param[in] x_ceed `CeedVector` for operator active input. May be `CEED_VECTOR_NONE` or `NULL`. If `NULL`, `CeedVector` will be automatically 52 * generated. 53 * @param[in] y_ceed `CeedVector` for operator active output. May be `CEED_VECTOR_NONE` or `NULL`. If `NULL`, `CeedVector` will be automatically 54 * generated. 55 * @param[in] X_loc Local `Vec` for operator active input. If `NULL`, vector will be obtained if needed at ApplyCeedOperator time. 56 * @param[in] Y_loc Local `Vec` for operator active output. If `NULL`, vector will be obtained if needed at ApplyCeedOperator time. 57 * @param[out] ctx Struct containing all data necessary for applying the operator 58 */ 59 PetscErrorCode OperatorApplyContextCreate(DM dm_x, DM dm_y, Ceed ceed, CeedOperator op_apply, CeedVector x_ceed, CeedVector y_ceed, Vec X_loc, 60 Vec Y_loc, OperatorApplyContext *ctx) { 61 CeedSize x_size, y_size; 62 63 PetscFunctionBeginUser; 64 CeedOperatorGetActiveVectorLengths(op_apply, &x_size, &y_size); 65 { // Verify sizes 66 PetscInt X_size, Y_size, dm_X_size, dm_Y_size; 67 CeedSize x_ceed_size, y_ceed_size; 68 if (dm_x) PetscCall(DMGetLocalVectorInfo(dm_x, &dm_X_size, NULL, NULL)); 69 if (dm_y) PetscCall(DMGetLocalVectorInfo(dm_y, &dm_Y_size, NULL, NULL)); 70 if (X_loc) { 71 PetscCall(VecGetLocalSize(X_loc, &X_size)); 72 PetscCheck(X_size == x_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 73 "X_loc (%" PetscInt_FMT ") not correct size for CeedOperator active input size (%" CeedSize_FMT ")", X_size, x_size); 74 if (dm_x) 75 PetscCheck(X_size == dm_X_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 76 "X_loc size (%" PetscInt_FMT ") does not match dm_x local vector size (%" PetscInt_FMT ")", X_size, dm_X_size); 77 } 78 if (Y_loc) { 79 PetscCall(VecGetLocalSize(Y_loc, &Y_size)); 80 PetscCheck(Y_size == y_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 81 "Y_loc (%" PetscInt_FMT ") not correct size for CeedOperator active output size (%" CeedSize_FMT ")", Y_size, y_size); 82 if (dm_y) 83 PetscCheck(Y_size == dm_Y_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 84 "Y_loc size (%" PetscInt_FMT ") does not match dm_y local vector size (%" PetscInt_FMT ")", Y_size, dm_Y_size); 85 } 86 if (x_ceed && x_ceed != CEED_VECTOR_NONE) { 87 CeedVectorGetLength(x_ceed, &x_ceed_size); 88 PetscCheck(x_size >= 0 ? x_ceed_size == x_size : true, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 89 "x_ceed (%" CeedSize_FMT ") not correct size for CeedOperator active input size (%" CeedSize_FMT ")", x_ceed_size, x_size); 90 if (dm_x) 91 PetscCheck(x_ceed_size == dm_X_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 92 "x_ceed size (%" CeedSize_FMT ") does not match dm_x local vector size (%" PetscInt_FMT ")", x_ceed_size, dm_X_size); 93 } 94 if (y_ceed && y_ceed != CEED_VECTOR_NONE) { 95 CeedVectorGetLength(y_ceed, &y_ceed_size); 96 PetscCheck(y_ceed_size == y_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 97 "y_ceed (%" CeedSize_FMT ") not correct size for CeedOperator active input size (%" CeedSize_FMT ")", y_ceed_size, y_size); 98 if (dm_y) 99 PetscCheck(y_ceed_size == dm_Y_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 100 "y_ceed size (%" CeedSize_FMT ") does not match dm_y local vector size (%" PetscInt_FMT ")", y_ceed_size, dm_Y_size); 101 } 102 } 103 104 PetscCall(PetscNew(ctx)); 105 106 // Copy PETSc Objects 107 if (dm_x) PetscCall(PetscObjectReference((PetscObject)dm_x)); 108 (*ctx)->dm_x = dm_x; 109 if (dm_y) PetscCall(PetscObjectReference((PetscObject)dm_y)); 110 (*ctx)->dm_y = dm_y; 111 112 if (X_loc) PetscCall(PetscObjectReference((PetscObject)X_loc)); 113 (*ctx)->X_loc = X_loc; 114 if (Y_loc) PetscCall(PetscObjectReference((PetscObject)Y_loc)); 115 (*ctx)->Y_loc = Y_loc; 116 117 // Copy libCEED objects 118 if (x_ceed) CeedVectorReferenceCopy(x_ceed, &(*ctx)->x_ceed); 119 else CeedVectorCreate(ceed, x_size, &(*ctx)->x_ceed); 120 121 if (y_ceed) CeedVectorReferenceCopy(y_ceed, &(*ctx)->y_ceed); 122 else CeedVectorCreate(ceed, y_size, &(*ctx)->y_ceed); 123 124 CeedOperatorReferenceCopy(op_apply, &(*ctx)->op); 125 CeedReferenceCopy(ceed, &(*ctx)->ceed); 126 127 PetscFunctionReturn(PETSC_SUCCESS); 128 } 129 130 /** 131 * @brief Destroy OperatorApplyContext struct 132 * 133 * @param[in,out] ctx Context to destroy 134 */ 135 PetscErrorCode OperatorApplyContextDestroy(OperatorApplyContext ctx) { 136 PetscFunctionBeginUser; 137 138 if (!ctx) PetscFunctionReturn(PETSC_SUCCESS); 139 140 // Destroy PETSc Objects 141 PetscCall(DMDestroy(&ctx->dm_x)); 142 PetscCall(DMDestroy(&ctx->dm_y)); 143 PetscCall(VecDestroy(&ctx->X_loc)); 144 PetscCall(VecDestroy(&ctx->Y_loc)); 145 146 // Destroy libCEED Objects 147 CeedVectorDestroy(&ctx->x_ceed); 148 CeedVectorDestroy(&ctx->y_ceed); 149 CeedOperatorDestroy(&ctx->op); 150 CeedDestroy(&ctx->ceed); 151 152 PetscCall(PetscFree(ctx)); 153 154 PetscFunctionReturn(PETSC_SUCCESS); 155 } 156 157 /** 158 @brief Transfer array from PETSc Vec to CeedVector 159 160 @param[in] X_petsc PETSc Vec 161 @param[out] mem_type PETSc MemType 162 @param[out] x_ceed CeedVector 163 164 @return An error code: 0 - success, otherwise - failure 165 **/ 166 PetscErrorCode VecP2C(Vec X_petsc, PetscMemType *mem_type, CeedVector x_ceed) { 167 PetscScalar *x; 168 PetscInt X_size; 169 CeedSize x_size; 170 171 PetscFunctionBeginUser; 172 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 173 CeedVectorGetLength(x_ceed, &x_size); 174 PetscCheck(X_size == x_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, "X_petsc (%" PetscInt_FMT ") and x_ceed (%" CeedSize_FMT ") must be same size", 175 X_size, x_size); 176 177 PetscCall(VecGetArrayAndMemType(X_petsc, &x, mem_type)); 178 CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x); 179 180 PetscFunctionReturn(PETSC_SUCCESS); 181 } 182 183 /** 184 @brief Transfer array from CeedVector to PETSc Vec 185 186 @param[in] x_ceed CeedVector 187 @param[in] mem_type PETSc MemType 188 @param[out] X_petsc PETSc Vec 189 190 @return An error code: 0 - success, otherwise - failure 191 **/ 192 PetscErrorCode VecC2P(CeedVector x_ceed, PetscMemType mem_type, Vec X_petsc) { 193 PetscScalar *x; 194 PetscInt X_size; 195 CeedSize x_size; 196 197 PetscFunctionBeginUser; 198 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 199 CeedVectorGetLength(x_ceed, &x_size); 200 PetscCheck(X_size == x_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, "X_petsc (%" PetscInt_FMT ") and x_ceed (%" CeedSize_FMT ") must be same size", 201 X_size, x_size); 202 203 CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x); 204 PetscCall(VecRestoreArrayAndMemType(X_petsc, &x)); 205 206 PetscFunctionReturn(PETSC_SUCCESS); 207 } 208 209 /** 210 @brief Transfer read-only array from PETSc Vec to CeedVector 211 212 @param[in] X_petsc PETSc Vec 213 @param[out] mem_type PETSc MemType 214 @param[out] x_ceed CeedVector 215 216 @return An error code: 0 - success, otherwise - failure 217 **/ 218 PetscErrorCode VecReadP2C(Vec X_petsc, PetscMemType *mem_type, CeedVector x_ceed) { 219 PetscScalar *x; 220 PetscInt X_size; 221 CeedSize x_size; 222 223 PetscFunctionBeginUser; 224 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 225 CeedVectorGetLength(x_ceed, &x_size); 226 PetscCheck(X_size == x_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, "X_petsc (%" PetscInt_FMT ") and x_ceed (%" CeedSize_FMT ") must be same size", 227 X_size, x_size); 228 229 PetscCall(VecGetArrayReadAndMemType(X_petsc, (const PetscScalar **)&x, mem_type)); 230 CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x); 231 232 PetscFunctionReturn(PETSC_SUCCESS); 233 } 234 235 /** 236 @brief Transfer read-only array from CeedVector to PETSc Vec 237 238 @param[in] x_ceed CeedVector 239 @param[in] mem_type PETSc MemType 240 @param[out] X_petsc PETSc Vec 241 242 @return An error code: 0 - success, otherwise - failure 243 **/ 244 PetscErrorCode VecReadC2P(CeedVector x_ceed, PetscMemType mem_type, Vec X_petsc) { 245 PetscScalar *x; 246 PetscInt X_size; 247 CeedSize x_size; 248 249 PetscFunctionBeginUser; 250 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 251 CeedVectorGetLength(x_ceed, &x_size); 252 PetscCheck(X_size == x_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, "X_petsc (%" PetscInt_FMT ") and x_ceed (%" CeedSize_FMT ") must be same size", 253 X_size, x_size); 254 255 CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x); 256 PetscCall(VecRestoreArrayReadAndMemType(X_petsc, (const PetscScalar **)&x)); 257 258 PetscFunctionReturn(PETSC_SUCCESS); 259 } 260 261 /** 262 @brief Copy PETSc Vec data into CeedVector 263 264 @param[in] X_petsc PETSc Vec 265 @param[out] x_ceed CeedVector 266 267 @return An error code: 0 - success, otherwise - failure 268 **/ 269 PetscErrorCode VecCopyP2C(Vec X_petsc, CeedVector x_ceed) { 270 PetscScalar *x; 271 PetscMemType mem_type; 272 PetscInt X_size; 273 CeedSize x_size; 274 275 PetscFunctionBeginUser; 276 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 277 CeedVectorGetLength(x_ceed, &x_size); 278 PetscCheck(X_size == x_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, "X_petsc (%" PetscInt_FMT ") and x_ceed (%" CeedSize_FMT ") must be same size", 279 X_size, x_size); 280 281 PetscCall(VecGetArrayReadAndMemType(X_petsc, (const PetscScalar **)&x, &mem_type)); 282 CeedVectorSetArray(x_ceed, MemTypeP2C(mem_type), CEED_COPY_VALUES, x); 283 PetscCall(VecRestoreArrayReadAndMemType(X_petsc, (const PetscScalar **)&x)); 284 285 PetscFunctionReturn(PETSC_SUCCESS); 286 } 287 288 //@brief Return VecType for the given DM 289 VecType DMReturnVecType(DM dm) { 290 VecType vec_type; 291 DMGetVecType(dm, &vec_type); 292 return vec_type; 293 } 294 295 /** 296 * @brief Create local PETSc Vecs for CeedOperator's active input/outputs 297 * 298 * This is primarily used for when the active input/ouput vector does not correspond to a `DM` object, and thus `DMCreateLocalVector` or 299 * `DMGetLocalVector` are not applicable. 300 * For example, if statitics are being store at quadrature points, a `DM`-created `Vec` will not have the 301 * correct size. 302 * 303 * @param[in] dm DM overwhich the Vecs would be used 304 * @param[in] op Operator to make the Vecs for 305 * @param[out] input Vec for CeedOperator active input 306 * @param[out] output Vec for CeedOperator active output 307 */ 308 PetscErrorCode CeedOperatorCreateLocalVecs(CeedOperator op, VecType vec_type, MPI_Comm comm, Vec *input, Vec *output) { 309 CeedSize input_size, output_size; 310 311 PetscFunctionBeginUser; 312 CeedOperatorGetActiveVectorLengths(op, &input_size, &output_size); 313 if (input) { 314 PetscCall(VecCreate(comm, input)); 315 PetscCall(VecSetType(*input, vec_type)); 316 PetscCall(VecSetSizes(*input, input_size, input_size)); 317 } 318 if (output) { 319 PetscCall(VecCreate(comm, output)); 320 PetscCall(VecSetType(*output, vec_type)); 321 PetscCall(VecSetSizes(*output, output_size, output_size)); 322 } 323 324 PetscFunctionReturn(PETSC_SUCCESS); 325 } 326 327 /** 328 * @brief Apply FEM Operator defined by `OperatorApplyContext` to various input and output vectors 329 * 330 * @param[in] X Input global `Vec`, maybe `NULL` 331 * @param[in] X_loc Input local `Vec`, maybe `NULL` 332 * @param[in] x_ceed Input `CeedVector`, maybe `CEED_VECTOR_NONE` 333 * @param[in,out] y_ceed Output `CeedVector`, maybe `CEED_VECTOR_NONE` 334 * @param[in,out] Y_loc Output local `Vec`, maybe `NULL` 335 * @param[in,out] Y Output global `Vec`, maybe `NULL` 336 * @param[in] ctx Context for the operator apply 337 * @param[in] use_apply_add Whether to use `CeedOperatorApply` or `CeedOperatorApplyAdd` 338 */ 339 PetscErrorCode ApplyCeedOperator_Core(Vec X, Vec X_loc, CeedVector x_ceed, CeedVector y_ceed, Vec Y_loc, Vec Y, OperatorApplyContext ctx, 340 bool use_apply_add) { 341 PetscMemType x_mem_type, y_mem_type; 342 343 PetscFunctionBeginUser; 344 if (X) PetscCall(DMGlobalToLocal(ctx->dm_x, X, INSERT_VALUES, X_loc)); 345 if (X_loc) PetscCall(VecReadP2C(X_loc, &x_mem_type, x_ceed)); 346 347 if (Y_loc) PetscCall(VecP2C(Y_loc, &y_mem_type, y_ceed)); 348 349 PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorApply, X, Y, 0, 0)); 350 PetscCall(PetscLogGpuTimeBegin()); 351 if (use_apply_add) CeedOperatorApplyAdd(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE); 352 else CeedOperatorApply(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE); 353 PetscCall(PetscLogGpuTimeEnd()); 354 PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorApply, X, Y, 0, 0)); 355 356 if (X_loc) PetscCall(VecReadC2P(ctx->x_ceed, x_mem_type, X_loc)); 357 358 if (Y_loc) PetscCall(VecC2P(ctx->y_ceed, y_mem_type, Y_loc)); 359 if (Y) PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, Y)); 360 361 PetscFunctionReturn(PETSC_SUCCESS); 362 }; 363 364 PetscErrorCode ApplyCeedOperatorGlobalToGlobal(Vec X, Vec Y, OperatorApplyContext ctx) { 365 Vec X_loc = ctx->X_loc, Y_loc = ctx->Y_loc; 366 367 PetscFunctionBeginUser; 368 PetscCall(VecZeroEntries(Y)); 369 370 // Get local vectors (if needed) 371 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); 372 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 373 374 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false)); 375 376 // Restore local vector (if needed) 377 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); 378 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 379 380 PetscFunctionReturn(PETSC_SUCCESS); 381 } 382 383 PetscErrorCode ApplyCeedOperatorLocalToGlobal(Vec X_loc, Vec Y, OperatorApplyContext ctx) { 384 Vec Y_loc = ctx->Y_loc; 385 386 PetscFunctionBeginUser; 387 PetscCall(VecZeroEntries(Y)); 388 389 // Get local vectors (if needed) 390 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 391 392 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false)); 393 394 // Restore local vectors (if needed) 395 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 396 397 PetscFunctionReturn(PETSC_SUCCESS); 398 } 399 400 PetscErrorCode ApplyCeedOperatorGlobalToLocal(Vec X, Vec Y_loc, OperatorApplyContext ctx) { 401 Vec X_loc = ctx->X_loc; 402 403 PetscFunctionBeginUser; 404 // Get local vectors (if needed) 405 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); 406 407 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, false)); 408 409 // Restore local vector (if needed) 410 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); 411 412 PetscFunctionReturn(PETSC_SUCCESS); 413 } 414 415 PetscErrorCode ApplyCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorApplyContext ctx) { 416 PetscFunctionBeginUser; 417 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, false)); 418 PetscFunctionReturn(PETSC_SUCCESS); 419 } 420 421 PetscErrorCode ApplyAddCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorApplyContext ctx) { 422 PetscFunctionBeginUser; 423 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, true)); 424 PetscFunctionReturn(PETSC_SUCCESS); 425 } 426 427 // ----------------------------------------------------------------------------- 428 // Wraps the libCEED operator for a MatShell 429 // ----------------------------------------------------------------------------- 430 PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { 431 OperatorApplyContext ctx; 432 PetscFunctionBeginUser; 433 434 PetscCall(MatShellGetContext(A, &ctx)); 435 PetscCall(ApplyCeedOperatorGlobalToGlobal(X, Y, ctx)); 436 437 PetscFunctionReturn(PETSC_SUCCESS); 438 }; 439 440 // ----------------------------------------------------------------------------- 441 // Returns the computed diagonal of the operator 442 // ----------------------------------------------------------------------------- 443 PetscErrorCode MatGetDiag_Ceed(Mat A, Vec D) { 444 OperatorApplyContext ctx; 445 Vec Y_loc; 446 PetscMemType mem_type; 447 PetscFunctionBeginUser; 448 449 PetscCall(MatShellGetContext(A, &ctx)); 450 if (ctx->Y_loc) Y_loc = ctx->Y_loc; 451 else PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 452 453 // -- Place PETSc vector in libCEED vector 454 PetscCall(VecP2C(Y_loc, &mem_type, ctx->y_ceed)); 455 456 // -- Compute Diagonal 457 PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorAssembleDiagonal, A, D, 0, 0)); 458 PetscCall(PetscLogGpuTimeBegin()); 459 CeedOperatorLinearAssembleDiagonal(ctx->op, ctx->y_ceed, CEED_REQUEST_IMMEDIATE); 460 PetscCall(PetscLogGpuTimeEnd()); 461 PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorAssembleDiagonal, A, D, 0, 0)); 462 463 // -- Local-to-Global 464 PetscCall(VecC2P(ctx->y_ceed, mem_type, Y_loc)); 465 PetscCall(VecZeroEntries(D)); 466 PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, D)); 467 468 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 469 PetscFunctionReturn(PETSC_SUCCESS); 470 }; 471 472 /** 473 * @brief Create PETSc MatShell object for the corresponding OperatorApplyContext 474 * 475 * @param[in] ctx Context that does the action of the operator 476 * @param[out] mat MatShell for the operator 477 */ 478 PetscErrorCode CreateMatShell_Ceed(OperatorApplyContext ctx, Mat *mat) { 479 MPI_Comm comm_x = PetscObjectComm((PetscObject)(ctx->dm_x)); 480 MPI_Comm comm_y = PetscObjectComm((PetscObject)(ctx->dm_y)); 481 PetscInt X_loc_size, X_size, Y_size, Y_loc_size; 482 VecType X_vec_type, Y_vec_type; 483 484 PetscFunctionBeginUser; 485 PetscCheck(comm_x == comm_y, PETSC_COMM_WORLD, PETSC_ERR_ARG_NOTSAMECOMM, "Input and output DM must have the same comm"); 486 487 PetscCall(DMGetGlobalVectorInfo(ctx->dm_x, &X_loc_size, &X_size, &X_vec_type)); 488 PetscCall(DMGetGlobalVectorInfo(ctx->dm_y, &Y_loc_size, &Y_size, &Y_vec_type)); 489 490 PetscCall(MatCreateShell(comm_x, Y_loc_size, X_loc_size, Y_size, X_size, ctx, mat)); 491 PetscCall(MatShellSetContextDestroy(*mat, (PetscErrorCode(*)(void *))OperatorApplyContextDestroy)); 492 PetscCall(MatShellSetOperation(*mat, MATOP_MULT, (void (*)(void))MatMult_Ceed)); 493 PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiag_Ceed)); 494 495 PetscCheck(X_vec_type == Y_vec_type, PETSC_COMM_WORLD, PETSC_ERR_ARG_NOTSAMETYPE, "Vec_type of ctx->dm_x (%s) and ctx->dm_y (%s) must be the same", 496 X_vec_type, Y_vec_type); 497 PetscCall(MatShellSetVecType(*mat, X_vec_type)); 498 499 PetscFunctionReturn(PETSC_SUCCESS); 500 } 501 502 // ----------------------------------------------------------------------------- 503