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 /** 217 * @brief Apply FEM Operator defined by `OperatorApplyContext` to various input and output vectors 218 * 219 * @param X Input global `Vec`, maybe `NULL` 220 * @param X_loc Input local `Vec`, maybe `NULL` 221 * @param x_ceed Input `CeedVector`, maybe `CEED_VECTOR_NONE` 222 * @param y_ceed Output `CeedVector`, maybe `CEED_VECTOR_NONE` 223 * @param Y_loc Output local `Vec`, maybe `NULL` 224 * @param Y Output global `Vec`, maybe `NULL` 225 * @param ctx Context for the operator apply 226 * @param use_apply_add Whether to use `CeedOperatorApply` or `CeedOperatorApplyAdd` 227 */ 228 PetscErrorCode ApplyCeedOperator_Core(Vec X, Vec X_loc, CeedVector x_ceed, CeedVector y_ceed, Vec Y_loc, Vec Y, OperatorApplyContext ctx, 229 bool use_apply_add) { 230 PetscMemType x_mem_type, y_mem_type; 231 232 PetscFunctionBeginUser; 233 if (X) PetscCall(DMGlobalToLocal(ctx->dm_x, X, INSERT_VALUES, X_loc)); 234 if (X_loc) PetscCall(VecReadP2C(X_loc, &x_mem_type, x_ceed)); 235 236 if (Y_loc) PetscCall(VecP2C(Y_loc, &y_mem_type, y_ceed)); 237 238 if (use_apply_add) CeedOperatorApplyAdd(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE); 239 else CeedOperatorApply(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE); 240 241 if (X_loc) PetscCall(VecReadC2P(ctx->x_ceed, x_mem_type, X_loc)); 242 243 if (Y_loc) PetscCall(VecC2P(ctx->y_ceed, y_mem_type, Y_loc)); 244 if (Y) PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, Y)); 245 246 PetscFunctionReturn(0); 247 }; 248 249 PetscErrorCode ApplyCeedOperatorGlobalToGlobal(Vec X, Vec Y, OperatorApplyContext ctx) { 250 Vec X_loc = ctx->X_loc, Y_loc = ctx->Y_loc; 251 252 PetscFunctionBeginUser; 253 PetscCall(VecZeroEntries(Y)); 254 255 // Get local vectors (if needed) 256 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); 257 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 258 259 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false)); 260 261 // Restore local vector (if needed) 262 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); 263 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 264 265 PetscFunctionReturn(0); 266 } 267 268 PetscErrorCode ApplyCeedOperatorLocalToGlobal(Vec X_loc, Vec Y, OperatorApplyContext ctx) { 269 Vec Y_loc = ctx->Y_loc; 270 271 PetscFunctionBeginUser; 272 PetscCall(VecZeroEntries(Y)); 273 274 // Get local vectors (if needed) 275 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 276 277 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false)); 278 279 // Restore local vectors (if needed) 280 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 281 282 PetscFunctionReturn(0); 283 } 284 285 PetscErrorCode ApplyCeedOperatorGlobalToLocal(Vec X, Vec Y_loc, OperatorApplyContext ctx) { 286 Vec X_loc = ctx->X_loc; 287 288 PetscFunctionBeginUser; 289 // Get local vectors (if needed) 290 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc)); 291 292 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, false)); 293 294 // Restore local vector (if needed) 295 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc)); 296 297 PetscFunctionReturn(0); 298 } 299 300 PetscErrorCode ApplyAddCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorApplyContext ctx) { 301 PetscFunctionBeginUser; 302 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, true)); 303 PetscFunctionReturn(0); 304 } 305 306 // ----------------------------------------------------------------------------- 307 // Wraps the libCEED operator for a MatShell 308 // ----------------------------------------------------------------------------- 309 PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) { 310 OperatorApplyContext ctx; 311 PetscFunctionBeginUser; 312 313 PetscCall(MatShellGetContext(A, &ctx)); 314 PetscCall(ApplyCeedOperatorGlobalToGlobal(X, Y, ctx)); 315 316 PetscFunctionReturn(0); 317 }; 318 319 // ----------------------------------------------------------------------------- 320 // Returns the computed diagonal of the operator 321 // ----------------------------------------------------------------------------- 322 PetscErrorCode MatGetDiag_Ceed(Mat A, Vec D) { 323 OperatorApplyContext ctx; 324 Vec Y_loc; 325 PetscMemType mem_type; 326 PetscFunctionBeginUser; 327 328 PetscCall(MatShellGetContext(A, &ctx)); 329 if (ctx->Y_loc) Y_loc = ctx->Y_loc; 330 else PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc)); 331 332 // -- Place PETSc vector in libCEED vector 333 PetscCall(VecP2C(Y_loc, &mem_type, ctx->y_ceed)); 334 335 // -- Compute Diagonal 336 CeedOperatorLinearAssembleDiagonal(ctx->op, ctx->y_ceed, CEED_REQUEST_IMMEDIATE); 337 338 // -- Local-to-Global 339 PetscCall(VecC2P(ctx->y_ceed, mem_type, Y_loc)); 340 PetscCall(VecZeroEntries(D)); 341 PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, D)); 342 343 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc)); 344 PetscFunctionReturn(0); 345 }; 346 347 // ----------------------------------------------------------------------------- 348