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