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 /** 16 * @brief Create OperatorApplyContext struct for applying FEM operator in a PETSc context 17 * 18 * All passed in objects are reference copied and may be destroyed if desired (with the exception of `CEED_VECTOR_NONE`). 19 * Resulting context should be destroyed with `OperatorApplyContextDestroy()`. 20 * 21 * @param[in] dm_x `DM` associated with the operator active input. May be `NULL` 22 * @param[in] dm_y `DM` associated with the operator active output. May be `NULL` 23 * @param[in] ceed `Ceed` object 24 * @param[in] op_apply `CeedOperator` representing the local action of the FEM operator 25 * @param[in] x_ceed `CeedVector` for operator active input. May be `CEED_VECTOR_NONE` or `NULL`. If `NULL`, `CeedVector` will be automatically 26 * generated. 27 * @param[in] y_ceed `CeedVector` for operator active output. May be `CEED_VECTOR_NONE` or `NULL`. If `NULL`, `CeedVector` will be automatically 28 * generated. 29 * @param[in] X_loc Local `Vec` for operator active input. If `NULL`, vector will be obtained if needed at ApplyCeedOperator time. 30 * @param[in] Y_loc Local `Vec` for operator active output. If `NULL`, vector will be obtained if needed at ApplyCeedOperator time. 31 * @param[out] ctx Struct containing all data necessary for applying the operator 32 */ 33 PetscErrorCode OperatorApplyContextCreate(DM dm_x, DM dm_y, Ceed ceed, CeedOperator op_apply, CeedVector x_ceed, CeedVector y_ceed, Vec X_loc, 34 Vec Y_loc, OperatorApplyContext *ctx) { 35 CeedSize x_size, y_size; 36 37 PetscFunctionBeginUser; 38 CeedOperatorGetActiveVectorLengths(op_apply, &x_size, &y_size); 39 { // Verify sizes 40 PetscInt X_size, Y_size; 41 if (X_loc) { 42 PetscCall(VecGetLocalSize(X_loc, &X_size)); 43 PetscCheck(X_size == x_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 44 "X_loc (%" PetscInt_FMT ") not correct size for CeedOperator active input size (%" CeedSize_FMT ")", X_size, x_size); 45 } 46 if (Y_loc) { 47 PetscCall(VecGetLocalSize(Y_loc, &Y_size)); 48 PetscCheck(Y_size == y_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, 49 "Y_loc (%" PetscInt_FMT ") not correct size for CeedOperator active output size (%" CeedSize_FMT ")", Y_size, y_size); 50 } 51 } 52 53 PetscCall(PetscNew(ctx)); 54 55 // Copy PETSc Objects 56 if (dm_x) PetscCall(PetscObjectReference((PetscObject)dm_x)); 57 (*ctx)->dm_x = dm_x; 58 if (dm_y) PetscCall(PetscObjectReference((PetscObject)dm_y)); 59 (*ctx)->dm_y = dm_y; 60 61 if (X_loc) PetscCall(PetscObjectReference((PetscObject)X_loc)); 62 (*ctx)->X_loc = X_loc; 63 if (Y_loc) PetscCall(PetscObjectReference((PetscObject)Y_loc)); 64 (*ctx)->Y_loc = Y_loc; 65 66 // Copy libCEED objects 67 if (x_ceed) CeedVectorReferenceCopy(x_ceed, &(*ctx)->x_ceed); 68 else CeedVectorCreate(ceed, x_size, &(*ctx)->x_ceed); 69 70 if (y_ceed) CeedVectorReferenceCopy(y_ceed, &(*ctx)->y_ceed); 71 else CeedVectorCreate(ceed, y_size, &(*ctx)->y_ceed); 72 73 CeedOperatorReferenceCopy(op_apply, &(*ctx)->op); 74 CeedReferenceCopy(ceed, &(*ctx)->ceed); 75 76 PetscFunctionReturn(0); 77 } 78 79 /** 80 * @brief Destroy OperatorApplyContext struct 81 * 82 * @param[in,out] ctx Context to destroy 83 */ 84 PetscErrorCode OperatorApplyContextDestroy(OperatorApplyContext ctx) { 85 PetscFunctionBeginUser; 86 87 if (!ctx) PetscFunctionReturn(0); 88 89 // Destroy PETSc Objects 90 PetscCall(DMDestroy(&ctx->dm_x)); 91 PetscCall(DMDestroy(&ctx->dm_y)); 92 PetscCall(VecDestroy(&ctx->X_loc)); 93 PetscCall(VecDestroy(&ctx->Y_loc)); 94 95 // Destroy libCEED Objects 96 CeedVectorDestroy(&ctx->x_ceed); 97 CeedVectorDestroy(&ctx->y_ceed); 98 CeedOperatorDestroy(&ctx->op); 99 CeedDestroy(&ctx->ceed); 100 101 PetscCall(PetscFree(ctx)); 102 103 PetscFunctionReturn(0); 104 } 105 106 /** 107 @brief Transfer array from PETSc Vec to CeedVector 108 109 @param[in] X_petsc PETSc Vec 110 @param[out] mem_type PETSc MemType 111 @param[out] x_ceed CeedVector 112 113 @return An error code: 0 - success, otherwise - failure 114 **/ 115 PetscErrorCode VecP2C(Vec X_petsc, PetscMemType *mem_type, CeedVector x_ceed) { 116 PetscScalar *x; 117 PetscInt X_size; 118 CeedSize x_size; 119 120 PetscFunctionBeginUser; 121 PetscCall(VecGetLocalSize(X_petsc, &X_size)); 122 CeedVectorGetLength(x_ceed, &x_size); 123 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", 124 X_size, x_size); 125 126 PetscCall(VecGetArrayAndMemType(X_petsc, &x, mem_type)); 127 CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x); 128 129 PetscFunctionReturn(0); 130 } 131 132 /** 133 @brief Transfer array from CeedVector to PETSc Vec 134 135 @param[in] x_ceed CeedVector 136 @param[in] mem_type PETSc MemType 137 @param[out] X_petsc PETSc Vec 138 139 @return An error code: 0 - success, otherwise - failure 140 **/ 141 PetscErrorCode VecC2P(CeedVector x_ceed, PetscMemType mem_type, Vec X_petsc) { 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 CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x); 153 PetscCall(VecRestoreArrayAndMemType(X_petsc, &x)); 154 155 PetscFunctionReturn(0); 156 } 157 158 /** 159 @brief Transfer read-only array from PETSc Vec to CeedVector 160 161 @param[in] X_petsc PETSc Vec 162 @param[out] mem_type PETSc MemType 163 @param[out] x_ceed CeedVector 164 165 @return An error code: 0 - success, otherwise - failure 166 **/ 167 PetscErrorCode VecReadP2C(Vec X_petsc, PetscMemType *mem_type, CeedVector x_ceed) { 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 PetscCall(VecGetArrayReadAndMemType(X_petsc, (const PetscScalar **)&x, mem_type)); 179 CeedVectorSetArray(x_ceed, MemTypeP2C(*mem_type), CEED_USE_POINTER, x); 180 181 PetscFunctionReturn(0); 182 } 183 184 /** 185 @brief Transfer read-only array from CeedVector to PETSc Vec 186 187 @param[in] x_ceed CeedVector 188 @param[in] mem_type PETSc MemType 189 @param[out] X_petsc PETSc Vec 190 191 @return An error code: 0 - success, otherwise - failure 192 **/ 193 PetscErrorCode VecReadC2P(CeedVector x_ceed, PetscMemType mem_type, Vec X_petsc) { 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 CeedVectorTakeArray(x_ceed, MemTypeP2C(mem_type), &x); 205 PetscCall(VecRestoreArrayReadAndMemType(X_petsc, (const PetscScalar **)&x)); 206 207 PetscFunctionReturn(0); 208 } 209 210 /** 211 @brief Copy PETSc Vec data into CeedVector 212 213 @param[in] X_petsc PETSc Vec 214 @param[out] x_ceed CeedVector 215 216 @return An error code: 0 - success, otherwise - failure 217 **/ 218 PetscErrorCode VecCopyP2C(Vec X_petsc, CeedVector x_ceed) { 219 PetscScalar *x; 220 PetscMemType mem_type; 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 PetscCall(VecGetArrayReadAndMemType(X_petsc, (const PetscScalar **)&x, &mem_type)); 231 CeedVectorSetArray(x_ceed, MemTypeP2C(mem_type), CEED_COPY_VALUES, x); 232 PetscCall(VecRestoreArrayReadAndMemType(X_petsc, (const PetscScalar **)&x)); 233 234 PetscFunctionReturn(0); 235 } 236 237 //@brief Return VecType for the given DM 238 VecType DMReturnVecType(DM dm) { 239 VecType vec_type; 240 DMGetVecType(dm, &vec_type); 241 return vec_type; 242 } 243 244 /** 245 * @brief Create local PETSc Vecs for CeedOperator's active input/outputs 246 * 247 * This is primarily used for when the active input/ouput vector does not correspond to a `DM` object, and thus `DMCreateLocalVector` or 248 * `DMGetLocalVector` are not applicable. 249 * For example, if statitics are being store at quadrature points, a `DM`-created `Vec` will not have the 250 * correct size. 251 * 252 * @param[in] dm DM overwhich the Vecs would be used 253 * @param[in] op Operator to make the Vecs for 254 * @param[out] input Vec for CeedOperator active input 255 * @param[out] output Vec for CeedOperator active output 256 */ 257 PetscErrorCode CeedOperatorCreateLocalVecs(CeedOperator op, VecType vec_type, MPI_Comm comm, Vec *input, Vec *output) { 258 CeedSize input_size, output_size; 259 260 PetscFunctionBeginUser; 261 CeedOperatorGetActiveVectorLengths(op, &input_size, &output_size); 262 if (input) { 263 PetscCall(VecCreate(comm, input)); 264 PetscCall(VecSetType(*input, vec_type)); 265 PetscCall(VecSetSizes(*input, input_size, input_size)); 266 } 267 if (output) { 268 PetscCall(VecCreate(comm, output)); 269 PetscCall(VecSetType(*output, vec_type)); 270 PetscCall(VecSetSizes(*output, output_size, output_size)); 271 } 272 273 PetscFunctionReturn(0); 274 } 275 276 /** 277 * @brief Apply FEM Operator defined by `OperatorApplyContext` to various input and output vectors 278 * 279 * @param X Input global `Vec`, maybe `NULL` 280 * @param X_loc Input local `Vec`, maybe `NULL` 281 * @param x_ceed Input `CeedVector`, maybe `CEED_VECTOR_NONE` 282 * @param y_ceed Output `CeedVector`, maybe `CEED_VECTOR_NONE` 283 * @param Y_loc Output local `Vec`, maybe `NULL` 284 * @param Y Output global `Vec`, maybe `NULL` 285 * @param ctx Context for the operator apply 286 * @param use_apply_add Whether to use `CeedOperatorApply` or `CeedOperatorApplyAdd` 287 */ 288 PetscErrorCode ApplyCeedOperator_Core(Vec X, Vec X_loc, CeedVector x_ceed, CeedVector y_ceed, Vec Y_loc, Vec Y, OperatorApplyContext ctx, 289 bool use_apply_add) { 290 PetscMemType x_mem_type, y_mem_type; 291 292 PetscFunctionBeginUser; 293 if (X) PetscCall(DMGlobalToLocal(ctx->dm_x, X, INSERT_VALUES, X_loc)); 294 if (X_loc) PetscCall(VecReadP2C(X_loc, &x_mem_type, x_ceed)); 295 296 if (Y_loc) PetscCall(VecP2C(Y_loc, &y_mem_type, y_ceed)); 297 298 if (use_apply_add) CeedOperatorApplyAdd(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE); 299 else CeedOperatorApply(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE); 300 301 if (X_loc) PetscCall(VecReadC2P(ctx->x_ceed, x_mem_type, X_loc)); 302 303 if (Y_loc) PetscCall(VecC2P(ctx->y_ceed, y_mem_type, Y_loc)); 304 if (Y) PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, Y)); 305 306 PetscFunctionReturn(0); 307 }; 308 309 PetscErrorCode ApplyCeedOperatorGlobalToGlobal(Vec X, Vec Y, OperatorApplyContext ctx) { 310 Vec X_loc = ctx->X_loc, Y_loc = ctx->Y_loc; 311 312 PetscFunctionBeginUser; 313 PetscCall(VecZeroEntries(Y)); 314 315 // Get local vectors (if needed) 316 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); 317 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 318 319 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false)); 320 321 // Restore local vector (if needed) 322 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); 323 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 324 325 PetscFunctionReturn(0); 326 } 327 328 PetscErrorCode ApplyCeedOperatorLocalToGlobal(Vec X_loc, Vec Y, OperatorApplyContext ctx) { 329 Vec Y_loc = ctx->Y_loc; 330 331 PetscFunctionBeginUser; 332 PetscCall(VecZeroEntries(Y)); 333 334 // Get local vectors (if needed) 335 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 336 337 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false)); 338 339 // Restore local vectors (if needed) 340 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 341 342 PetscFunctionReturn(0); 343 } 344 345 PetscErrorCode ApplyCeedOperatorGlobalToLocal(Vec X, Vec Y_loc, OperatorApplyContext ctx) { 346 Vec X_loc = ctx->X_loc; 347 348 PetscFunctionBeginUser; 349 // Get local vectors (if needed) 350 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); 351 352 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, false)); 353 354 // Restore local vector (if needed) 355 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); 356 357 PetscFunctionReturn(0); 358 } 359 360 PetscErrorCode ApplyAddCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorApplyContext ctx) { 361 PetscFunctionBeginUser; 362 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, true)); 363 PetscFunctionReturn(0); 364 } 365 366 // ----------------------------------------------------------------------------- 367 // Wraps the libCEED operator for a MatShell 368 // ----------------------------------------------------------------------------- 369 PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { 370 OperatorApplyContext ctx; 371 PetscFunctionBeginUser; 372 373 PetscCall(MatShellGetContext(A, &ctx)); 374 PetscCall(ApplyCeedOperatorGlobalToGlobal(X, Y, ctx)); 375 376 PetscFunctionReturn(0); 377 }; 378 379 // ----------------------------------------------------------------------------- 380 // Returns the computed diagonal of the operator 381 // ----------------------------------------------------------------------------- 382 PetscErrorCode MatGetDiag_Ceed(Mat A, Vec D) { 383 OperatorApplyContext ctx; 384 Vec Y_loc; 385 PetscMemType mem_type; 386 PetscFunctionBeginUser; 387 388 PetscCall(MatShellGetContext(A, &ctx)); 389 if (ctx->Y_loc) Y_loc = ctx->Y_loc; 390 else PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 391 392 // -- Place PETSc vector in libCEED vector 393 PetscCall(VecP2C(Y_loc, &mem_type, ctx->y_ceed)); 394 395 // -- Compute Diagonal 396 CeedOperatorLinearAssembleDiagonal(ctx->op, ctx->y_ceed, CEED_REQUEST_IMMEDIATE); 397 398 // -- Local-to-Global 399 PetscCall(VecC2P(ctx->y_ceed, mem_type, Y_loc)); 400 PetscCall(VecZeroEntries(D)); 401 PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, D)); 402 403 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 404 PetscFunctionReturn(0); 405 }; 406 407 // ----------------------------------------------------------------------------- 408