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 PetscCallCeed(ceed, 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 PetscCallCeed(ceed, 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 PetscCallCeed(ceed, 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) PetscCallCeed(ceed, CeedVectorReferenceCopy(x_ceed, &(*ctx)->x_ceed)); 119 else PetscCallCeed(ceed, CeedVectorCreate(ceed, x_size, &(*ctx)->x_ceed)); 120 121 if (y_ceed) PetscCallCeed(ceed, CeedVectorReferenceCopy(y_ceed, &(*ctx)->y_ceed)); 122 else PetscCallCeed(ceed, CeedVectorCreate(ceed, y_size, &(*ctx)->y_ceed)); 123 124 PetscCallCeed(ceed, CeedOperatorReferenceCopy(op_apply, &(*ctx)->op)); 125 PetscCallCeed(ceed, 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 if (!ctx) PetscFunctionReturn(PETSC_SUCCESS); 138 Ceed ceed = ctx->ceed; 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 PetscCallCeed(ceed, CeedVectorDestroy(&ctx->x_ceed)); 148 PetscCallCeed(ceed, CeedVectorDestroy(&ctx->y_ceed)); 149 PetscCallCeed(ceed, CeedOperatorDestroy(&ctx->op)); 150 PetscCallCeed(ceed, 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 Ceed ceed; 171 172 PetscFunctionBeginUser; 173 PetscCall(CeedVectorGetCeed(x_ceed, &ceed)); 174 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 175 PetscCallCeed(ceed, CeedVectorGetLength(x_ceed, &x_size)); 176 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", 177 X_size, x_size); 178 179 PetscCall(VecGetArrayAndMemType(X_petsc, &x, mem_type)); 180 PetscCallCeed(ceed, CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x)); 181 182 PetscFunctionReturn(PETSC_SUCCESS); 183 } 184 185 /** 186 @brief Transfer array from CeedVector to PETSc Vec 187 188 @param[in] x_ceed CeedVector 189 @param[in] mem_type PETSc MemType 190 @param[out] X_petsc PETSc Vec 191 192 @return An error code: 0 - success, otherwise - failure 193 **/ 194 PetscErrorCode VecC2P(CeedVector x_ceed, PetscMemType mem_type, Vec X_petsc) { 195 PetscScalar *x; 196 PetscInt X_size; 197 CeedSize x_size; 198 Ceed ceed; 199 200 PetscFunctionBeginUser; 201 PetscCall(CeedVectorGetCeed(x_ceed, &ceed)); 202 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 203 PetscCallCeed(ceed, CeedVectorGetLength(x_ceed, &x_size)); 204 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", 205 X_size, x_size); 206 207 PetscCallCeed(ceed, CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x)); 208 PetscCall(VecRestoreArrayAndMemType(X_petsc, &x)); 209 210 PetscFunctionReturn(PETSC_SUCCESS); 211 } 212 213 /** 214 @brief Transfer read-only array from PETSc Vec to CeedVector 215 216 @param[in] X_petsc PETSc Vec 217 @param[out] mem_type PETSc MemType 218 @param[out] x_ceed CeedVector 219 220 @return An error code: 0 - success, otherwise - failure 221 **/ 222 PetscErrorCode VecReadP2C(Vec X_petsc, PetscMemType *mem_type, CeedVector x_ceed) { 223 PetscScalar *x; 224 PetscInt X_size; 225 CeedSize x_size; 226 Ceed ceed; 227 228 PetscFunctionBeginUser; 229 PetscCall(CeedVectorGetCeed(x_ceed, &ceed)); 230 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 231 PetscCallCeed(ceed, CeedVectorGetLength(x_ceed, &x_size)); 232 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", 233 X_size, x_size); 234 235 PetscCall(VecGetArrayReadAndMemType(X_petsc, (const PetscScalar **)&x, mem_type)); 236 PetscCallCeed(ceed, CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x)); 237 238 PetscFunctionReturn(PETSC_SUCCESS); 239 } 240 241 /** 242 @brief Transfer read-only array from CeedVector to PETSc Vec 243 244 @param[in] x_ceed CeedVector 245 @param[in] mem_type PETSc MemType 246 @param[out] X_petsc PETSc Vec 247 248 @return An error code: 0 - success, otherwise - failure 249 **/ 250 PetscErrorCode VecReadC2P(CeedVector x_ceed, PetscMemType mem_type, Vec X_petsc) { 251 PetscScalar *x; 252 PetscInt X_size; 253 CeedSize x_size; 254 Ceed ceed; 255 256 PetscFunctionBeginUser; 257 PetscCall(CeedVectorGetCeed(x_ceed, &ceed)); 258 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 259 PetscCallCeed(ceed, CeedVectorGetLength(x_ceed, &x_size)); 260 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", 261 X_size, x_size); 262 263 PetscCallCeed(ceed, CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x)); 264 PetscCall(VecRestoreArrayReadAndMemType(X_petsc, (const PetscScalar **)&x)); 265 266 PetscFunctionReturn(PETSC_SUCCESS); 267 } 268 269 /** 270 @brief Copy PETSc Vec data into CeedVector 271 272 @param[in] X_petsc PETSc Vec 273 @param[out] x_ceed CeedVector 274 275 @return An error code: 0 - success, otherwise - failure 276 **/ 277 PetscErrorCode VecCopyP2C(Vec X_petsc, CeedVector x_ceed) { 278 PetscScalar *x; 279 PetscMemType mem_type; 280 PetscInt X_size; 281 CeedSize x_size; 282 Ceed ceed; 283 284 PetscFunctionBeginUser; 285 PetscCall(CeedVectorGetCeed(x_ceed, &ceed)); 286 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 287 PetscCallCeed(ceed, CeedVectorGetLength(x_ceed, &x_size)); 288 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", 289 X_size, x_size); 290 291 PetscCall(VecGetArrayReadAndMemType(X_petsc, (const PetscScalar **)&x, &mem_type)); 292 PetscCallCeed(ceed, CeedVectorSetArray(x_ceed, MemTypeP2C(mem_type), CEED_COPY_VALUES, x)); 293 PetscCall(VecRestoreArrayReadAndMemType(X_petsc, (const PetscScalar **)&x)); 294 295 PetscFunctionReturn(PETSC_SUCCESS); 296 } 297 298 //@brief Return VecType for the given DM 299 VecType DMReturnVecType(DM dm) { 300 VecType vec_type; 301 DMGetVecType(dm, &vec_type); 302 return vec_type; 303 } 304 305 /** 306 * @brief Create local PETSc Vecs for CeedOperator's active input/outputs 307 * 308 * This is primarily used for when the active input/ouput vector does not correspond to a `DM` object, and thus `DMCreateLocalVector` or 309 * `DMGetLocalVector` are not applicable. 310 * For example, if statitics are being store at quadrature points, a `DM`-created `Vec` will not have the 311 * correct size. 312 * 313 * @param[in] dm DM overwhich the Vecs would be used 314 * @param[in] op Operator to make the Vecs for 315 * @param[out] input Vec for CeedOperator active input 316 * @param[out] output Vec for CeedOperator active output 317 */ 318 PetscErrorCode CeedOperatorCreateLocalVecs(CeedOperator op, VecType vec_type, MPI_Comm comm, Vec *input, Vec *output) { 319 CeedSize input_size, output_size; 320 Ceed ceed; 321 322 PetscFunctionBeginUser; 323 PetscCall(CeedOperatorGetCeed(op, &ceed)); 324 PetscCallCeed(ceed, CeedOperatorGetActiveVectorLengths(op, &input_size, &output_size)); 325 if (input) { 326 PetscCall(VecCreate(comm, input)); 327 PetscCall(VecSetType(*input, vec_type)); 328 PetscCall(VecSetSizes(*input, input_size, input_size)); 329 } 330 if (output) { 331 PetscCall(VecCreate(comm, output)); 332 PetscCall(VecSetType(*output, vec_type)); 333 PetscCall(VecSetSizes(*output, output_size, output_size)); 334 } 335 336 PetscFunctionReturn(PETSC_SUCCESS); 337 } 338 339 /** 340 * @brief Apply FEM Operator defined by `OperatorApplyContext` to various input and output vectors 341 * 342 * @param[in] X Input global `Vec`, maybe `NULL` 343 * @param[in] X_loc Input local `Vec`, maybe `NULL` 344 * @param[in] x_ceed Input `CeedVector`, maybe `CEED_VECTOR_NONE` 345 * @param[in,out] y_ceed Output `CeedVector`, maybe `CEED_VECTOR_NONE` 346 * @param[in,out] Y_loc Output local `Vec`, maybe `NULL` 347 * @param[in,out] Y Output global `Vec`, maybe `NULL` 348 * @param[in] ctx Context for the operator apply 349 * @param[in] use_apply_add Whether to use `CeedOperatorApply` or `CeedOperatorApplyAdd` 350 */ 351 PetscErrorCode ApplyCeedOperator_Core(Vec X, Vec X_loc, CeedVector x_ceed, CeedVector y_ceed, Vec Y_loc, Vec Y, OperatorApplyContext ctx, 352 bool use_apply_add) { 353 PetscMemType x_mem_type, y_mem_type; 354 Ceed ceed = ctx->ceed; 355 356 PetscFunctionBeginUser; 357 if (X) PetscCall(DMGlobalToLocal(ctx->dm_x, X, INSERT_VALUES, X_loc)); 358 if (X_loc) PetscCall(VecReadP2C(X_loc, &x_mem_type, x_ceed)); 359 360 if (Y_loc) PetscCall(VecP2C(Y_loc, &y_mem_type, y_ceed)); 361 362 PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorApply, X, Y, 0, 0)); 363 PetscCall(PetscLogGpuTimeBegin()); 364 if (use_apply_add) PetscCallCeed(ceed, CeedOperatorApplyAdd(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE)); 365 else PetscCallCeed(ceed, CeedOperatorApply(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE)); 366 PetscCall(PetscLogGpuTimeEnd()); 367 PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorApply, X, Y, 0, 0)); 368 369 if (X_loc) PetscCall(VecReadC2P(ctx->x_ceed, x_mem_type, X_loc)); 370 371 if (Y_loc) PetscCall(VecC2P(ctx->y_ceed, y_mem_type, Y_loc)); 372 if (Y) PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, Y)); 373 374 PetscFunctionReturn(PETSC_SUCCESS); 375 }; 376 377 PetscErrorCode ApplyCeedOperatorGlobalToGlobal(Vec X, Vec Y, OperatorApplyContext ctx) { 378 Vec X_loc = ctx->X_loc, Y_loc = ctx->Y_loc; 379 380 PetscFunctionBeginUser; 381 PetscCall(VecZeroEntries(Y)); 382 383 // Get local vectors (if needed) 384 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); 385 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 386 387 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false)); 388 389 // Restore local vector (if needed) 390 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); 391 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 392 393 PetscFunctionReturn(PETSC_SUCCESS); 394 } 395 396 PetscErrorCode ApplyCeedOperatorLocalToGlobal(Vec X_loc, Vec Y, OperatorApplyContext ctx) { 397 Vec Y_loc = ctx->Y_loc; 398 399 PetscFunctionBeginUser; 400 PetscCall(VecZeroEntries(Y)); 401 402 // Get local vectors (if needed) 403 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 404 405 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false)); 406 407 // Restore local vectors (if needed) 408 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 409 410 PetscFunctionReturn(PETSC_SUCCESS); 411 } 412 413 PetscErrorCode ApplyCeedOperatorGlobalToLocal(Vec X, Vec Y_loc, OperatorApplyContext ctx) { 414 Vec X_loc = ctx->X_loc; 415 416 PetscFunctionBeginUser; 417 // Get local vectors (if needed) 418 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); 419 420 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, false)); 421 422 // Restore local vector (if needed) 423 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); 424 425 PetscFunctionReturn(PETSC_SUCCESS); 426 } 427 428 PetscErrorCode ApplyCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorApplyContext ctx) { 429 PetscFunctionBeginUser; 430 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, false)); 431 PetscFunctionReturn(PETSC_SUCCESS); 432 } 433 434 PetscErrorCode ApplyAddCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorApplyContext ctx) { 435 PetscFunctionBeginUser; 436 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, true)); 437 PetscFunctionReturn(PETSC_SUCCESS); 438 } 439 440 // ----------------------------------------------------------------------------- 441 // Wraps the libCEED operator for a MatShell 442 // ----------------------------------------------------------------------------- 443 PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { 444 OperatorApplyContext ctx; 445 PetscFunctionBeginUser; 446 447 PetscCall(MatShellGetContext(A, &ctx)); 448 PetscCall(ApplyCeedOperatorGlobalToGlobal(X, Y, ctx)); 449 450 PetscFunctionReturn(PETSC_SUCCESS); 451 }; 452 453 // ----------------------------------------------------------------------------- 454 // Returns the computed diagonal of the operator 455 // ----------------------------------------------------------------------------- 456 PetscErrorCode MatGetDiag_Ceed(Mat A, Vec D) { 457 OperatorApplyContext ctx; 458 Vec Y_loc; 459 PetscMemType mem_type; 460 PetscFunctionBeginUser; 461 462 PetscCall(MatShellGetContext(A, &ctx)); 463 Ceed ceed = ctx->ceed; 464 if (ctx->Y_loc) Y_loc = ctx->Y_loc; 465 else PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 466 467 // -- Place PETSc vector in libCEED vector 468 PetscCall(VecP2C(Y_loc, &mem_type, ctx->y_ceed)); 469 470 // -- Compute Diagonal 471 PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorAssembleDiagonal, A, D, 0, 0)); 472 PetscCall(PetscLogGpuTimeBegin()); 473 PetscCallCeed(ceed, CeedOperatorLinearAssembleDiagonal(ctx->op, ctx->y_ceed, CEED_REQUEST_IMMEDIATE)); 474 PetscCall(PetscLogGpuTimeEnd()); 475 PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorAssembleDiagonal, A, D, 0, 0)); 476 477 // -- Local-to-Global 478 PetscCall(VecC2P(ctx->y_ceed, mem_type, Y_loc)); 479 PetscCall(VecZeroEntries(D)); 480 PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, D)); 481 482 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 483 PetscFunctionReturn(PETSC_SUCCESS); 484 }; 485 486 /** 487 * @brief Create PETSc MatShell object for the corresponding OperatorApplyContext 488 * 489 * @param[in] ctx Context that does the action of the operator 490 * @param[out] mat MatShell for the operator 491 */ 492 PetscErrorCode CreateMatShell_Ceed(OperatorApplyContext ctx, Mat *mat) { 493 MPI_Comm comm_x = PetscObjectComm((PetscObject)(ctx->dm_x)); 494 MPI_Comm comm_y = PetscObjectComm((PetscObject)(ctx->dm_y)); 495 PetscInt X_loc_size, X_size, Y_size, Y_loc_size; 496 VecType X_vec_type, Y_vec_type; 497 498 PetscFunctionBeginUser; 499 PetscCheck(comm_x == comm_y, PETSC_COMM_WORLD, PETSC_ERR_ARG_NOTSAMECOMM, "Input and output DM must have the same comm"); 500 501 PetscCall(DMGetGlobalVectorInfo(ctx->dm_x, &X_loc_size, &X_size, &X_vec_type)); 502 PetscCall(DMGetGlobalVectorInfo(ctx->dm_y, &Y_loc_size, &Y_size, &Y_vec_type)); 503 504 PetscCall(MatCreateShell(comm_x, Y_loc_size, X_loc_size, Y_size, X_size, ctx, mat)); 505 PetscCall(MatShellSetContextDestroy(*mat, (PetscErrorCode(*)(void *))OperatorApplyContextDestroy)); 506 PetscCall(MatShellSetOperation(*mat, MATOP_MULT, (void (*)(void))MatMult_Ceed)); 507 PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiag_Ceed)); 508 509 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", 510 X_vec_type, Y_vec_type); 511 PetscCall(MatShellSetVecType(*mat, X_vec_type)); 512 513 PetscFunctionReturn(PETSC_SUCCESS); 514 } 515 516 // ----------------------------------------------------------------------------- 517