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