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(0); 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(0); 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; 67 if (X_loc) { 68 PetscCall(VecGetLocalSize(X_loc, &X_size)); 69 PetscCheck(X_size == x_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 70 "X_loc (%" PetscInt_FMT ") not correct size for CeedOperator active input size (%" CeedSize_FMT ")", X_size, x_size); 71 } 72 if (Y_loc) { 73 PetscCall(VecGetLocalSize(Y_loc, &Y_size)); 74 PetscCheck(Y_size == y_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 75 "Y_loc (%" PetscInt_FMT ") not correct size for CeedOperator active output size (%" CeedSize_FMT ")", Y_size, y_size); 76 } 77 } 78 79 PetscCall(PetscNew(ctx)); 80 81 // Copy PETSc Objects 82 if (dm_x) PetscCall(PetscObjectReference((PetscObject)dm_x)); 83 (*ctx)->dm_x = dm_x; 84 if (dm_y) PetscCall(PetscObjectReference((PetscObject)dm_y)); 85 (*ctx)->dm_y = dm_y; 86 87 if (X_loc) PetscCall(PetscObjectReference((PetscObject)X_loc)); 88 (*ctx)->X_loc = X_loc; 89 if (Y_loc) PetscCall(PetscObjectReference((PetscObject)Y_loc)); 90 (*ctx)->Y_loc = Y_loc; 91 92 // Copy libCEED objects 93 if (x_ceed) CeedVectorReferenceCopy(x_ceed, &(*ctx)->x_ceed); 94 else CeedVectorCreate(ceed, x_size, &(*ctx)->x_ceed); 95 96 if (y_ceed) CeedVectorReferenceCopy(y_ceed, &(*ctx)->y_ceed); 97 else CeedVectorCreate(ceed, y_size, &(*ctx)->y_ceed); 98 99 CeedOperatorReferenceCopy(op_apply, &(*ctx)->op); 100 CeedReferenceCopy(ceed, &(*ctx)->ceed); 101 102 PetscFunctionReturn(0); 103 } 104 105 /** 106 * @brief Destroy OperatorApplyContext struct 107 * 108 * @param[in,out] ctx Context to destroy 109 */ 110 PetscErrorCode OperatorApplyContextDestroy(OperatorApplyContext ctx) { 111 PetscFunctionBeginUser; 112 113 if (!ctx) PetscFunctionReturn(0); 114 115 // Destroy PETSc Objects 116 PetscCall(DMDestroy(&ctx->dm_x)); 117 PetscCall(DMDestroy(&ctx->dm_y)); 118 PetscCall(VecDestroy(&ctx->X_loc)); 119 PetscCall(VecDestroy(&ctx->Y_loc)); 120 121 // Destroy libCEED Objects 122 CeedVectorDestroy(&ctx->x_ceed); 123 CeedVectorDestroy(&ctx->y_ceed); 124 CeedOperatorDestroy(&ctx->op); 125 CeedDestroy(&ctx->ceed); 126 127 PetscCall(PetscFree(ctx)); 128 129 PetscFunctionReturn(0); 130 } 131 132 /** 133 @brief Transfer array from PETSc Vec to CeedVector 134 135 @param[in] X_petsc PETSc Vec 136 @param[out] mem_type PETSc MemType 137 @param[out] x_ceed CeedVector 138 139 @return An error code: 0 - success, otherwise - failure 140 **/ 141 PetscErrorCode VecP2C(Vec X_petsc, PetscMemType *mem_type, CeedVector x_ceed) { 142 PetscScalar *x; 143 PetscInt X_size; 144 CeedSize x_size; 145 146 PetscFunctionBeginUser; 147 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 148 CeedVectorGetLength(x_ceed, &x_size); 149 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", 150 X_size, x_size); 151 152 PetscCall(VecGetArrayAndMemType(X_petsc, &x, mem_type)); 153 CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x); 154 155 PetscFunctionReturn(0); 156 } 157 158 /** 159 @brief Transfer array from CeedVector to PETSc Vec 160 161 @param[in] x_ceed CeedVector 162 @param[in] mem_type PETSc MemType 163 @param[out] X_petsc PETSc Vec 164 165 @return An error code: 0 - success, otherwise - failure 166 **/ 167 PetscErrorCode VecC2P(CeedVector x_ceed, PetscMemType mem_type, Vec X_petsc) { 168 PetscScalar *x; 169 PetscInt X_size; 170 CeedSize x_size; 171 172 PetscFunctionBeginUser; 173 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 174 CeedVectorGetLength(x_ceed, &x_size); 175 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", 176 X_size, x_size); 177 178 CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x); 179 PetscCall(VecRestoreArrayAndMemType(X_petsc, &x)); 180 181 PetscFunctionReturn(0); 182 } 183 184 /** 185 @brief Transfer read-only array from PETSc Vec to CeedVector 186 187 @param[in] X_petsc PETSc Vec 188 @param[out] mem_type PETSc MemType 189 @param[out] x_ceed CeedVector 190 191 @return An error code: 0 - success, otherwise - failure 192 **/ 193 PetscErrorCode VecReadP2C(Vec X_petsc, PetscMemType *mem_type, CeedVector x_ceed) { 194 PetscScalar *x; 195 PetscInt X_size; 196 CeedSize x_size; 197 198 PetscFunctionBeginUser; 199 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 200 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 PetscCall(VecGetArrayReadAndMemType(X_petsc, (const PetscScalar **)&x, mem_type)); 205 CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x); 206 207 PetscFunctionReturn(0); 208 } 209 210 /** 211 @brief Transfer read-only array from CeedVector to PETSc Vec 212 213 @param[in] x_ceed CeedVector 214 @param[in] mem_type PETSc MemType 215 @param[out] X_petsc PETSc Vec 216 217 @return An error code: 0 - success, otherwise - failure 218 **/ 219 PetscErrorCode VecReadC2P(CeedVector x_ceed, PetscMemType mem_type, Vec X_petsc) { 220 PetscScalar *x; 221 PetscInt X_size; 222 CeedSize x_size; 223 224 PetscFunctionBeginUser; 225 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 226 CeedVectorGetLength(x_ceed, &x_size); 227 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", 228 X_size, x_size); 229 230 CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x); 231 PetscCall(VecRestoreArrayReadAndMemType(X_petsc, (const PetscScalar **)&x)); 232 233 PetscFunctionReturn(0); 234 } 235 236 /** 237 @brief Copy PETSc Vec data into CeedVector 238 239 @param[in] X_petsc PETSc Vec 240 @param[out] x_ceed CeedVector 241 242 @return An error code: 0 - success, otherwise - failure 243 **/ 244 PetscErrorCode VecCopyP2C(Vec X_petsc, CeedVector x_ceed) { 245 PetscScalar *x; 246 PetscMemType mem_type; 247 PetscInt X_size; 248 CeedSize x_size; 249 250 PetscFunctionBeginUser; 251 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 252 CeedVectorGetLength(x_ceed, &x_size); 253 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", 254 X_size, x_size); 255 256 PetscCall(VecGetArrayReadAndMemType(X_petsc, (const PetscScalar **)&x, &mem_type)); 257 CeedVectorSetArray(x_ceed, MemTypeP2C(mem_type), CEED_COPY_VALUES, x); 258 PetscCall(VecRestoreArrayReadAndMemType(X_petsc, (const PetscScalar **)&x)); 259 260 PetscFunctionReturn(0); 261 } 262 263 //@brief Return VecType for the given DM 264 VecType DMReturnVecType(DM dm) { 265 VecType vec_type; 266 DMGetVecType(dm, &vec_type); 267 return vec_type; 268 } 269 270 /** 271 * @brief Create local PETSc Vecs for CeedOperator's active input/outputs 272 * 273 * This is primarily used for when the active input/ouput vector does not correspond to a `DM` object, and thus `DMCreateLocalVector` or 274 * `DMGetLocalVector` are not applicable. 275 * For example, if statitics are being store at quadrature points, a `DM`-created `Vec` will not have the 276 * correct size. 277 * 278 * @param[in] dm DM overwhich the Vecs would be used 279 * @param[in] op Operator to make the Vecs for 280 * @param[out] input Vec for CeedOperator active input 281 * @param[out] output Vec for CeedOperator active output 282 */ 283 PetscErrorCode CeedOperatorCreateLocalVecs(CeedOperator op, VecType vec_type, MPI_Comm comm, Vec *input, Vec *output) { 284 CeedSize input_size, output_size; 285 286 PetscFunctionBeginUser; 287 CeedOperatorGetActiveVectorLengths(op, &input_size, &output_size); 288 if (input) { 289 PetscCall(VecCreate(comm, input)); 290 PetscCall(VecSetType(*input, vec_type)); 291 PetscCall(VecSetSizes(*input, input_size, input_size)); 292 } 293 if (output) { 294 PetscCall(VecCreate(comm, output)); 295 PetscCall(VecSetType(*output, vec_type)); 296 PetscCall(VecSetSizes(*output, output_size, output_size)); 297 } 298 299 PetscFunctionReturn(0); 300 } 301 302 /** 303 * @brief Apply FEM Operator defined by `OperatorApplyContext` to various input and output vectors 304 * 305 * @param X Input global `Vec`, maybe `NULL` 306 * @param X_loc Input local `Vec`, maybe `NULL` 307 * @param x_ceed Input `CeedVector`, maybe `CEED_VECTOR_NONE` 308 * @param y_ceed Output `CeedVector`, maybe `CEED_VECTOR_NONE` 309 * @param Y_loc Output local `Vec`, maybe `NULL` 310 * @param Y Output global `Vec`, maybe `NULL` 311 * @param ctx Context for the operator apply 312 * @param use_apply_add Whether to use `CeedOperatorApply` or `CeedOperatorApplyAdd` 313 */ 314 PetscErrorCode ApplyCeedOperator_Core(Vec X, Vec X_loc, CeedVector x_ceed, CeedVector y_ceed, Vec Y_loc, Vec Y, OperatorApplyContext ctx, 315 bool use_apply_add) { 316 PetscMemType x_mem_type, y_mem_type; 317 318 PetscFunctionBeginUser; 319 if (X) PetscCall(DMGlobalToLocal(ctx->dm_x, X, INSERT_VALUES, X_loc)); 320 if (X_loc) PetscCall(VecReadP2C(X_loc, &x_mem_type, x_ceed)); 321 322 if (Y_loc) PetscCall(VecP2C(Y_loc, &y_mem_type, y_ceed)); 323 324 if (use_apply_add) CeedOperatorApplyAdd(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE); 325 else CeedOperatorApply(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE); 326 327 if (X_loc) PetscCall(VecReadC2P(ctx->x_ceed, x_mem_type, X_loc)); 328 329 if (Y_loc) PetscCall(VecC2P(ctx->y_ceed, y_mem_type, Y_loc)); 330 if (Y) PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, Y)); 331 332 PetscFunctionReturn(0); 333 }; 334 335 PetscErrorCode ApplyCeedOperatorGlobalToGlobal(Vec X, Vec Y, OperatorApplyContext ctx) { 336 Vec X_loc = ctx->X_loc, Y_loc = ctx->Y_loc; 337 338 PetscFunctionBeginUser; 339 PetscCall(VecZeroEntries(Y)); 340 341 // Get local vectors (if needed) 342 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); 343 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 344 345 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false)); 346 347 // Restore local vector (if needed) 348 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); 349 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 350 351 PetscFunctionReturn(0); 352 } 353 354 PetscErrorCode ApplyCeedOperatorLocalToGlobal(Vec X_loc, Vec Y, OperatorApplyContext ctx) { 355 Vec Y_loc = ctx->Y_loc; 356 357 PetscFunctionBeginUser; 358 PetscCall(VecZeroEntries(Y)); 359 360 // Get local vectors (if needed) 361 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 362 363 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false)); 364 365 // Restore local vectors (if needed) 366 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 367 368 PetscFunctionReturn(0); 369 } 370 371 PetscErrorCode ApplyCeedOperatorGlobalToLocal(Vec X, Vec Y_loc, OperatorApplyContext ctx) { 372 Vec X_loc = ctx->X_loc; 373 374 PetscFunctionBeginUser; 375 // Get local vectors (if needed) 376 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); 377 378 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, false)); 379 380 // Restore local vector (if needed) 381 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); 382 383 PetscFunctionReturn(0); 384 } 385 386 PetscErrorCode ApplyCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorApplyContext ctx) { 387 PetscFunctionBeginUser; 388 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, false)); 389 PetscFunctionReturn(0); 390 } 391 392 PetscErrorCode ApplyAddCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorApplyContext ctx) { 393 PetscFunctionBeginUser; 394 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, true)); 395 PetscFunctionReturn(0); 396 } 397 398 // ----------------------------------------------------------------------------- 399 // Wraps the libCEED operator for a MatShell 400 // ----------------------------------------------------------------------------- 401 PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { 402 OperatorApplyContext ctx; 403 PetscFunctionBeginUser; 404 405 PetscCall(MatShellGetContext(A, &ctx)); 406 PetscCall(ApplyCeedOperatorGlobalToGlobal(X, Y, ctx)); 407 408 PetscFunctionReturn(0); 409 }; 410 411 // ----------------------------------------------------------------------------- 412 // Returns the computed diagonal of the operator 413 // ----------------------------------------------------------------------------- 414 PetscErrorCode MatGetDiag_Ceed(Mat A, Vec D) { 415 OperatorApplyContext ctx; 416 Vec Y_loc; 417 PetscMemType mem_type; 418 PetscFunctionBeginUser; 419 420 PetscCall(MatShellGetContext(A, &ctx)); 421 if (ctx->Y_loc) Y_loc = ctx->Y_loc; 422 else PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 423 424 // -- Place PETSc vector in libCEED vector 425 PetscCall(VecP2C(Y_loc, &mem_type, ctx->y_ceed)); 426 427 // -- Compute Diagonal 428 CeedOperatorLinearAssembleDiagonal(ctx->op, ctx->y_ceed, CEED_REQUEST_IMMEDIATE); 429 430 // -- Local-to-Global 431 PetscCall(VecC2P(ctx->y_ceed, mem_type, Y_loc)); 432 PetscCall(VecZeroEntries(D)); 433 PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, D)); 434 435 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 436 PetscFunctionReturn(0); 437 }; 438 439 /** 440 * @brief Create PETSc MatShell object for the corresponding OperatorApplyContext 441 * 442 * @param[in] ctx Context that does the action of the operator 443 * @param[out] mat MatShell for the operator 444 */ 445 PetscErrorCode CreateMatShell_Ceed(OperatorApplyContext ctx, Mat *mat) { 446 MPI_Comm comm_x = PetscObjectComm((PetscObject)(ctx->dm_x)); 447 MPI_Comm comm_y = PetscObjectComm((PetscObject)(ctx->dm_y)); 448 PetscInt X_loc_size, X_size, Y_size, Y_loc_size; 449 VecType X_vec_type, Y_vec_type; 450 451 PetscFunctionBeginUser; 452 PetscCheck(comm_x == comm_y, PETSC_COMM_WORLD, PETSC_ERR_ARG_NOTSAMECOMM, "Input and output DM must have the same comm"); 453 454 PetscCall(DMGetGlobalVectorInfo(ctx->dm_x, &X_loc_size, &X_size, &X_vec_type)); 455 PetscCall(DMGetGlobalVectorInfo(ctx->dm_y, &Y_loc_size, &Y_size, &Y_vec_type)); 456 457 PetscCall(MatCreateShell(comm_x, Y_loc_size, X_loc_size, Y_size, X_size, ctx, mat)); 458 PetscCall(MatShellSetContextDestroy(*mat, (PetscErrorCode(*)(void *))OperatorApplyContextDestroy)); 459 PetscCall(MatShellSetOperation(*mat, MATOP_MULT, (void (*)(void))MatMult_Ceed)); 460 PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiag_Ceed)); 461 462 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", 463 X_vec_type, Y_vec_type); 464 PetscCall(MatShellSetVecType(*mat, X_vec_type)); 465 466 PetscFunctionReturn(0); 467 } 468 469 // ----------------------------------------------------------------------------- 470