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