1 // Copyright (c) 2017-2026, 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
DMGetLocalVectorInfo(DM dm,PetscInt * local_size,PetscInt * global_size,VecType * vec_type)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
DMGetGlobalVectorInfo(DM dm,PetscInt * local_size,PetscInt * global_size,VecType * vec_type)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 */
OperatorApplyContextCreate(DM dm_x,DM dm_y,Ceed ceed,CeedOperator op_apply,CeedVector x_ceed,CeedVector y_ceed,Vec X_loc,Vec Y_loc,OperatorApplyContext * ctx)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 }
79 if (Y_loc) {
80 PetscCall(VecGetLocalSize(Y_loc, &Y_size));
81 PetscCheck(Y_size == y_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ,
82 "Y_loc (%" PetscInt_FMT ") not correct size for CeedOperator active output size (%" CeedSize_FMT ")", Y_size, y_size);
83 if (dm_y) {
84 PetscCheck(Y_size == dm_Y_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ,
85 "Y_loc size (%" PetscInt_FMT ") does not match dm_y local vector size (%" PetscInt_FMT ")", Y_size, dm_Y_size);
86 }
87 }
88 if (x_ceed && x_ceed != CEED_VECTOR_NONE) {
89 PetscCallCeed(ceed, CeedVectorGetLength(x_ceed, &x_ceed_size));
90 PetscCheck(x_size >= 0 ? x_ceed_size == x_size : true, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ,
91 "x_ceed (%" CeedSize_FMT ") not correct size for CeedOperator active input size (%" CeedSize_FMT ")", x_ceed_size, x_size);
92 if (dm_x) {
93 PetscCheck(x_ceed_size == dm_X_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ,
94 "x_ceed size (%" CeedSize_FMT ") does not match dm_x local vector size (%" PetscInt_FMT ")", x_ceed_size, dm_X_size);
95 }
96 }
97 if (y_ceed && y_ceed != CEED_VECTOR_NONE) {
98 PetscCallCeed(ceed, CeedVectorGetLength(y_ceed, &y_ceed_size));
99 PetscCheck(y_ceed_size == y_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ,
100 "y_ceed (%" CeedSize_FMT ") not correct size for CeedOperator active input size (%" CeedSize_FMT ")", y_ceed_size, y_size);
101 if (dm_y) {
102 PetscCheck(y_ceed_size == dm_Y_size, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ,
103 "y_ceed size (%" CeedSize_FMT ") does not match dm_y local vector size (%" PetscInt_FMT ")", y_ceed_size, dm_Y_size);
104 }
105 }
106 }
107
108 PetscCall(PetscNew(ctx));
109
110 // Copy PETSc Objects
111 if (dm_x) PetscCall(PetscObjectReference((PetscObject)dm_x));
112 (*ctx)->dm_x = dm_x;
113 if (dm_y) PetscCall(PetscObjectReference((PetscObject)dm_y));
114 (*ctx)->dm_y = dm_y;
115
116 if (X_loc) PetscCall(PetscObjectReference((PetscObject)X_loc));
117 (*ctx)->X_loc = X_loc;
118 if (Y_loc) PetscCall(PetscObjectReference((PetscObject)Y_loc));
119 (*ctx)->Y_loc = Y_loc;
120
121 // Copy libCEED objects
122 if (x_ceed) PetscCallCeed(ceed, CeedVectorReferenceCopy(x_ceed, &(*ctx)->x_ceed));
123 else PetscCallCeed(ceed, CeedVectorCreate(ceed, x_size, &(*ctx)->x_ceed));
124
125 if (y_ceed) PetscCallCeed(ceed, CeedVectorReferenceCopy(y_ceed, &(*ctx)->y_ceed));
126 else PetscCallCeed(ceed, CeedVectorCreate(ceed, y_size, &(*ctx)->y_ceed));
127
128 PetscCallCeed(ceed, CeedOperatorReferenceCopy(op_apply, &(*ctx)->op));
129 PetscCallCeed(ceed, CeedReferenceCopy(ceed, &(*ctx)->ceed));
130 PetscFunctionReturn(PETSC_SUCCESS);
131 }
132
133 /**
134 * @brief Destroy OperatorApplyContext struct
135 *
136 * @param[in,out] ctx Context to destroy
137 */
OperatorApplyContextDestroy(OperatorApplyContext ctx)138 PetscErrorCode OperatorApplyContextDestroy(OperatorApplyContext ctx) {
139 PetscFunctionBeginUser;
140 if (!ctx) PetscFunctionReturn(PETSC_SUCCESS);
141 Ceed ceed = ctx->ceed;
142
143 // Destroy PETSc Objects
144 PetscCall(DMDestroy(&ctx->dm_x));
145 PetscCall(DMDestroy(&ctx->dm_y));
146 PetscCall(VecDestroy(&ctx->X_loc));
147 PetscCall(VecDestroy(&ctx->Y_loc));
148
149 // Destroy libCEED Objects
150 PetscCallCeed(ceed, CeedVectorDestroy(&ctx->x_ceed));
151 PetscCallCeed(ceed, CeedVectorDestroy(&ctx->y_ceed));
152 PetscCallCeed(ceed, CeedOperatorDestroy(&ctx->op));
153 PetscCallCeed(ceed, CeedDestroy(&ctx->ceed));
154
155 PetscCall(PetscFree(ctx));
156 PetscFunctionReturn(PETSC_SUCCESS);
157 }
158
159 //@brief Return VecType for the given DM
DMReturnVecType(DM dm)160 VecType DMReturnVecType(DM dm) {
161 VecType vec_type;
162 DMGetVecType(dm, &vec_type);
163 return vec_type;
164 }
165
166 /**
167 * @brief Create local PETSc Vecs for CeedOperator's active input/outputs
168 *
169 * This is primarily used for when the active input/ouput vector does not correspond to a `DM` object, and thus `DMCreateLocalVector` or
170 * `DMGetLocalVector` are not applicable.
171 * For example, if statitics are being store at quadrature points, a `DM`-created `Vec` will not have the
172 * correct size.
173 *
174 * @param[in] op Operator to make the Vecs for
175 * @param[in] vec_type `VecType` for the new Vecs
176 * @param[in] comm `MPI_Comm` for the new Vecs
177 * @param[out] input Vec for CeedOperator active input
178 * @param[out] output Vec for CeedOperator active output
179 */
CeedOperatorCreateLocalVecs(CeedOperator op,VecType vec_type,MPI_Comm comm,Vec * input,Vec * output)180 PetscErrorCode CeedOperatorCreateLocalVecs(CeedOperator op, VecType vec_type, MPI_Comm comm, Vec *input, Vec *output) {
181 CeedSize input_size, output_size;
182 Ceed ceed;
183 int comm_size;
184
185 PetscFunctionBeginUser;
186 PetscCall(CeedOperatorGetCeed(op, &ceed));
187 PetscCallMPI(MPI_Comm_size(comm, &comm_size));
188 PetscCheck(comm_size == 1, PETSC_COMM_WORLD, PETSC_ERR_ARG_SIZ, "MPI_Comm must be of size 1, recieved comm of size %d", comm_size);
189 PetscCallCeed(ceed, CeedOperatorGetActiveVectorLengths(op, &input_size, &output_size));
190 if (input) {
191 PetscCall(VecCreate(comm, input));
192 PetscCall(VecSetType(*input, vec_type));
193 PetscCall(VecSetSizes(*input, input_size, input_size));
194 }
195 if (output) {
196 PetscCall(VecCreate(comm, output));
197 PetscCall(VecSetType(*output, vec_type));
198 PetscCall(VecSetSizes(*output, output_size, output_size));
199 }
200 PetscCheck(CeedDestroy(&ceed) == CEED_ERROR_SUCCESS, comm, PETSC_ERR_LIB, "Destroying Ceed object failed");
201 PetscFunctionReturn(PETSC_SUCCESS);
202 }
203
204 /**
205 * @brief Apply FEM Operator defined by `OperatorApplyContext` to various input and output vectors
206 *
207 * @param[in] X Input global `Vec`, maybe `NULL`
208 * @param[in] X_loc Input local `Vec`, maybe `NULL`
209 * @param[in] x_ceed Input `CeedVector`, maybe `CEED_VECTOR_NONE`
210 * @param[in,out] y_ceed Output `CeedVector`, maybe `CEED_VECTOR_NONE`
211 * @param[in,out] Y_loc Output local `Vec`, maybe `NULL`
212 * @param[in,out] Y Output global `Vec`, maybe `NULL`
213 * @param[in] ctx Context for the operator apply
214 * @param[in] use_apply_add Whether to use `CeedOperatorApply` or `CeedOperatorApplyAdd`
215 */
ApplyCeedOperator_Core(Vec X,Vec X_loc,CeedVector x_ceed,CeedVector y_ceed,Vec Y_loc,Vec Y,OperatorApplyContext ctx,bool use_apply_add)216 PetscErrorCode ApplyCeedOperator_Core(Vec X, Vec X_loc, CeedVector x_ceed, CeedVector y_ceed, Vec Y_loc, Vec Y, OperatorApplyContext ctx,
217 bool use_apply_add) {
218 PetscMemType x_mem_type, y_mem_type;
219 Ceed ceed = ctx->ceed;
220
221 PetscFunctionBeginUser;
222 if (X) PetscCall(DMGlobalToLocal(ctx->dm_x, X, INSERT_VALUES, X_loc));
223 if (X_loc) PetscCall(VecReadPetscToCeed(X_loc, &x_mem_type, x_ceed));
224
225 if (Y_loc) PetscCall(VecPetscToCeed(Y_loc, &y_mem_type, y_ceed));
226
227 PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorApply, X, Y, 0, 0));
228 PetscCall(PetscLogGpuTimeBegin());
229 if (use_apply_add) PetscCallCeed(ceed, CeedOperatorApplyAdd(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE));
230 else PetscCallCeed(ceed, CeedOperatorApply(ctx->op, x_ceed, y_ceed, CEED_REQUEST_IMMEDIATE));
231 PetscCall(PetscLogGpuTimeEnd());
232 PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorApply, X, Y, 0, 0));
233
234 if (X_loc) PetscCall(VecReadCeedToPetsc(ctx->x_ceed, x_mem_type, X_loc));
235
236 if (Y_loc) PetscCall(VecCeedToPetsc(ctx->y_ceed, y_mem_type, Y_loc));
237 if (Y) PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, Y));
238 PetscFunctionReturn(PETSC_SUCCESS);
239 };
240
ApplyCeedOperatorGlobalToGlobal(Vec X,Vec Y,OperatorApplyContext ctx)241 PetscErrorCode ApplyCeedOperatorGlobalToGlobal(Vec X, Vec Y, OperatorApplyContext ctx) {
242 Vec X_loc = ctx->X_loc, Y_loc = ctx->Y_loc;
243
244 PetscFunctionBeginUser;
245 PetscCall(VecZeroEntries(Y));
246
247 // Get local vectors (if needed)
248 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc));
249 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc));
250
251 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false));
252
253 // Restore local vector (if needed)
254 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc));
255 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc));
256 PetscFunctionReturn(PETSC_SUCCESS);
257 }
258
ApplyCeedOperatorLocalToGlobal(Vec X_loc,Vec Y,OperatorApplyContext ctx)259 PetscErrorCode ApplyCeedOperatorLocalToGlobal(Vec X_loc, Vec Y, OperatorApplyContext ctx) {
260 Vec Y_loc = ctx->Y_loc;
261
262 PetscFunctionBeginUser;
263 PetscCall(VecZeroEntries(Y));
264
265 // Get local vectors (if needed)
266 if (!ctx->Y_loc) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc));
267
268 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, Y, ctx, false));
269
270 // Restore local vectors (if needed)
271 if (!ctx->Y_loc) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc));
272 PetscFunctionReturn(PETSC_SUCCESS);
273 }
274
ApplyCeedOperatorGlobalToLocal(Vec X,Vec Y_loc,OperatorApplyContext ctx)275 PetscErrorCode ApplyCeedOperatorGlobalToLocal(Vec X, Vec Y_loc, OperatorApplyContext ctx) {
276 Vec X_loc = ctx->X_loc;
277
278 PetscFunctionBeginUser;
279 // Get local vectors (if needed)
280 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc));
281
282 PetscCall(ApplyCeedOperator_Core(X, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, false));
283
284 // Restore local vector (if needed)
285 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc));
286 PetscFunctionReturn(PETSC_SUCCESS);
287 }
288
ApplyCeedOperatorLocalToLocal(Vec X_loc,Vec Y_loc,OperatorApplyContext ctx)289 PetscErrorCode ApplyCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorApplyContext ctx) {
290 PetscFunctionBeginUser;
291 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, false));
292 PetscFunctionReturn(PETSC_SUCCESS);
293 }
294
ApplyAddCeedOperatorLocalToLocal(Vec X_loc,Vec Y_loc,OperatorApplyContext ctx)295 PetscErrorCode ApplyAddCeedOperatorLocalToLocal(Vec X_loc, Vec Y_loc, OperatorApplyContext ctx) {
296 PetscFunctionBeginUser;
297 PetscCall(ApplyCeedOperator_Core(NULL, X_loc, ctx->x_ceed, ctx->y_ceed, Y_loc, NULL, ctx, true));
298 PetscFunctionReturn(PETSC_SUCCESS);
299 }
300
301 /**
302 * @brief Return Mats for KSP solve
303 *
304 * Uses command-line flag with `ksp`'s prefix to determine if mat_ceed should be used directly or whether it should be assembled.
305 * If `Amat` is to be assembled, then `Pmat` is equal to `Amat`.
306 *
307 * If `Amat` uses `mat_ceed`, then `Pmat` is either assembled or uses `mat_ceed` based on the preconditioner choice in `ksp`.
308 *
309 * @param[in] ksp `KSP` object for used for solving
310 * @param[in] mat_ceed `MATCEED` for the linear operator
311 * @param[in] assemble Whether to assemble `Amat` and `Pmat` if they are not `mat_ceed`
312 * @param[out] Amat `Mat` to be used for the solver `Amat`
313 * @param[out] Pmat `Mat` to be used for the solver `Pmat`
314 */
CreateSolveOperatorsFromMatCeed(KSP ksp,Mat mat_ceed,PetscBool assemble,Mat * Amat,Mat * Pmat)315 PetscErrorCode CreateSolveOperatorsFromMatCeed(KSP ksp, Mat mat_ceed, PetscBool assemble, Mat *Amat, Mat *Pmat) {
316 PetscBool use_matceed_pmat, assemble_amat = PETSC_FALSE;
317
318 PetscFunctionBeginUser;
319 { // Determine if Amat should be MATCEED or assembled
320 const char *ksp_prefix = NULL;
321
322 PetscCall(KSPGetOptionsPrefix(ksp, &ksp_prefix));
323 PetscOptionsBegin(PetscObjectComm((PetscObject)mat_ceed), ksp_prefix, "", NULL);
324 PetscCall(PetscOptionsBool("-matceed_assemble_amat", "Assemble the A matrix for KSP solve", NULL, assemble_amat, &assemble_amat, NULL));
325 PetscOptionsEnd();
326 }
327
328 if (assemble_amat) {
329 PetscCall(MatCeedCreateMatCOO(mat_ceed, Amat));
330 if (assemble) PetscCall(MatCeedAssembleCOO(mat_ceed, *Amat));
331
332 PetscCall(PetscObjectReference((PetscObject)*Amat));
333 *Pmat = *Amat;
334 PetscFunctionReturn(PETSC_SUCCESS);
335 } else {
336 PetscCall(PetscObjectReference((PetscObject)mat_ceed));
337 *Amat = mat_ceed;
338 }
339
340 { // Determine if Pmat should be MATCEED or assembled
341 PC pc;
342 PCType pc_type;
343
344 PetscCall(KSPGetPC(ksp, &pc));
345 PetscCall(PCGetType(pc, &pc_type));
346 PetscCall(PetscStrcmpAny(pc_type, &use_matceed_pmat, PCNONE, PCJACOBI, PCVPBJACOBI, PCPBJACOBI, ""));
347 }
348
349 if (use_matceed_pmat) {
350 PetscCall(PetscObjectReference((PetscObject)mat_ceed));
351 *Pmat = mat_ceed;
352 } else {
353 PetscCall(MatCeedCreateMatCOO(mat_ceed, Pmat));
354 if (assemble) PetscCall(MatCeedAssembleCOO(mat_ceed, *Pmat));
355 }
356 PetscFunctionReturn(PETSC_SUCCESS);
357 }
358
359 /**
360 * @brief Runs KSPSetFromOptions and sets Operators based on mat_ceed
361 *
362 * See CreateSolveOperatorsFromMatCeed for details on how the KSPSolve operators are set.
363 *
364 * @param[in] ksp `KSP` of the solve
365 * @param[in] mat_ceed `MatCeed` linear operator to solve for
366 */
KSPSetFromOptions_WithMatCeed(KSP ksp,Mat mat_ceed)367 PetscErrorCode KSPSetFromOptions_WithMatCeed(KSP ksp, Mat mat_ceed) {
368 Mat Amat, Pmat;
369
370 PetscFunctionBeginUser;
371 PetscCall(KSPSetFromOptions(ksp));
372 PetscCall(CreateSolveOperatorsFromMatCeed(ksp, mat_ceed, PETSC_TRUE, &Amat, &Pmat));
373 PetscCall(KSPSetOperators(ksp, Amat, Pmat));
374 PetscCall(MatDestroy(&Amat));
375 PetscCall(MatDestroy(&Pmat));
376 PetscFunctionReturn(PETSC_SUCCESS);
377 }
378 // -----------------------------------------------------------------------------
379