1 // SPDX-FileCopyrightText: Copyright (c) 2017-2024, HONEE contributors.
2 // SPDX-License-Identifier: Apache-2.0 OR BSD-2-Clause
3 /// @file
4 /// MatCEED implementation
5
6 #include <ceed.h>
7 #include <ceed/backend.h>
8 #include <mat-ceed-impl.h>
9 #include <mat-ceed.h>
10 #include <petsc-ceed-utils.h>
11 #include <petsc-ceed.h>
12 #include <petscdm.h>
13 #include <petscmat.h>
14 #include <stdbool.h>
15 #include <stdio.h>
16 #include <stdlib.h>
17 #include <string.h>
18
19 #if PETSC_VERSION_LT(3, 24, 0)
20 typedef void (*MatSetOpFn)(void);
21 #else
22 typedef PetscErrorCodeFn *MatSetOpFn;
23 #endif
24
25 PetscClassId MATCEED_CLASSID;
26 PetscLogEvent MATCEED_MULT, MATCEED_MULT_CEEDOP, MATCEED_MULT_TRANSPOSE, MATCEED_MULT_TRANSPOSE_CEEDOP, MATCEED_ASSEMBLE_DIAGONAL,
27 MATCEED_ASSEMBLE_DIAGONAL_CEEDOP, MATCEED_SETUP_PBDIAGONAL, MATCEED_SETUP_PBDIAGONAL_CEEDOP, MATCEED_ASSEMBLE_PBDIAGONAL,
28 MATCEED_ASSEMBLE_PBDIAGONAL_CEEDOP, MATCEED_SETUP_FULL, MATCEED_SETUP_FULL_CEEDOP, MATCEED_ASSEMBLE_FULL, MATCEED_ASSEMBLE_FULL_CEEDOP;
29
30 /**
31 @brief Register MATCEED log events.
32
33 Not collective across MPI processes.
34
35 @return An error code: 0 - success, otherwise - failure
36 **/
MatCeedRegisterLogEvents()37 static PetscErrorCode MatCeedRegisterLogEvents() {
38 static PetscBool registered = PETSC_FALSE;
39
40 PetscFunctionBeginUser;
41 if (registered) PetscFunctionReturn(PETSC_SUCCESS);
42 PetscCall(PetscClassIdRegister("MatCEED", &MATCEED_CLASSID));
43 PetscCall(PetscLogEventRegister("MatCEEDMul", MATCEED_CLASSID, &MATCEED_MULT));
44 PetscCall(PetscLogEventRegister("MatCEEDMulCeed", MATCEED_CLASSID, &MATCEED_MULT_CEEDOP));
45 PetscCall(PetscLogEventRegister("MatCEEDMulT", MATCEED_CLASSID, &MATCEED_MULT_TRANSPOSE));
46 PetscCall(PetscLogEventRegister("MatCEEDMulTCeed", MATCEED_CLASSID, &MATCEED_MULT_TRANSPOSE_CEEDOP));
47 PetscCall(PetscLogEventRegister("MatCEEDAsmDiag", MATCEED_CLASSID, &MATCEED_ASSEMBLE_DIAGONAL));
48 PetscCall(PetscLogEventRegister("MatCEEDAsmPBDSU", MATCEED_CLASSID, &MATCEED_SETUP_PBDIAGONAL));
49 PetscCall(PetscLogEventRegister("MatCEEDAsmPBDSUCeed", MATCEED_CLASSID, &MATCEED_SETUP_PBDIAGONAL_CEEDOP));
50 PetscCall(PetscLogEventRegister("MatCEEDAsmPBD", MATCEED_CLASSID, &MATCEED_ASSEMBLE_PBDIAGONAL));
51 PetscCall(PetscLogEventRegister("MatCEEDAsmPBDCeed", MATCEED_CLASSID, &MATCEED_ASSEMBLE_PBDIAGONAL_CEEDOP));
52 PetscCall(PetscLogEventRegister("MatCEEDAsmSU", MATCEED_CLASSID, &MATCEED_SETUP_FULL));
53 PetscCall(PetscLogEventRegister("MatCEEDAsmSUCeed", MATCEED_CLASSID, &MATCEED_SETUP_FULL_CEEDOP));
54 PetscCall(PetscLogEventRegister("MatCEEDAsm", MATCEED_CLASSID, &MATCEED_ASSEMBLE_FULL));
55 PetscCall(PetscLogEventRegister("MatCEEDAsmCeed", MATCEED_CLASSID, &MATCEED_ASSEMBLE_FULL_CEEDOP));
56 registered = PETSC_TRUE;
57 PetscFunctionReturn(PETSC_SUCCESS);
58 }
59
60 /**
61 @brief Assemble the point block diagonal of a `MATCEED` into a `MATAIJ` or similar.
62 The `mat_coo` preallocation is set to match the sparsity pattern of `mat_ceed`.
63 The caller is responsible for assuring the global and local sizes are compatible, otherwise this function will fail.
64
65 Collective across MPI processes.
66
67 @param[in] mat_ceed `MATCEED` to assemble
68 @param[in,out] mat_coo `MATAIJ` or similar to assemble into
69
70 @return An error code: 0 - success, otherwise - failure
71 **/
MatCeedAssemblePointBlockDiagonalCOO(Mat mat_ceed,Mat mat_coo)72 static PetscErrorCode MatCeedAssemblePointBlockDiagonalCOO(Mat mat_ceed, Mat mat_coo) {
73 MatCeedContext ctx;
74
75 PetscFunctionBeginUser;
76 PetscCall(MatShellGetContext(mat_ceed, &ctx));
77
78 // Check if COO pattern set
79 {
80 PetscInt index = -1;
81
82 for (PetscInt i = 0; i < ctx->num_mats_assembled_pbd; i++) {
83 if (ctx->mats_assembled_pbd[i] == mat_coo) index = i;
84 }
85 if (index == -1) {
86 PetscInt *rows_petsc = NULL, *cols_petsc = NULL;
87 CeedInt *rows_ceed, *cols_ceed;
88 PetscCount num_entries;
89 PetscLogStage stage_amg_setup;
90
91 // -- Assemble sparsity pattern if mat hasn't been assembled before
92 PetscCall(PetscLogStageGetId("MatCEED Asm Setup", &stage_amg_setup));
93 if (stage_amg_setup == -1) {
94 PetscCall(PetscLogStageRegister("MatCEED Asm Setup", &stage_amg_setup));
95 }
96 PetscCall(PetscLogStagePush(stage_amg_setup));
97 PetscCall(PetscLogEventBegin(MATCEED_SETUP_PBDIAGONAL, mat_ceed, mat_coo, NULL, NULL));
98 PetscCall(PetscLogEventBegin(MATCEED_SETUP_PBDIAGONAL_CEEDOP, mat_ceed, mat_coo, NULL, NULL));
99 PetscCallCeed(ctx->ceed, CeedOperatorLinearAssemblePointBlockDiagonalSymbolic(ctx->op_mult, &num_entries, &rows_ceed, &cols_ceed));
100 PetscCall(PetscLogEventEnd(MATCEED_SETUP_PBDIAGONAL_CEEDOP, mat_ceed, mat_coo, NULL, NULL));
101 PetscCall(IntArrayCeedToPetsc(num_entries, &rows_ceed, &rows_petsc));
102 PetscCall(IntArrayCeedToPetsc(num_entries, &cols_ceed, &cols_petsc));
103 PetscCall(MatSetPreallocationCOOLocal(mat_coo, num_entries, rows_petsc, cols_petsc));
104 free(rows_petsc);
105 free(cols_petsc);
106 if (!ctx->coo_values_pbd) PetscCallCeed(ctx->ceed, CeedVectorCreate(ctx->ceed, num_entries, &ctx->coo_values_pbd));
107 PetscCall(PetscRealloc(++ctx->num_mats_assembled_pbd * sizeof(Mat), &ctx->mats_assembled_pbd));
108 ctx->mats_assembled_pbd[ctx->num_mats_assembled_pbd - 1] = mat_coo;
109 PetscCall(PetscLogEventEnd(MATCEED_SETUP_PBDIAGONAL, mat_ceed, mat_coo, NULL, NULL));
110 PetscCall(PetscLogStagePop());
111 }
112 }
113
114 // Assemble mat_ceed
115 PetscCall(PetscLogEventBegin(MATCEED_ASSEMBLE_PBDIAGONAL, mat_ceed, mat_coo, NULL, NULL));
116 PetscCall(MatAssemblyBegin(mat_coo, MAT_FINAL_ASSEMBLY));
117 {
118 const CeedScalar *values;
119 PetscMemType mem_type_petsc;
120 CeedMemType mem_type = CEED_MEM_HOST;
121 PetscBool is_spd, is_spd_known;
122
123 PetscCall(MatGetMemTypeFromVecType(mat_coo, &mem_type_petsc));
124 mem_type = MemTypePetscToCeed(mem_type_petsc);
125 PetscCall(PetscLogEventBegin(MATCEED_ASSEMBLE_PBDIAGONAL_CEEDOP, mat_ceed, mat_coo, NULL, NULL));
126 PetscCallCeed(ctx->ceed, CeedOperatorLinearAssemblePointBlockDiagonal(ctx->op_mult, ctx->coo_values_pbd, CEED_REQUEST_IMMEDIATE));
127 PetscCall(PetscLogEventEnd(MATCEED_ASSEMBLE_PBDIAGONAL_CEEDOP, mat_ceed, mat_coo, NULL, NULL));
128 PetscCallCeed(ctx->ceed, CeedVectorGetArrayRead(ctx->coo_values_pbd, mem_type, &values));
129 PetscCall(MatSetValuesCOO(mat_coo, values, INSERT_VALUES));
130 PetscCall(MatIsSPDKnown(mat_ceed, &is_spd_known, &is_spd));
131 if (is_spd_known) PetscCall(MatSetOption(mat_coo, MAT_SPD, is_spd));
132 PetscCallCeed(ctx->ceed, CeedVectorRestoreArrayRead(ctx->coo_values_pbd, &values));
133 }
134 PetscCall(MatAssemblyEnd(mat_coo, MAT_FINAL_ASSEMBLY));
135 PetscCall(PetscLogEventEnd(MATCEED_ASSEMBLE_PBDIAGONAL, mat_ceed, mat_coo, NULL, NULL));
136 PetscFunctionReturn(PETSC_SUCCESS);
137 }
138
139 /**
140 @brief Assemble inner `Mat` for diagonal `PC` operations
141
142 Collective across MPI processes.
143
144 @param[in] mat_ceed `MATCEED` to assemble block diagonal of
145 @param[in] use_ceed_pbd Boolean flag to use libCEED PBD assembly
146 @param[out] mat_inner Inner `Mat` for diagonal operations
147
148 @return An error code: 0 - success, otherwise - failure
149 **/
MatCeedAssembleInnerBlockDiagonalMat(Mat mat_ceed,PetscBool use_ceed_pbd,Mat * mat_inner)150 static PetscErrorCode MatCeedAssembleInnerBlockDiagonalMat(Mat mat_ceed, PetscBool use_ceed_pbd, Mat *mat_inner) {
151 MatCeedContext ctx;
152
153 PetscFunctionBeginUser;
154 PetscCall(MatShellGetContext(mat_ceed, &ctx));
155 if (use_ceed_pbd) {
156 // Check if COO pattern set
157 if (!ctx->mat_assembled_pbd_internal) PetscCall(MatCeedCreateMatCOO(mat_ceed, &ctx->mat_assembled_pbd_internal));
158
159 // Assemble mat_assembled_full_internal
160 PetscCall(MatCeedAssemblePointBlockDiagonalCOO(mat_ceed, ctx->mat_assembled_pbd_internal));
161 if (mat_inner) *mat_inner = ctx->mat_assembled_pbd_internal;
162 } else {
163 // Check if COO pattern set
164 if (!ctx->mat_assembled_full_internal) PetscCall(MatCeedCreateMatCOO(mat_ceed, &ctx->mat_assembled_full_internal));
165
166 // Assemble mat_assembled_full_internal
167 PetscCall(MatCeedAssembleCOO(mat_ceed, ctx->mat_assembled_full_internal));
168 if (mat_inner) *mat_inner = ctx->mat_assembled_full_internal;
169 }
170 PetscFunctionReturn(PETSC_SUCCESS);
171 }
172
173 /**
174 @brief Get `MATCEED` variable block diagonal for Jacobi.
175
176 Collective across MPI processes.
177
178 @param[in] mat_ceed `MATCEED` to get variable block diagonal of
179 @param[out] mat_vblock The variable diagonal block matrix
180
181 @return An error code: 0 - success, otherwise - failure
182 **/
MatGetVariableBlockDiagonal_Ceed(Mat mat_ceed,Mat * mat_vblock)183 static PetscErrorCode MatGetVariableBlockDiagonal_Ceed(Mat mat_ceed, Mat *mat_vblock) {
184 MatCeedContext ctx;
185
186 PetscFunctionBeginUser;
187 PetscCall(MatShellGetContext(mat_ceed, &ctx));
188
189 // Assemble inner mat if needed
190 PetscCall(MatCeedAssembleInnerBlockDiagonalMat(mat_ceed, ctx->is_ceed_vpbd_valid, mat_vblock));
191 PetscCall(PetscObjectReference((PetscObject)*mat_vblock));
192 PetscFunctionReturn(PETSC_SUCCESS);
193 }
194
195 /**
196 @brief Get `MATCEED` block diagonal for Jacobi.
197
198 Collective across MPI processes.
199
200 @param[in] mat_ceed `MATCEED` to to get block diagonal of
201 @param[out] mat_block The diagonal block matrix
202
203 @return An error code: 0 - success, otherwise - failure
204 **/
MatGetBlockDiagonal_Ceed(Mat mat_ceed,Mat * mat_block)205 static PetscErrorCode MatGetBlockDiagonal_Ceed(Mat mat_ceed, Mat *mat_block) {
206 MatCeedContext ctx;
207
208 PetscFunctionBeginUser;
209 PetscCall(MatShellGetContext(mat_ceed, &ctx));
210
211 // Assemble inner mat if needed
212 PetscCall(MatCeedAssembleInnerBlockDiagonalMat(mat_ceed, ctx->is_ceed_pbd_valid, mat_block));
213 PetscCall(PetscObjectReference((PetscObject)*mat_block));
214 PetscFunctionReturn(PETSC_SUCCESS);
215 }
216
217 /**
218 @brief Get on-process diagonal block of `MATCEED`.
219
220 This is a block per-process of the diagonal of the global matrix.
221 This is NOT the diagonal blocks associated with the block size of the matrix (i.e. `MatSetBlockSize()` has no effect on this function).
222
223 Collective across MPI processes.
224
225 @param[in] mat_ceed `MATCEED` to get the block per-proces of
226 @param[out] mat_block The diagonal block matrix
227
228 @return An error code: 0 - success, otherwise - failure
229 **/
MatGetDiagonalBlock_Ceed(Mat mat_ceed,Mat * mat_block)230 static PetscErrorCode MatGetDiagonalBlock_Ceed(Mat mat_ceed, Mat *mat_block) {
231 MatCeedContext ctx;
232
233 PetscFunctionBeginUser;
234 PetscCall(MatShellGetContext(mat_ceed, &ctx));
235
236 // Check if COO pattern set
237 if (!ctx->mat_assembled_full_internal) PetscCall(MatCeedCreateMatCOO(mat_ceed, &ctx->mat_assembled_full_internal));
238
239 // Assemble mat_assembled_full_internal
240 PetscCall(MatCeedAssembleCOO(mat_ceed, ctx->mat_assembled_full_internal));
241
242 // Get diagonal block
243 PetscCall(MatGetDiagonalBlock(ctx->mat_assembled_full_internal, mat_block));
244 PetscFunctionReturn(PETSC_SUCCESS);
245 }
246
247 /**
248 @brief View `MATCEED`.
249
250 Collective across MPI processes.
251
252 @param[in] mat_ceed `MATCEED` to view
253 @param[in] viewer The visualization context
254
255 @return An error code: 0 - success, otherwise - failure
256 **/
MatView_Ceed(Mat mat_ceed,PetscViewer viewer)257 static PetscErrorCode MatView_Ceed(Mat mat_ceed, PetscViewer viewer) {
258 PetscBool is_ascii;
259 PetscViewerFormat format;
260 PetscMPIInt size, rank;
261 MatCeedContext ctx;
262
263 PetscFunctionBeginUser;
264 PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
265 PetscCall(MatShellGetContext(mat_ceed, &ctx));
266 if (!viewer) PetscCall(PetscViewerASCIIGetStdout(PetscObjectComm((PetscObject)mat_ceed), &viewer));
267
268 PetscCall(PetscViewerGetFormat(viewer, &format));
269 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat_ceed), &size));
270 if (size == 1 && format == PETSC_VIEWER_LOAD_BALANCE) PetscFunctionReturn(PETSC_SUCCESS);
271
272 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat_ceed), &rank));
273 if (rank != 0) PetscFunctionReturn(PETSC_SUCCESS);
274
275 PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &is_ascii));
276 {
277 PetscBool is_detailed = format == PETSC_VIEWER_ASCII_INFO_DETAIL;
278 char rank_string[16] = {'\0'};
279 PetscInt num_tabs = 0;
280 CeedInt prev_num_tabs;
281 FILE *file;
282
283 PetscCall(PetscViewerASCIIPrintf(viewer, "MatCEED:\n"));
284 PetscCall(PetscViewerASCIIPushTab(viewer)); // MatCEED
285 PetscCall(PetscViewerASCIIPrintf(viewer, "Default COO MatType: %s\n", ctx->coo_mat_type));
286 PetscCall(PetscSNPrintf(rank_string, 16, "on Rank %d", rank));
287 PetscCall(PetscViewerASCIIPrintf(viewer, "CeedOperator Apply %s:\n", is_detailed ? rank_string : "Summary"));
288 PetscCall(PetscViewerASCIIPrintf(viewer, "libCEED PB Diagonal Assembly: %s\n", ctx->is_ceed_pbd_valid ? "True" : "False"));
289 PetscCall(PetscViewerASCIIPrintf(viewer, "libCEED VPB Diagonal Assembly: %s\n", ctx->is_ceed_vpbd_valid ? "True" : "False"));
290 PetscCall(PetscViewerASCIIPrintf(viewer, "CeedOperator Apply %s:\n", is_detailed ? rank_string : "Summary"));
291 PetscCall(PetscViewerASCIIGetPointer(viewer, &file));
292 PetscCall(PetscViewerASCIIPushTab(viewer)); // CeedOperator
293 PetscCall(PetscViewerASCIIGetTab(viewer, &num_tabs));
294 PetscCallCeed(ctx->ceed, CeedOperatorSetNumViewTabs(ctx->op_mult, num_tabs));
295 if (is_detailed) PetscCallCeed(ctx->ceed, CeedOperatorView(ctx->op_mult, file));
296 else PetscCallCeed(ctx->ceed, CeedOperatorViewTerse(ctx->op_mult, file));
297 PetscCallCeed(ctx->ceed, CeedOperatorSetNumViewTabs(ctx->op_mult, 0));
298 PetscCall(PetscViewerASCIIPopTab(viewer)); // CeedOperator
299 if (ctx->op_mult_transpose) {
300 PetscCall(PetscViewerASCIIPrintf(viewer, "CeedOperator ApplyTranspose %s:\n", is_detailed ? rank_string : "Summary"));
301 PetscCall(PetscViewerASCIIPushTab(viewer)); // CeedOperator
302 PetscCall(PetscViewerASCIIGetTab(viewer, &num_tabs));
303 PetscCallCeed(ctx->ceed, CeedOperatorGetNumViewTabs(ctx->op_mult_transpose, &prev_num_tabs));
304 PetscCallCeed(ctx->ceed, CeedOperatorSetNumViewTabs(ctx->op_mult_transpose, num_tabs));
305 if (is_detailed) PetscCallCeed(ctx->ceed, CeedOperatorView(ctx->op_mult_transpose, file));
306 else PetscCallCeed(ctx->ceed, CeedOperatorViewTerse(ctx->op_mult_transpose, file));
307 PetscCallCeed(ctx->ceed, CeedOperatorSetNumViewTabs(ctx->op_mult_transpose, prev_num_tabs));
308 PetscCall(PetscViewerASCIIPopTab(viewer)); // CeedOperator
309 }
310 PetscCall(PetscViewerASCIIPopTab(viewer)); // MatCEED
311 }
312 PetscFunctionReturn(PETSC_SUCCESS);
313 }
314
315 /**
316 @brief Set options for MatCeed from command-line interface
317
318 @note This is a PETSc interface, thus the odd signature
319 @note The name of the second parameter must be `PetscOptionsObject` due to abuse of PETSc macros
320
321 Collective across MPI processes.
322
323 @param[in] mat_ceed `Mat` object to set up
324 @param[in] PetscOptionsObject `PetscOptionItems` object to read options from
325
326 @return An error code: 0 - success, otherwise - failure
327 **/
MatSetFromOptions_Ceed(Mat mat_ceed,PetscOptionItems PetscOptionsObject)328 static PetscErrorCode MatSetFromOptions_Ceed(Mat mat_ceed, PetscOptionItems PetscOptionsObject) {
329 MatCeedContext ctx;
330 char coo_mat_type_cl[PETSC_MAX_OPTION_NAME];
331 PetscBool is_coo_mat_type_cl;
332
333 PetscFunctionBegin;
334 PetscCall(MatShellGetContext(mat_ceed, &ctx));
335
336 PetscOptionsHeadBegin(PetscOptionsObject, "MATCEED options");
337 // Check for specific CL coo mat type for this Mat
338 PetscCall(PetscOptionsFList("-mat_ceed_coo_mat_type", "Default MATCEED COO assembly MatType", NULL, MatList, coo_mat_type_cl, coo_mat_type_cl,
339 sizeof(coo_mat_type_cl), &is_coo_mat_type_cl));
340 if (is_coo_mat_type_cl) {
341 PetscCall(PetscFree(ctx->coo_mat_type));
342 PetscCall(PetscStrallocpy(coo_mat_type_cl, &ctx->coo_mat_type));
343 }
344 // Check for COO mat reuse flag
345 PetscCall(PetscOptionsBool("-mat_ceed_coo_reuse_preallocation",
346 "Flag to allow the MatCeed to reuse existing COO preallocations, even if not set by this MatCeed", NULL,
347 ctx->coo_reuse_preallocation, &ctx->coo_reuse_preallocation, NULL));
348 PetscOptionsHeadEnd();
349 PetscFunctionReturn(PETSC_SUCCESS);
350 }
351
352 // -----------------------------------------------------------------------------
353 // MatCeed
354 // -----------------------------------------------------------------------------
355
356 /**
357 @brief Create PETSc `Mat` from libCEED operators.
358
359 Collective across MPI processes.
360
361 @param[in] dm_x Input `DM`
362 @param[in] dm_y Output `DM`
363 @param[in] op_mult `CeedOperator` for forward evaluation
364 @param[in] op_mult_transpose `CeedOperator` for transpose evaluation
365 @param[out] mat New MatCeed
366
367 @return An error code: 0 - success, otherwise - failure
368 **/
MatCreateCeed(DM dm_x,DM dm_y,CeedOperator op_mult,CeedOperator op_mult_transpose,Mat * mat)369 PetscErrorCode MatCreateCeed(DM dm_x, DM dm_y, CeedOperator op_mult, CeedOperator op_mult_transpose, Mat *mat) {
370 PetscInt X_l_size, X_g_size, Y_l_size, Y_g_size;
371 VecType vec_type;
372 MatCeedContext ctx;
373
374 PetscFunctionBeginUser;
375 PetscCall(MatCeedRegisterLogEvents());
376
377 // Collect context data
378 PetscCall(DMGetVecType(dm_x, &vec_type));
379 {
380 Vec X;
381
382 PetscCall(DMGetGlobalVector(dm_x, &X));
383 PetscCall(VecGetSize(X, &X_g_size));
384 PetscCall(VecGetLocalSize(X, &X_l_size));
385 PetscCall(DMRestoreGlobalVector(dm_x, &X));
386 }
387 if (dm_y) {
388 Vec Y;
389
390 PetscCall(DMGetGlobalVector(dm_y, &Y));
391 PetscCall(VecGetSize(Y, &Y_g_size));
392 PetscCall(VecGetLocalSize(Y, &Y_l_size));
393 PetscCall(DMRestoreGlobalVector(dm_y, &Y));
394 } else {
395 dm_y = dm_x;
396 Y_g_size = X_g_size;
397 Y_l_size = X_l_size;
398 }
399
400 // Create context
401 {
402 Vec X_loc, Y_loc_transpose = NULL;
403
404 PetscCall(DMCreateLocalVector(dm_x, &X_loc));
405 PetscCall(VecZeroEntries(X_loc));
406 if (op_mult_transpose) {
407 PetscCall(DMCreateLocalVector(dm_y, &Y_loc_transpose));
408 PetscCall(VecZeroEntries(Y_loc_transpose));
409 }
410 PetscCall(MatCeedContextCreate(dm_x, dm_y, X_loc, Y_loc_transpose, op_mult, op_mult_transpose, MATCEED_MULT, MATCEED_MULT_TRANSPOSE,
411 MATCEED_MULT_CEEDOP, MATCEED_MULT_TRANSPOSE_CEEDOP, &ctx));
412 PetscCall(VecDestroy(&X_loc));
413 PetscCall(VecDestroy(&Y_loc_transpose));
414 }
415
416 // Create mat
417 PetscCall(MatCreateShell(PetscObjectComm((PetscObject)dm_x), Y_l_size, X_l_size, Y_g_size, X_g_size, ctx, mat));
418 PetscCall(PetscObjectChangeTypeName((PetscObject)*mat, MATCEED));
419 // -- Set block and variable block sizes
420 if (dm_x == dm_y) {
421 MatType dm_mat_type, dm_mat_type_copy;
422 Mat temp_mat;
423
424 PetscCall(DMGetMatType(dm_x, &dm_mat_type));
425 PetscCall(PetscStrallocpy(dm_mat_type, (char **)&dm_mat_type_copy));
426 PetscCall(DMSetMatType(dm_x, MATAIJ));
427 PetscCall(DMCreateMatrix(dm_x, &temp_mat));
428 PetscCall(DMSetMatType(dm_x, dm_mat_type_copy));
429 PetscCall(PetscFree(dm_mat_type_copy));
430
431 {
432 PetscInt block_size, num_blocks, max_vblock_size = PETSC_INT_MAX;
433 const PetscInt *vblock_sizes;
434
435 // -- Get block sizes
436 PetscCall(MatGetBlockSize(temp_mat, &block_size));
437 PetscCall(MatGetVariableBlockSizes(temp_mat, &num_blocks, &vblock_sizes));
438 {
439 PetscInt local_min_max[2] = {0}, global_min_max[2] = {0, PETSC_INT_MAX};
440
441 for (PetscInt i = 0; i < num_blocks; i++) local_min_max[1] = PetscMax(local_min_max[1], vblock_sizes[i]);
442 PetscCall(PetscGlobalMinMaxInt(PetscObjectComm((PetscObject)dm_x), local_min_max, global_min_max));
443 max_vblock_size = global_min_max[1];
444 }
445
446 // -- Copy block sizes
447 if (block_size > 1) PetscCall(MatSetBlockSize(*mat, block_size));
448 if (num_blocks) PetscCall(MatSetVariableBlockSizes(*mat, num_blocks, (PetscInt *)vblock_sizes));
449
450 // -- Check libCEED compatibility
451 {
452 bool is_composite;
453
454 ctx->is_ceed_pbd_valid = PETSC_TRUE;
455 ctx->is_ceed_vpbd_valid = PETSC_TRUE;
456 PetscCallCeed(ctx->ceed, CeedOperatorIsComposite(op_mult, &is_composite));
457 if (is_composite) {
458 CeedInt num_sub_operators;
459 CeedOperator *sub_operators;
460
461 PetscCallCeed(ctx->ceed, CeedOperatorCompositeGetNumSub(op_mult, &num_sub_operators));
462 PetscCallCeed(ctx->ceed, CeedOperatorCompositeGetSubList(op_mult, &sub_operators));
463 for (CeedInt i = 0; i < num_sub_operators; i++) {
464 CeedInt num_bases, num_comp;
465 CeedBasis *active_bases;
466 CeedOperatorAssemblyData assembly_data;
467
468 PetscCallCeed(ctx->ceed, CeedOperatorGetOperatorAssemblyData(sub_operators[i], &assembly_data));
469 PetscCallCeed(ctx->ceed, CeedOperatorAssemblyDataGetBases(assembly_data, &num_bases, &active_bases, NULL, NULL, NULL, NULL));
470 PetscCallCeed(ctx->ceed, CeedBasisGetNumComponents(active_bases[0], &num_comp));
471 if (num_bases > 1) {
472 ctx->is_ceed_pbd_valid = PETSC_FALSE;
473 ctx->is_ceed_vpbd_valid = PETSC_FALSE;
474 }
475 if (num_comp != block_size) ctx->is_ceed_pbd_valid = PETSC_FALSE;
476 if (num_comp < max_vblock_size) ctx->is_ceed_vpbd_valid = PETSC_FALSE;
477 }
478 } else {
479 // LCOV_EXCL_START
480 CeedInt num_bases, num_comp;
481 CeedBasis *active_bases;
482 CeedOperatorAssemblyData assembly_data;
483
484 PetscCallCeed(ctx->ceed, CeedOperatorGetOperatorAssemblyData(op_mult, &assembly_data));
485 PetscCallCeed(ctx->ceed, CeedOperatorAssemblyDataGetBases(assembly_data, &num_bases, &active_bases, NULL, NULL, NULL, NULL));
486 PetscCallCeed(ctx->ceed, CeedBasisGetNumComponents(active_bases[0], &num_comp));
487 if (num_bases > 1) {
488 ctx->is_ceed_pbd_valid = PETSC_FALSE;
489 ctx->is_ceed_vpbd_valid = PETSC_FALSE;
490 }
491 if (num_comp != block_size) ctx->is_ceed_pbd_valid = PETSC_FALSE;
492 if (num_comp < max_vblock_size) ctx->is_ceed_vpbd_valid = PETSC_FALSE;
493 // LCOV_EXCL_STOP
494 }
495 {
496 PetscInt local_is_valid[2], global_is_valid[2];
497
498 local_is_valid[0] = local_is_valid[1] = ctx->is_ceed_pbd_valid;
499 PetscCall(PetscGlobalMinMaxInt(PetscObjectComm((PetscObject)dm_x), local_is_valid, global_is_valid));
500 ctx->is_ceed_pbd_valid = global_is_valid[0];
501 local_is_valid[0] = local_is_valid[1] = ctx->is_ceed_vpbd_valid;
502 PetscCall(PetscGlobalMinMaxInt(PetscObjectComm((PetscObject)dm_x), local_is_valid, global_is_valid));
503 ctx->is_ceed_vpbd_valid = global_is_valid[0];
504 }
505 }
506 }
507 PetscCall(MatDestroy(&temp_mat));
508 }
509 // -- Set internal mat type
510 {
511 VecType vec_type;
512 MatType coo_mat_type;
513
514 PetscCall(VecGetType(ctx->X_loc, &vec_type));
515 PetscCall(DefaultMatTypeFromVecType(vec_type, &coo_mat_type));
516 PetscCall(PetscStrallocpy(coo_mat_type, &ctx->coo_mat_type));
517 }
518 // -- Set mat operations
519 PetscCall(MatShellSetContextDestroy(*mat, (PetscCtxDestroyFn *)MatCeedContextDestroy));
520 PetscCall(MatShellSetOperation(*mat, MATOP_VIEW, (MatSetOpFn)MatView_Ceed));
521 PetscCall(MatShellSetOperation(*mat, MATOP_MULT, (MatSetOpFn)MatMult_Ceed));
522 if (op_mult_transpose) PetscCall(MatShellSetOperation(*mat, MATOP_MULT_TRANSPOSE, (MatSetOpFn)MatMultTranspose_Ceed));
523 PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL, (MatSetOpFn)MatGetDiagonal_Ceed));
524 PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL_BLOCK, (MatSetOpFn)MatGetDiagonalBlock_Ceed));
525 PetscCall(MatShellSetOperation(*mat, MATOP_GET_BLOCK_DIAGONAL, (MatSetOpFn)MatGetBlockDiagonal_Ceed));
526 PetscCall(MatShellSetOperation(*mat, MATOP_GET_VBLOCK_DIAGONAL, (MatSetOpFn)MatGetVariableBlockDiagonal_Ceed));
527 PetscCall(MatShellSetOperation(*mat, MATOP_SET_FROM_OPTIONS, (MatSetOpFn)MatSetFromOptions_Ceed));
528 PetscCall(MatShellSetVecType(*mat, vec_type));
529 PetscFunctionReturn(PETSC_SUCCESS);
530 }
531
532 /**
533 @brief Copy `MATCEED` into a compatible `Mat` with type `MatShell` or `MATCEED`.
534
535 Collective across MPI processes.
536
537 @param[in] mat_ceed `MATCEED` to copy from
538 @param[out] mat_other `MatShell` or `MATCEED` to copy into
539
540 @return An error code: 0 - success, otherwise - failure
541 **/
MatCeedCopy(Mat mat_ceed,Mat mat_other)542 PetscErrorCode MatCeedCopy(Mat mat_ceed, Mat mat_other) {
543 PetscFunctionBeginUser;
544 PetscCall(MatCeedRegisterLogEvents());
545
546 // Check type compatibility
547 {
548 PetscBool is_matceed = PETSC_FALSE, is_matshell = PETSC_FALSE;
549 MatType mat_type_ceed, mat_type_other;
550
551 PetscCall(MatGetType(mat_ceed, &mat_type_ceed));
552 PetscCall(PetscStrcmp(mat_type_ceed, MATCEED, &is_matceed));
553 PetscCheck(is_matceed, PETSC_COMM_SELF, PETSC_ERR_LIB, "mat_ceed must have type " MATCEED);
554 PetscCall(MatGetType(mat_other, &mat_type_other));
555 PetscCall(PetscStrcmp(mat_type_other, MATCEED, &is_matceed));
556 PetscCall(PetscStrcmp(mat_type_other, MATSHELL, &is_matshell));
557 PetscCheck(is_matceed || is_matshell, PETSC_COMM_SELF, PETSC_ERR_LIB, "mat_other must have type " MATCEED " or " MATSHELL);
558 }
559
560 // Check dimension compatibility
561 {
562 PetscInt X_l_ceed_size, X_g_ceed_size, Y_l_ceed_size, Y_g_ceed_size, X_l_other_size, X_g_other_size, Y_l_other_size, Y_g_other_size;
563
564 PetscCall(MatGetSize(mat_ceed, &Y_g_ceed_size, &X_g_ceed_size));
565 PetscCall(MatGetLocalSize(mat_ceed, &Y_l_ceed_size, &X_l_ceed_size));
566 PetscCall(MatGetSize(mat_ceed, &Y_g_other_size, &X_g_other_size));
567 PetscCall(MatGetLocalSize(mat_ceed, &Y_l_other_size, &X_l_other_size));
568 PetscCheck((Y_g_ceed_size == Y_g_other_size) && (X_g_ceed_size == X_g_other_size) && (Y_l_ceed_size == Y_l_other_size) &&
569 (X_l_ceed_size == X_l_other_size),
570 PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ,
571 "mat_ceed and mat_other must have compatible sizes; found mat_ceed (Global: %" PetscInt_FMT ", %" PetscInt_FMT
572 "; Local: %" PetscInt_FMT ", %" PetscInt_FMT ") mat_other (Global: %" PetscInt_FMT ", %" PetscInt_FMT "; Local: %" PetscInt_FMT
573 ", %" PetscInt_FMT ")",
574 Y_g_ceed_size, X_g_ceed_size, Y_l_ceed_size, X_l_ceed_size, Y_g_other_size, X_g_other_size, Y_l_other_size, X_l_other_size);
575 }
576
577 // Convert
578 {
579 VecType vec_type;
580 MatCeedContext ctx;
581
582 PetscCall(PetscObjectChangeTypeName((PetscObject)mat_other, MATCEED));
583 PetscCall(MatShellGetContext(mat_ceed, &ctx));
584 PetscCall(MatCeedContextReference(ctx));
585 PetscCall(MatShellSetContext(mat_other, ctx));
586 PetscCall(MatShellSetContextDestroy(mat_other, (PetscCtxDestroyFn *)MatCeedContextDestroy));
587 PetscCall(MatShellSetOperation(mat_other, MATOP_VIEW, (MatSetOpFn)MatView_Ceed));
588 PetscCall(MatShellSetOperation(mat_other, MATOP_MULT, (MatSetOpFn)MatMult_Ceed));
589 if (ctx->op_mult_transpose) PetscCall(MatShellSetOperation(mat_other, MATOP_MULT_TRANSPOSE, (MatSetOpFn)MatMultTranspose_Ceed));
590 PetscCall(MatShellSetOperation(mat_other, MATOP_GET_DIAGONAL, (MatSetOpFn)MatGetDiagonal_Ceed));
591 PetscCall(MatShellSetOperation(mat_other, MATOP_GET_DIAGONAL_BLOCK, (MatSetOpFn)MatGetDiagonalBlock_Ceed));
592 PetscCall(MatShellSetOperation(mat_other, MATOP_GET_BLOCK_DIAGONAL, (MatSetOpFn)MatGetBlockDiagonal_Ceed));
593 PetscCall(MatShellSetOperation(mat_other, MATOP_GET_VBLOCK_DIAGONAL, (MatSetOpFn)MatGetVariableBlockDiagonal_Ceed));
594 {
595 PetscInt block_size;
596
597 PetscCall(MatGetBlockSize(mat_ceed, &block_size));
598 if (block_size > 1) PetscCall(MatSetBlockSize(mat_other, block_size));
599 }
600 {
601 PetscInt num_blocks;
602 const PetscInt *block_sizes;
603
604 PetscCall(MatGetVariableBlockSizes(mat_ceed, &num_blocks, &block_sizes));
605 if (num_blocks) PetscCall(MatSetVariableBlockSizes(mat_other, num_blocks, (PetscInt *)block_sizes));
606 }
607 PetscCall(DMGetVecType(ctx->dm_x, &vec_type));
608 PetscCall(MatShellSetVecType(mat_other, vec_type));
609 }
610 PetscFunctionReturn(PETSC_SUCCESS);
611 }
612
613 /**
614 @brief Mark `CeedQFunction` data as updated and the `CeedQFunction` as requiring re-assembly for a `MatCEED`.
615
616 Collective across MPI processes.
617
618 @param[in] mat_ceed `MATCEED`
619 @param[out] update_needed Boolean flag indicating `CeedQFunction` update needed
620
621 @return An error code: 0 - success, otherwise - failure
622 **/
MatCeedSetAssemblyDataUpdateNeeded(Mat mat_ceed,PetscBool update_needed)623 PetscErrorCode MatCeedSetAssemblyDataUpdateNeeded(Mat mat_ceed, PetscBool update_needed) {
624 MatCeedContext ctx;
625
626 PetscFunctionBeginUser;
627 PetscCall(MatShellGetContext(mat_ceed, &ctx));
628 PetscCallCeed(ctx->ceed, CeedOperatorSetQFunctionAssemblyDataUpdateNeeded(ctx->op_mult, update_needed));
629 if (ctx->op_mult_transpose) {
630 PetscCallCeed(ctx->ceed, CeedOperatorSetQFunctionAssemblyDataUpdateNeeded(ctx->op_mult_transpose, update_needed));
631 }
632 if (update_needed) {
633 PetscCall(MatAssemblyBegin(mat_ceed, MAT_FINAL_ASSEMBLY));
634 PetscCall(MatAssemblyEnd(mat_ceed, MAT_FINAL_ASSEMBLY));
635 }
636 PetscFunctionReturn(PETSC_SUCCESS);
637 }
638
639 /**
640 @brief Setup a `Mat` with the same COO pattern as a `MatCEED`.
641
642 Collective across MPI processes.
643
644 @param[in] mat_ceed `MATCEED`
645 @param[out] mat_coo Sparse `Mat` with same COO pattern
646
647 @return An error code: 0 - success, otherwise - failure
648 **/
MatCeedCreateMatCOO(Mat mat_ceed,Mat * mat_coo)649 PetscErrorCode MatCeedCreateMatCOO(Mat mat_ceed, Mat *mat_coo) {
650 MatCeedContext ctx;
651
652 PetscFunctionBeginUser;
653 PetscCall(MatShellGetContext(mat_ceed, &ctx));
654
655 PetscCheck(ctx->dm_x == ctx->dm_y, PetscObjectComm((PetscObject)mat_ceed), PETSC_ERR_SUP, "COO assembly only supported for MATCEED on a single DM");
656
657 // Create sparse matrix
658 {
659 MatType dm_mat_type, dm_mat_type_copy;
660
661 PetscCall(DMGetMatType(ctx->dm_x, &dm_mat_type));
662 PetscCall(PetscStrallocpy(dm_mat_type, (char **)&dm_mat_type_copy));
663 PetscCall(DMSetMatType(ctx->dm_x, ctx->coo_mat_type));
664 PetscCall(DMCreateMatrix(ctx->dm_x, mat_coo));
665 PetscCall(DMSetMatType(ctx->dm_x, dm_mat_type_copy));
666 PetscCall(PetscFree(dm_mat_type_copy));
667 }
668 PetscFunctionReturn(PETSC_SUCCESS);
669 }
670
671 /**
672 @brief Setup the COO preallocation `MATCEED` into a `MATAIJ` or similar.
673 The caller is responsible for assuring the global and local sizes are compatible, otherwise this function will fail.
674
675 Collective across MPI processes.
676
677 @param[in] mat_ceed `MATCEED` to assemble
678 @param[in,out] mat_coo `MATAIJ` or similar to assemble into
679
680 @return An error code: 0 - success, otherwise - failure
681 **/
MatCeedSetPreallocationCOO(Mat mat_ceed,Mat mat_coo)682 PetscErrorCode MatCeedSetPreallocationCOO(Mat mat_ceed, Mat mat_coo) {
683 MatCeedContext ctx;
684
685 PetscFunctionBeginUser;
686 PetscCall(MatShellGetContext(mat_ceed, &ctx));
687
688 {
689 PetscCount num_entries;
690 PetscLogStage stage_amg_setup;
691 PetscObject coo_struct;
692
693 PetscCall(PetscLogStageGetId("MatCEED Asm Setup", &stage_amg_setup));
694 if (stage_amg_setup == -1) {
695 PetscCall(PetscLogStageRegister("MatCEED Asm Setup", &stage_amg_setup));
696 }
697 PetscCall(PetscLogStagePush(stage_amg_setup));
698 PetscCall(PetscLogEventBegin(MATCEED_SETUP_FULL, mat_ceed, mat_coo, NULL, NULL));
699 PetscCall(PetscObjectQuery((PetscObject)mat_coo, "__PETSc_MatCOOStruct_Host", (PetscObject *)&coo_struct));
700 // -- Check for existing assembly data
701 if (!ctx->coo_reuse_preallocation || !coo_struct) {
702 PetscInt *rows_petsc = NULL, *cols_petsc = NULL;
703 CeedInt *rows_ceed, *cols_ceed;
704
705 // -- Assemble sparsity pattern if mat hasn't been assembled before
706 PetscCall(PetscLogEventBegin(MATCEED_SETUP_FULL_CEEDOP, mat_ceed, mat_coo, NULL, NULL));
707 PetscCallCeed(ctx->ceed, CeedOperatorLinearAssembleSymbolic(ctx->op_mult, &num_entries, &rows_ceed, &cols_ceed));
708 PetscCall(PetscLogEventEnd(MATCEED_SETUP_FULL_CEEDOP, mat_ceed, mat_coo, NULL, NULL));
709 PetscCall(IntArrayCeedToPetsc(num_entries, &rows_ceed, &rows_petsc));
710 PetscCall(IntArrayCeedToPetsc(num_entries, &cols_ceed, &cols_petsc));
711 PetscCall(MatSetPreallocationCOOLocal(mat_coo, num_entries, rows_petsc, cols_petsc));
712 free(rows_petsc);
713 free(cols_petsc);
714 } else {
715 PetscCallCeed(ctx->ceed, CeedOperatorLinearAssembleGetNumEntries(ctx->op_mult, &num_entries));
716 }
717
718 if (!ctx->coo_values_full) PetscCallCeed(ctx->ceed, CeedVectorCreate(ctx->ceed, num_entries, &ctx->coo_values_full));
719 PetscCall(PetscRealloc(++ctx->num_mats_assembled_full * sizeof(Mat), &ctx->mats_assembled_full));
720 ctx->mats_assembled_full[ctx->num_mats_assembled_full - 1] = mat_coo;
721 PetscCall(PetscLogEventEnd(MATCEED_SETUP_FULL, mat_ceed, mat_coo, NULL, NULL));
722 PetscCall(PetscLogStagePop());
723 }
724 PetscFunctionReturn(PETSC_SUCCESS);
725 }
726
727 /**
728 @brief Assemble a `MATCEED` into a `MATAIJ` or similar.
729 The `mat_coo` preallocation is set to match the sparsity pattern of `mat_ceed`.
730 The caller is responsible for assuring the global and local sizes are compatible, otherwise this function will fail.
731
732 Collective across MPI processes.
733
734 @param[in] mat_ceed `MATCEED` to assemble
735 @param[in,out] mat_coo `MATAIJ` or similar to assemble into
736
737 @return An error code: 0 - success, otherwise - failure
738 **/
MatCeedAssembleCOO(Mat mat_ceed,Mat mat_coo)739 PetscErrorCode MatCeedAssembleCOO(Mat mat_ceed, Mat mat_coo) {
740 MatCeedContext ctx;
741
742 PetscFunctionBeginUser;
743 PetscCall(MatShellGetContext(mat_ceed, &ctx));
744
745 // Set COO pattern if needed
746 {
747 CeedInt index = -1;
748
749 for (PetscInt i = 0; i < ctx->num_mats_assembled_full; i++) {
750 if (ctx->mats_assembled_full[i] == mat_coo) index = i;
751 }
752 if (index == -1) PetscCall(MatCeedSetPreallocationCOO(mat_ceed, mat_coo));
753 }
754
755 // Assemble mat_ceed
756 PetscCall(PetscLogEventBegin(MATCEED_ASSEMBLE_FULL, mat_ceed, mat_coo, NULL, NULL));
757 {
758 const CeedScalar *values;
759 CeedMemType mem_type = CEED_MEM_HOST;
760 PetscBool is_spd, is_spd_known;
761 PetscMemType mem_type_petsc;
762
763 PetscCall(MatGetMemTypeFromVecType(mat_coo, &mem_type_petsc));
764 mem_type = MemTypePetscToCeed(mem_type_petsc);
765 PetscCall(PetscLogEventBegin(MATCEED_ASSEMBLE_FULL_CEEDOP, mat_ceed, mat_coo, NULL, NULL));
766 PetscCallCeed(ctx->ceed, CeedOperatorLinearAssemble(ctx->op_mult, ctx->coo_values_full));
767 PetscCall(PetscLogEventEnd(MATCEED_ASSEMBLE_FULL_CEEDOP, mat_ceed, mat_coo, NULL, NULL));
768 PetscCallCeed(ctx->ceed, CeedVectorGetArrayRead(ctx->coo_values_full, mem_type, &values));
769 PetscCall(MatSetValuesCOO(mat_coo, values, INSERT_VALUES));
770 PetscCall(MatIsSPDKnown(mat_ceed, &is_spd_known, &is_spd));
771 if (is_spd_known) PetscCall(MatSetOption(mat_coo, MAT_SPD, is_spd));
772 PetscCallCeed(ctx->ceed, CeedVectorRestoreArrayRead(ctx->coo_values_full, &values));
773 }
774 PetscCall(MatAssemblyBegin(mat_coo, MAT_FINAL_ASSEMBLY));
775 PetscCall(MatAssemblyEnd(mat_coo, MAT_FINAL_ASSEMBLY));
776 PetscCall(PetscLogEventEnd(MATCEED_ASSEMBLE_FULL, mat_ceed, mat_coo, NULL, NULL));
777 PetscFunctionReturn(PETSC_SUCCESS);
778 }
779
780 /**
781 @brief Set the current value of a context field for a `MatCEED`.
782
783 Not collective across MPI processes.
784
785 @param[in,out] mat `MatCEED`
786 @param[in] name Name of the context field
787 @param[in] value New context field value
788
789 @return An error code: 0 - success, otherwise - failure
790 **/
MatCeedSetContextDouble(Mat mat,const char * name,double value)791 PetscErrorCode MatCeedSetContextDouble(Mat mat, const char *name, double value) {
792 PetscBool was_updated = PETSC_FALSE;
793 MatCeedContext ctx;
794
795 PetscFunctionBeginUser;
796 PetscCall(MatShellGetContext(mat, &ctx));
797 {
798 CeedContextFieldLabel label = NULL;
799
800 PetscCallCeed(ctx->ceed, CeedOperatorGetContextFieldLabel(ctx->op_mult, name, &label));
801 if (label) {
802 double set_value = 2 * value + 1.0;
803
804 PetscCall(MatCeedGetContextDouble(mat, name, &set_value));
805 if (set_value != value) {
806 PetscCallCeed(ctx->ceed, CeedOperatorSetContextDouble(ctx->op_mult, label, &value));
807 was_updated = PETSC_TRUE;
808 }
809 }
810 if (ctx->op_mult_transpose) {
811 label = NULL;
812 PetscCallCeed(ctx->ceed, CeedOperatorGetContextFieldLabel(ctx->op_mult_transpose, name, &label));
813 if (label) {
814 double set_value = 2 * value + 1.0;
815
816 PetscCall(MatCeedGetContextDouble(mat, name, &set_value));
817 if (set_value != value) {
818 PetscCallCeed(ctx->ceed, CeedOperatorSetContextDouble(ctx->op_mult_transpose, label, &value));
819 was_updated = PETSC_TRUE;
820 }
821 }
822 }
823 }
824 if (was_updated) {
825 PetscCall(MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY));
826 PetscCall(MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY));
827 }
828 PetscFunctionReturn(PETSC_SUCCESS);
829 }
830
831 /**
832 @brief Get the current value of a context field for a `MatCEED`.
833
834 Not collective across MPI processes.
835
836 @param[in] mat `MatCEED`
837 @param[in] name Name of the context field
838 @param[out] value Current context field value
839
840 @return An error code: 0 - success, otherwise - failure
841 **/
MatCeedGetContextDouble(Mat mat,const char * name,double * value)842 PetscErrorCode MatCeedGetContextDouble(Mat mat, const char *name, double *value) {
843 MatCeedContext ctx;
844
845 PetscFunctionBeginUser;
846 PetscCall(MatShellGetContext(mat, &ctx));
847 {
848 CeedContextFieldLabel label = NULL;
849 CeedOperator op = ctx->op_mult;
850
851 PetscCallCeed(ctx->ceed, CeedOperatorGetContextFieldLabel(op, name, &label));
852 if (!label && ctx->op_mult_transpose) {
853 op = ctx->op_mult_transpose;
854 PetscCallCeed(ctx->ceed, CeedOperatorGetContextFieldLabel(op, name, &label));
855 }
856 if (label) {
857 PetscSizeT num_values;
858 const double *values_ceed;
859
860 PetscCallCeed(ctx->ceed, CeedOperatorGetContextDoubleRead(op, label, &num_values, &values_ceed));
861 *value = values_ceed[0];
862 PetscCallCeed(ctx->ceed, CeedOperatorRestoreContextDoubleRead(op, label, &values_ceed));
863 }
864 }
865 PetscFunctionReturn(PETSC_SUCCESS);
866 }
867
868 /**
869 @brief Set the current `PetscReal` value of a context field for a `MatCEED`.
870
871 Not collective across MPI processes.
872
873 @param[in,out] mat `MatCEED`
874 @param[in] name Name of the context field
875 @param[in] value New context field value
876
877 @return An error code: 0 - success, otherwise - failure
878 **/
MatCeedSetContextReal(Mat mat,const char * name,PetscReal value)879 PetscErrorCode MatCeedSetContextReal(Mat mat, const char *name, PetscReal value) {
880 double value_double = value;
881
882 PetscFunctionBeginUser;
883 PetscCall(MatCeedSetContextDouble(mat, name, value_double));
884 PetscFunctionReturn(PETSC_SUCCESS);
885 }
886
887 /**
888 @brief Get the current `PetscReal` value of a context field for a `MatCEED`.
889
890 Not collective across MPI processes.
891
892 @param[in] mat `MatCEED`
893 @param[in] name Name of the context field
894 @param[out] value Current context field value
895
896 @return An error code: 0 - success, otherwise - failure
897 **/
MatCeedGetContextReal(Mat mat,const char * name,PetscReal * value)898 PetscErrorCode MatCeedGetContextReal(Mat mat, const char *name, PetscReal *value) {
899 double value_double = 0.0;
900
901 PetscFunctionBeginUser;
902 PetscCall(MatCeedGetContextDouble(mat, name, &value_double));
903 *value = value_double;
904 PetscFunctionReturn(PETSC_SUCCESS);
905 }
906
907 /**
908 @brief Set the current time for a `MatCEED`.
909
910 Not collective across MPI processes.
911
912 @param[in,out] mat `MatCEED`
913 @param[in] time Current time
914
915 @return An error code: 0 - success, otherwise - failure
916 **/
MatCeedSetTime(Mat mat,PetscReal time)917 PetscErrorCode MatCeedSetTime(Mat mat, PetscReal time) {
918 PetscFunctionBeginUser;
919 {
920 double time_ceed = time;
921
922 PetscCall(MatCeedSetContextDouble(mat, "time", time_ceed));
923 }
924 PetscFunctionReturn(PETSC_SUCCESS);
925 }
926
927 /**
928 @brief Get the current time for a `MatCEED`.
929
930 Not collective across MPI processes.
931
932 @param[in] mat `MatCEED`
933 @param[out] time Current time, or -1.0 if the boundary evaluator has no time field
934
935 @return An error code: 0 - success, otherwise - failure
936 **/
MatCeedGetTime(Mat mat,PetscReal * time)937 PetscErrorCode MatCeedGetTime(Mat mat, PetscReal *time) {
938 PetscFunctionBeginUser;
939 *time = -1.0;
940 {
941 double time_ceed = -1.0;
942
943 PetscCall(MatCeedGetContextDouble(mat, "time", &time_ceed));
944 *time = time_ceed;
945 }
946 PetscFunctionReturn(PETSC_SUCCESS);
947 }
948
949 /**
950 @brief Set the current time step for a `MatCEED`.
951
952 Not collective across MPI processes.
953
954 @param[in,out] mat `MatCEED`
955 @param[in] dt Current time step
956
957 @return An error code: 0 - success, otherwise - failure
958 **/
MatCeedSetDt(Mat mat,PetscReal dt)959 PetscErrorCode MatCeedSetDt(Mat mat, PetscReal dt) {
960 PetscFunctionBeginUser;
961 {
962 double dt_ceed = dt;
963
964 PetscCall(MatCeedSetContextDouble(mat, "dt", dt_ceed));
965 }
966 PetscFunctionReturn(PETSC_SUCCESS);
967 }
968
969 /**
970 @brief Set the Jacobian shifts for a `MatCEED`.
971
972 Not collective across MPI processes.
973
974 @param[in,out] mat `MatCEED`
975 @param[in] shift_v Velocity shift
976 @param[in] shift_a Acceleration shift
977
978 @return An error code: 0 - success, otherwise - failure
979 **/
MatCeedSetShifts(Mat mat,PetscReal shift_v,PetscReal shift_a)980 PetscErrorCode MatCeedSetShifts(Mat mat, PetscReal shift_v, PetscReal shift_a) {
981 PetscFunctionBeginUser;
982 {
983 double shift_v_ceed = shift_v;
984
985 PetscCall(MatCeedSetContextDouble(mat, "shift v", shift_v_ceed));
986 }
987 if (shift_a) {
988 double shift_a_ceed = shift_a;
989
990 PetscCall(MatCeedSetContextDouble(mat, "shift a", shift_a_ceed));
991 }
992 PetscFunctionReturn(PETSC_SUCCESS);
993 }
994
995 /**
996 @brief Set user context for a `MATCEED`.
997
998 Collective across MPI processes.
999
1000 @param[in,out] mat `MATCEED`
1001 @param[in] f The context destroy function, or NULL
1002 @param[in] ctx User context, or NULL to unset
1003
1004 @return An error code: 0 - success, otherwise - failure
1005 **/
MatCeedSetContext(Mat mat,PetscCtxDestroyFn f,void * ctx)1006 PetscErrorCode MatCeedSetContext(Mat mat, PetscCtxDestroyFn f, void *ctx) {
1007 PetscContainer user_ctx = NULL;
1008
1009 PetscFunctionBeginUser;
1010 if (ctx) {
1011 PetscCall(PetscContainerCreate(PetscObjectComm((PetscObject)mat), &user_ctx));
1012 PetscCall(PetscContainerSetPointer(user_ctx, ctx));
1013 PetscCall(PetscContainerSetCtxDestroy(user_ctx, f));
1014 }
1015 PetscCall(PetscObjectCompose((PetscObject)mat, "MatCeed user context", (PetscObject)user_ctx));
1016 PetscCall(PetscContainerDestroy(&user_ctx));
1017 PetscFunctionReturn(PETSC_SUCCESS);
1018 }
1019
1020 /**
1021 @brief Retrieve the user context for a `MATCEED`.
1022
1023 Collective across MPI processes.
1024
1025 @param[in,out] mat `MATCEED`
1026 @param[in] ctx User context
1027
1028 @return An error code: 0 - success, otherwise - failure
1029 **/
MatCeedGetContext(Mat mat,void * ctx)1030 PetscErrorCode MatCeedGetContext(Mat mat, void *ctx) {
1031 PetscContainer user_ctx;
1032
1033 PetscFunctionBeginUser;
1034 PetscCall(PetscObjectQuery((PetscObject)mat, "MatCeed user context", (PetscObject *)&user_ctx));
1035 if (user_ctx) PetscCall(PetscContainerGetPointer(user_ctx, (void **)ctx));
1036 else *(void **)ctx = NULL;
1037 PetscFunctionReturn(PETSC_SUCCESS);
1038 }
1039 /**
1040 @brief Set a user defined matrix operation for a `MATCEED` matrix.
1041
1042 Within each user-defined routine, the user should call `MatCeedGetContext()` to obtain the user-defined context that was set by
1043 `MatCeedSetContext()`.
1044
1045 Collective across MPI processes.
1046
1047 @param[in,out] mat `MATCEED`
1048 @param[in] op Name of the `MatOperation`
1049 @param[in] g Function that provides the operation
1050
1051 @return An error code: 0 - success, otherwise - failure
1052 **/
MatCeedSetOperation(Mat mat,MatOperation op,void (* g)(void))1053 PetscErrorCode MatCeedSetOperation(Mat mat, MatOperation op, void (*g)(void)) {
1054 PetscFunctionBeginUser;
1055 PetscCall(MatShellSetOperation(mat, op, (MatSetOpFn)g));
1056 PetscFunctionReturn(PETSC_SUCCESS);
1057 }
1058
1059 /**
1060 @brief Set the flag for whether the `MATCEED` should reuse a COO matrix, even if not created by this `MATCEED`.
1061
1062 @note Only set this flag if you are certain that the COO layout is properly set on incoming matrices.
1063
1064 Collective across MPI processes.
1065
1066 @param[in,out] mat `MATCEED`
1067 @param[in] coo_reuse_preallocation Flag value to set
1068
1069 @return An error code: 0 - success, otherwise - failure
1070 **/
MatCeedSetReusePreallocationCOO(Mat mat,PetscBool coo_reuse_preallocation)1071 PetscErrorCode MatCeedSetReusePreallocationCOO(Mat mat, PetscBool coo_reuse_preallocation) {
1072 MatCeedContext ctx;
1073
1074 PetscFunctionBeginUser;
1075 PetscCall(MatShellGetContext(mat, &ctx));
1076 ctx->coo_reuse_preallocation = coo_reuse_preallocation;
1077 PetscFunctionReturn(PETSC_SUCCESS);
1078 }
1079
1080 /**
1081 @brief Gets the flag for whether the `MATCEED` should reuse a COO matrix, even if not created by this `MATCEED`.
1082
1083 Collective across MPI processes.
1084
1085 @param[in] mat `MATCEED`
1086 @param[in] coo_reuse_preallocation Address to store current flag value
1087
1088 @return An error code: 0 - success, otherwise - failure
1089 **/
MatCeedGetReusePreallocationCOO(Mat mat,PetscBool * coo_reuse_preallocation)1090 PetscErrorCode MatCeedGetReusePreallocationCOO(Mat mat, PetscBool *coo_reuse_preallocation) {
1091 MatCeedContext ctx;
1092
1093 PetscFunctionBeginUser;
1094 PetscCall(MatShellGetContext(mat, &ctx));
1095 *coo_reuse_preallocation = ctx->coo_reuse_preallocation;
1096 PetscFunctionReturn(PETSC_SUCCESS);
1097 }
1098
1099 /**
1100 @brief Sets the default COO matrix type as a string from the `MATCEED`.
1101
1102 Collective across MPI processes.
1103
1104 @param[in,out] mat `MATCEED`
1105 @param[in] type COO `MatType` to set
1106
1107 @return An error code: 0 - success, otherwise - failure
1108 **/
MatCeedSetCOOMatType(Mat mat,MatType type)1109 PetscErrorCode MatCeedSetCOOMatType(Mat mat, MatType type) {
1110 MatCeedContext ctx;
1111
1112 PetscFunctionBeginUser;
1113 PetscCall(MatShellGetContext(mat, &ctx));
1114 // Check if same
1115 {
1116 size_t len_old, len_new;
1117 PetscBool is_same = PETSC_FALSE;
1118
1119 PetscCall(PetscStrlen(ctx->coo_mat_type, &len_old));
1120 PetscCall(PetscStrlen(type, &len_new));
1121 if (len_old == len_new) PetscCall(PetscStrncmp(ctx->coo_mat_type, type, len_old, &is_same));
1122 if (is_same) PetscFunctionReturn(PETSC_SUCCESS);
1123 }
1124 // Clean up old mats in different format
1125 // LCOV_EXCL_START
1126 if (ctx->mat_assembled_full_internal) {
1127 for (PetscInt i = 0; i < ctx->num_mats_assembled_full; i++) {
1128 if (ctx->mats_assembled_full[i] == ctx->mat_assembled_full_internal) {
1129 for (PetscInt j = i + 1; j < ctx->num_mats_assembled_full; j++) {
1130 ctx->mats_assembled_full[j - 1] = ctx->mats_assembled_full[j];
1131 }
1132 ctx->num_mats_assembled_full--;
1133 // Note: we'll realloc this array again, so no need to shrink the allocation
1134 PetscCall(MatDestroy(&ctx->mat_assembled_full_internal));
1135 }
1136 }
1137 }
1138 if (ctx->mat_assembled_pbd_internal) {
1139 for (PetscInt i = 0; i < ctx->num_mats_assembled_pbd; i++) {
1140 if (ctx->mats_assembled_pbd[i] == ctx->mat_assembled_pbd_internal) {
1141 for (PetscInt j = i + 1; j < ctx->num_mats_assembled_pbd; j++) {
1142 ctx->mats_assembled_pbd[j - 1] = ctx->mats_assembled_pbd[j];
1143 }
1144 // Note: we'll realloc this array again, so no need to shrink the allocation
1145 ctx->num_mats_assembled_pbd--;
1146 PetscCall(MatDestroy(&ctx->mat_assembled_pbd_internal));
1147 }
1148 }
1149 }
1150 PetscCall(PetscFree(ctx->coo_mat_type));
1151 PetscCall(PetscStrallocpy(type, &ctx->coo_mat_type));
1152 PetscFunctionReturn(PETSC_SUCCESS);
1153 // LCOV_EXCL_STOP
1154 }
1155
1156 /**
1157 @brief Gets the default COO matrix type as a string from the `MATCEED`.
1158
1159 Collective across MPI processes.
1160
1161 @param[in,out] mat `MATCEED`
1162 @param[in] type COO `MatType`
1163
1164 @return An error code: 0 - success, otherwise - failure
1165 **/
MatCeedGetCOOMatType(Mat mat,MatType * type)1166 PetscErrorCode MatCeedGetCOOMatType(Mat mat, MatType *type) {
1167 MatCeedContext ctx;
1168
1169 PetscFunctionBeginUser;
1170 PetscCall(MatShellGetContext(mat, &ctx));
1171 *type = ctx->coo_mat_type;
1172 PetscFunctionReturn(PETSC_SUCCESS);
1173 }
1174
1175 /**
1176 @brief Set input local vectors for `MATCEED` `MatMult()` and `MatMultTranspose()` operations.
1177
1178 Not collective across MPI processes.
1179
1180 @param[in,out] mat `MATCEED`
1181 @param[in] X_loc Input PETSc local vector, or NULL
1182 @param[in] Y_loc_transpose Input PETSc local vector for transpose operation, or NULL
1183
1184 @return An error code: 0 - success, otherwise - failure
1185 **/
MatCeedSetLocalVectors(Mat mat,Vec X_loc,Vec Y_loc_transpose)1186 PetscErrorCode MatCeedSetLocalVectors(Mat mat, Vec X_loc, Vec Y_loc_transpose) {
1187 MatCeedContext ctx;
1188
1189 PetscFunctionBeginUser;
1190 PetscCall(MatShellGetContext(mat, &ctx));
1191 if (X_loc) {
1192 PetscInt len_old, len_new;
1193
1194 PetscCall(VecGetSize(ctx->X_loc, &len_old));
1195 PetscCall(VecGetSize(X_loc, &len_new));
1196 PetscCheck(len_old == len_new, PETSC_COMM_SELF, PETSC_ERR_LIB, "new X_loc length %" PetscInt_FMT " should match old X_loc length %" PetscInt_FMT,
1197 len_new, len_old);
1198 PetscCall(VecReferenceCopy(X_loc, &ctx->X_loc));
1199 }
1200 if (Y_loc_transpose) {
1201 PetscInt len_old, len_new;
1202
1203 PetscCall(VecGetSize(ctx->Y_loc_transpose, &len_old));
1204 PetscCall(VecGetSize(Y_loc_transpose, &len_new));
1205 PetscCheck(len_old == len_new, PETSC_COMM_SELF, PETSC_ERR_LIB,
1206 "new Y_loc_transpose length %" PetscInt_FMT " should match old Y_loc_transpose length %" PetscInt_FMT, len_new, len_old);
1207 PetscCall(VecReferenceCopy(Y_loc_transpose, &ctx->Y_loc_transpose));
1208 }
1209 PetscFunctionReturn(PETSC_SUCCESS);
1210 }
1211
1212 /**
1213 @brief Get input local vectors for `MATCEED` `MatMult()` and `MatMultTranspose()` operations.
1214
1215 Not collective across MPI processes.
1216
1217 @param[in,out] mat `MATCEED`
1218 @param[out] X_loc Input PETSc local vector, or NULL
1219 @param[out] Y_loc_transpose Input PETSc local vector for transpose operation, or NULL
1220
1221 @return An error code: 0 - success, otherwise - failure
1222 **/
MatCeedGetLocalVectors(Mat mat,Vec * X_loc,Vec * Y_loc_transpose)1223 PetscErrorCode MatCeedGetLocalVectors(Mat mat, Vec *X_loc, Vec *Y_loc_transpose) {
1224 MatCeedContext ctx;
1225
1226 PetscFunctionBeginUser;
1227 PetscCall(MatShellGetContext(mat, &ctx));
1228 if (X_loc) {
1229 *X_loc = NULL;
1230 PetscCall(VecReferenceCopy(ctx->X_loc, X_loc));
1231 }
1232 if (Y_loc_transpose) {
1233 *Y_loc_transpose = NULL;
1234 PetscCall(VecReferenceCopy(ctx->Y_loc_transpose, Y_loc_transpose));
1235 }
1236 PetscFunctionReturn(PETSC_SUCCESS);
1237 }
1238
1239 /**
1240 @brief Restore input local vectors for `MATCEED` `MatMult()` and `MatMultTranspose()` operations.
1241
1242 Not collective across MPI processes.
1243
1244 @param[in,out] mat MatCeed
1245 @param[out] X_loc Input PETSc local vector, or NULL
1246 @param[out] Y_loc_transpose Input PETSc local vector for transpose operation, or NULL
1247
1248 @return An error code: 0 - success, otherwise - failure
1249 **/
MatCeedRestoreLocalVectors(Mat mat,Vec * X_loc,Vec * Y_loc_transpose)1250 PetscErrorCode MatCeedRestoreLocalVectors(Mat mat, Vec *X_loc, Vec *Y_loc_transpose) {
1251 PetscFunctionBeginUser;
1252 if (X_loc) PetscCall(VecDestroy(X_loc));
1253 if (Y_loc_transpose) PetscCall(VecDestroy(Y_loc_transpose));
1254 PetscFunctionReturn(PETSC_SUCCESS);
1255 }
1256
1257 /**
1258 @brief Get libCEED `CeedOperator` for `MATCEED` `MatMult()` and `MatMultTranspose()` operations.
1259
1260 Not collective across MPI processes.
1261
1262 @param[in,out] mat MatCeed
1263 @param[out] op_mult libCEED `CeedOperator` for `MatMult()`, or NULL
1264 @param[out] op_mult_transpose libCEED `CeedOperator` for `MatMultTranspose()`, or NULL
1265
1266 @return An error code: 0 - success, otherwise - failure
1267 **/
MatCeedGetCeedOperators(Mat mat,CeedOperator * op_mult,CeedOperator * op_mult_transpose)1268 PetscErrorCode MatCeedGetCeedOperators(Mat mat, CeedOperator *op_mult, CeedOperator *op_mult_transpose) {
1269 MatCeedContext ctx;
1270
1271 PetscFunctionBeginUser;
1272 PetscCall(MatShellGetContext(mat, &ctx));
1273 if (op_mult) {
1274 *op_mult = NULL;
1275 PetscCallCeed(ctx->ceed, CeedOperatorReferenceCopy(ctx->op_mult, op_mult));
1276 }
1277 if (op_mult_transpose) {
1278 *op_mult_transpose = NULL;
1279 PetscCallCeed(ctx->ceed, CeedOperatorReferenceCopy(ctx->op_mult_transpose, op_mult_transpose));
1280 }
1281 PetscFunctionReturn(PETSC_SUCCESS);
1282 }
1283
1284 /**
1285 @brief Restore libCEED `CeedOperator` for `MATCEED` `MatMult()` and `MatMultTranspose()` operations.
1286
1287 Not collective across MPI processes.
1288
1289 @param[in,out] mat MatCeed
1290 @param[out] op_mult libCEED `CeedOperator` for `MatMult()`, or NULL
1291 @param[out] op_mult_transpose libCEED `CeedOperator` for `MatMultTranspose()`, or NULL
1292
1293 @return An error code: 0 - success, otherwise - failure
1294 **/
MatCeedRestoreCeedOperators(Mat mat,CeedOperator * op_mult,CeedOperator * op_mult_transpose)1295 PetscErrorCode MatCeedRestoreCeedOperators(Mat mat, CeedOperator *op_mult, CeedOperator *op_mult_transpose) {
1296 MatCeedContext ctx;
1297
1298 PetscFunctionBeginUser;
1299 PetscCall(MatShellGetContext(mat, &ctx));
1300 if (op_mult) PetscCallCeed(ctx->ceed, CeedOperatorDestroy(op_mult));
1301 if (op_mult_transpose) PetscCallCeed(ctx->ceed, CeedOperatorDestroy(op_mult_transpose));
1302 PetscFunctionReturn(PETSC_SUCCESS);
1303 }
1304
1305 /**
1306 @brief Set `PetscLogEvent` for `MATCEED` `MatMult()` and `MatMultTranspose()` operators.
1307
1308 Not collective across MPI processes.
1309
1310 @param[in,out] mat MatCeed
1311 @param[out] log_event_mult `PetscLogEvent` for forward evaluation, or NULL
1312 @param[out] log_event_mult_transpose `PetscLogEvent` for transpose evaluation, or NULL
1313
1314 @return An error code: 0 - success, otherwise - failure
1315 **/
MatCeedSetLogEvents(Mat mat,PetscLogEvent log_event_mult,PetscLogEvent log_event_mult_transpose)1316 PetscErrorCode MatCeedSetLogEvents(Mat mat, PetscLogEvent log_event_mult, PetscLogEvent log_event_mult_transpose) {
1317 MatCeedContext ctx;
1318
1319 PetscFunctionBeginUser;
1320 PetscCall(MatShellGetContext(mat, &ctx));
1321 if (log_event_mult) ctx->log_event_mult = log_event_mult;
1322 if (log_event_mult_transpose) ctx->log_event_mult_transpose = log_event_mult_transpose;
1323 PetscFunctionReturn(PETSC_SUCCESS);
1324 }
1325
1326 /**
1327 @brief Get `PetscLogEvent` for `MATCEED` `MatMult()` and `MatMultTranspose()` operators.
1328
1329 Not collective across MPI processes.
1330
1331 @param[in,out] mat MatCeed
1332 @param[out] log_event_mult `PetscLogEvent` for forward evaluation, or NULL
1333 @param[out] log_event_mult_transpose `PetscLogEvent` for transpose evaluation, or NULL
1334
1335 @return An error code: 0 - success, otherwise - failure
1336 **/
MatCeedGetLogEvents(Mat mat,PetscLogEvent * log_event_mult,PetscLogEvent * log_event_mult_transpose)1337 PetscErrorCode MatCeedGetLogEvents(Mat mat, PetscLogEvent *log_event_mult, PetscLogEvent *log_event_mult_transpose) {
1338 MatCeedContext ctx;
1339
1340 PetscFunctionBeginUser;
1341 PetscCall(MatShellGetContext(mat, &ctx));
1342 if (log_event_mult) *log_event_mult = ctx->log_event_mult;
1343 if (log_event_mult_transpose) *log_event_mult_transpose = ctx->log_event_mult_transpose;
1344 PetscFunctionReturn(PETSC_SUCCESS);
1345 }
1346
1347 /**
1348 @brief Set `CeedOperator` `PetscLogEvent` for `MATCEED` `MatMult()` and `MatMultTranspose()` operators.
1349
1350 Not collective across MPI processes.
1351
1352 @param[in,out] mat MatCeed
1353 @param[out] log_event_mult `PetscLogEvent` for forward `CeedOperator` evaluation, or NULL
1354 @param[out] log_event_mult_transpose `PetscLogEvent` for transpose `CeedOperator` evaluation, or NULL
1355
1356 @return An error code: 0 - success, otherwise - failure
1357 **/
MatCeedSetCeedOperatorLogEvents(Mat mat,PetscLogEvent log_event_mult,PetscLogEvent log_event_mult_transpose)1358 PetscErrorCode MatCeedSetCeedOperatorLogEvents(Mat mat, PetscLogEvent log_event_mult, PetscLogEvent log_event_mult_transpose) {
1359 MatCeedContext ctx;
1360
1361 PetscFunctionBeginUser;
1362 PetscCall(MatShellGetContext(mat, &ctx));
1363 if (log_event_mult) ctx->log_event_ceed_mult = log_event_mult;
1364 if (log_event_mult_transpose) ctx->log_event_ceed_mult_transpose = log_event_mult_transpose;
1365 PetscFunctionReturn(PETSC_SUCCESS);
1366 }
1367
1368 /**
1369 @brief Get `CeedOperator` `PetscLogEvent` for `MATCEED` `MatMult()` and `MatMultTranspose()` operators.
1370
1371 Not collective across MPI processes.
1372
1373 @param[in,out] mat MatCeed
1374 @param[out] log_event_mult `PetscLogEvent` for forward `CeedOperator` evaluation, or NULL
1375 @param[out] log_event_mult_transpose `PetscLogEvent` for transpose `CeedOperator` evaluation, or NULL
1376
1377 @return An error code: 0 - success, otherwise - failure
1378 **/
MatCeedGetCeedOperatorLogEvents(Mat mat,PetscLogEvent * log_event_mult,PetscLogEvent * log_event_mult_transpose)1379 PetscErrorCode MatCeedGetCeedOperatorLogEvents(Mat mat, PetscLogEvent *log_event_mult, PetscLogEvent *log_event_mult_transpose) {
1380 MatCeedContext ctx;
1381
1382 PetscFunctionBeginUser;
1383 PetscCall(MatShellGetContext(mat, &ctx));
1384 if (log_event_mult) *log_event_mult = ctx->log_event_ceed_mult;
1385 if (log_event_mult_transpose) *log_event_mult_transpose = ctx->log_event_ceed_mult_transpose;
1386 PetscFunctionReturn(PETSC_SUCCESS);
1387 }
1388
1389 // -----------------------------------------------------------------------------
1390 // Operator context data
1391 // -----------------------------------------------------------------------------
1392
1393 /**
1394 @brief Setup context data for operator application.
1395
1396 Collective across MPI processes.
1397
1398 @param[in] dm_x Input `DM`
1399 @param[in] dm_y Output `DM`
1400 @param[in] X_loc Input PETSc local vector, or NULL
1401 @param[in] Y_loc_transpose Input PETSc local vector for transpose operation, or NULL
1402 @param[in] op_mult `CeedOperator` for forward evaluation
1403 @param[in] op_mult_transpose `CeedOperator` for transpose evaluation
1404 @param[in] log_event_mult `PetscLogEvent` for forward evaluation
1405 @param[in] log_event_mult_transpose `PetscLogEvent` for transpose evaluation
1406 @param[in] log_event_ceed_mult `PetscLogEvent` for forward `CeedOperator` evaluation
1407 @param[in] log_event_ceed_mult_transpose `PetscLogEvent` for transpose `CeedOperator` evaluation
1408 @param[out] ctx Context data for operator evaluation
1409
1410 @return An error code: 0 - success, otherwise - failure
1411 **/
MatCeedContextCreate(DM dm_x,DM dm_y,Vec X_loc,Vec Y_loc_transpose,CeedOperator op_mult,CeedOperator op_mult_transpose,PetscLogEvent log_event_mult,PetscLogEvent log_event_mult_transpose,PetscLogEvent log_event_ceed_mult,PetscLogEvent log_event_ceed_mult_transpose,MatCeedContext * ctx)1412 PetscErrorCode MatCeedContextCreate(DM dm_x, DM dm_y, Vec X_loc, Vec Y_loc_transpose, CeedOperator op_mult, CeedOperator op_mult_transpose,
1413 PetscLogEvent log_event_mult, PetscLogEvent log_event_mult_transpose, PetscLogEvent log_event_ceed_mult,
1414 PetscLogEvent log_event_ceed_mult_transpose, MatCeedContext *ctx) {
1415 CeedSize x_loc_len, y_loc_len;
1416
1417 PetscFunctionBeginUser;
1418 // Allocate
1419 PetscCall(PetscNew(ctx));
1420 (*ctx)->ref_count = 1;
1421
1422 // Logging
1423 (*ctx)->log_event_mult = log_event_mult;
1424 (*ctx)->log_event_mult_transpose = log_event_mult_transpose;
1425 (*ctx)->log_event_ceed_mult = log_event_ceed_mult;
1426 (*ctx)->log_event_ceed_mult_transpose = log_event_ceed_mult_transpose;
1427
1428 // PETSc objects
1429 PetscCall(DMReferenceCopy(dm_x, &(*ctx)->dm_x));
1430 PetscCall(DMReferenceCopy(dm_y, &(*ctx)->dm_y));
1431 if (X_loc) PetscCall(VecReferenceCopy(X_loc, &(*ctx)->X_loc));
1432 if (Y_loc_transpose) PetscCall(VecReferenceCopy(Y_loc_transpose, &(*ctx)->Y_loc_transpose));
1433
1434 // Memtype
1435 {
1436 const PetscScalar *x;
1437 Vec X;
1438
1439 PetscCall(DMGetLocalVector(dm_x, &X));
1440 PetscCall(VecGetArrayReadAndMemType(X, &x, &(*ctx)->mem_type));
1441 PetscCall(VecRestoreArrayReadAndMemType(X, &x));
1442 PetscCall(DMRestoreLocalVector(dm_x, &X));
1443 }
1444
1445 // libCEED objects
1446 PetscCheck(CeedOperatorGetCeed(op_mult, &(*ctx)->ceed) == CEED_ERROR_SUCCESS, PETSC_COMM_SELF, PETSC_ERR_LIB,
1447 "retrieving Ceed context object failed");
1448 PetscCallCeed((*ctx)->ceed, CeedOperatorGetActiveVectorLengths(op_mult, &x_loc_len, &y_loc_len));
1449 PetscCallCeed((*ctx)->ceed, CeedOperatorReferenceCopy(op_mult, &(*ctx)->op_mult));
1450 if (op_mult_transpose) PetscCallCeed((*ctx)->ceed, CeedOperatorReferenceCopy(op_mult_transpose, &(*ctx)->op_mult_transpose));
1451 if (x_loc_len == -1) x_loc_len = 0;
1452 PetscCallCeed((*ctx)->ceed, CeedVectorCreate((*ctx)->ceed, x_loc_len, &(*ctx)->x_loc));
1453 if (y_loc_len == -1) y_loc_len = 0;
1454 PetscCallCeed((*ctx)->ceed, CeedVectorCreate((*ctx)->ceed, y_loc_len, &(*ctx)->y_loc));
1455
1456 // Flop counting
1457 {
1458 CeedSize ceed_flops_estimate = 0;
1459
1460 PetscCallCeed((*ctx)->ceed, CeedOperatorGetFlopsEstimate(op_mult, &ceed_flops_estimate));
1461 (*ctx)->flops_mult = ceed_flops_estimate;
1462 if (op_mult_transpose) {
1463 PetscCallCeed((*ctx)->ceed, CeedOperatorGetFlopsEstimate(op_mult_transpose, &ceed_flops_estimate));
1464 (*ctx)->flops_mult_transpose = ceed_flops_estimate;
1465 }
1466 }
1467
1468 // Check sizes
1469 if (x_loc_len > 0 || y_loc_len > 0) {
1470 CeedSize ctx_x_loc_len, ctx_y_loc_len;
1471 PetscInt X_loc_len, dm_x_loc_len, Y_loc_len, dm_y_loc_len;
1472 Vec dm_X_loc, dm_Y_loc;
1473
1474 // -- Input
1475 PetscCall(DMGetLocalVector(dm_x, &dm_X_loc));
1476 PetscCall(VecGetLocalSize(dm_X_loc, &dm_x_loc_len));
1477 PetscCall(DMRestoreLocalVector(dm_x, &dm_X_loc));
1478 PetscCallCeed((*ctx)->ceed, CeedVectorGetLength((*ctx)->x_loc, &ctx_x_loc_len));
1479 if (X_loc) {
1480 PetscCall(VecGetLocalSize(X_loc, &X_loc_len));
1481 PetscCheck(X_loc_len == dm_x_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB,
1482 "X_loc (%" PetscInt_FMT ") must match dm_x (%" PetscInt_FMT ") dimensions", X_loc_len, dm_x_loc_len);
1483 }
1484 PetscCheck(x_loc_len == dm_x_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "op (%" CeedSize_FMT ") must match dm_x (%" PetscInt_FMT ") dimensions",
1485 x_loc_len, dm_x_loc_len);
1486 PetscCheck(x_loc_len == ctx_x_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "x_loc (%" CeedSize_FMT ") must match op dimensions (%" CeedSize_FMT ")",
1487 x_loc_len, ctx_x_loc_len);
1488
1489 // -- Output
1490 PetscCall(DMGetLocalVector(dm_y, &dm_Y_loc));
1491 PetscCall(VecGetLocalSize(dm_Y_loc, &dm_y_loc_len));
1492 PetscCall(DMRestoreLocalVector(dm_y, &dm_Y_loc));
1493 PetscCallCeed((*ctx)->ceed, CeedVectorGetLength((*ctx)->y_loc, &ctx_y_loc_len));
1494 PetscCheck(ctx_y_loc_len == dm_y_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB, "op (%" CeedSize_FMT ") must match dm_y (%" PetscInt_FMT ") dimensions",
1495 ctx_y_loc_len, dm_y_loc_len);
1496
1497 // -- Transpose
1498 if (Y_loc_transpose) {
1499 PetscCall(VecGetLocalSize(Y_loc_transpose, &Y_loc_len));
1500 PetscCheck(Y_loc_len == dm_y_loc_len, PETSC_COMM_SELF, PETSC_ERR_LIB,
1501 "Y_loc_transpose (%" PetscInt_FMT ") must match dm_y (%" PetscInt_FMT ") dimensions", Y_loc_len, dm_y_loc_len);
1502 }
1503 }
1504 PetscFunctionReturn(PETSC_SUCCESS);
1505 }
1506
1507 /**
1508 @brief Increment reference counter for `MATCEED` context.
1509
1510 Not collective across MPI processes.
1511
1512 @param[in,out] ctx Context data
1513
1514 @return An error code: 0 - success, otherwise - failure
1515 **/
MatCeedContextReference(MatCeedContext ctx)1516 PetscErrorCode MatCeedContextReference(MatCeedContext ctx) {
1517 PetscFunctionBeginUser;
1518 ctx->ref_count++;
1519 PetscFunctionReturn(PETSC_SUCCESS);
1520 }
1521
1522 /**
1523 @brief Copy reference for `MATCEED`.
1524 Note: If `ctx_copy` is non-null, it is assumed to be a valid pointer to a `MatCeedContext`.
1525
1526 Not collective across MPI processes.
1527
1528 @param[in] ctx Context data
1529 @param[out] ctx_copy Copy of pointer to context data
1530
1531 @return An error code: 0 - success, otherwise - failure
1532 **/
MatCeedContextReferenceCopy(MatCeedContext ctx,MatCeedContext * ctx_copy)1533 PetscErrorCode MatCeedContextReferenceCopy(MatCeedContext ctx, MatCeedContext *ctx_copy) {
1534 PetscFunctionBeginUser;
1535 PetscCall(MatCeedContextReference(ctx));
1536 PetscCall(MatCeedContextDestroy(ctx_copy));
1537 *ctx_copy = ctx;
1538 PetscFunctionReturn(PETSC_SUCCESS);
1539 }
1540
1541 /**
1542 @brief Destroy context data for operator application.
1543
1544 Collective across MPI processes.
1545
1546 @param[in,out] ctx Context data for operator evaluation
1547
1548 @return An error code: 0 - success, otherwise - failure
1549 **/
MatCeedContextDestroy(MatCeedContext * ctx)1550 PetscErrorCode MatCeedContextDestroy(MatCeedContext *ctx) {
1551 PetscFunctionBeginUser;
1552 if (!ctx || --(*ctx)->ref_count > 0) PetscFunctionReturn(PETSC_SUCCESS);
1553
1554 // PETSc objects
1555 PetscCall(DMDestroy(&(*ctx)->dm_x));
1556 PetscCall(DMDestroy(&(*ctx)->dm_y));
1557 PetscCall(VecDestroy(&(*ctx)->X_loc));
1558 PetscCall(VecDestroy(&(*ctx)->Y_loc_transpose));
1559 PetscCall(MatDestroy(&(*ctx)->mat_assembled_full_internal));
1560 PetscCall(MatDestroy(&(*ctx)->mat_assembled_pbd_internal));
1561 PetscCall(PetscFree((*ctx)->coo_mat_type));
1562 PetscCall(PetscFree((*ctx)->mats_assembled_full));
1563 PetscCall(PetscFree((*ctx)->mats_assembled_pbd));
1564
1565 // libCEED objects
1566 PetscCallCeed((*ctx)->ceed, CeedVectorDestroy(&(*ctx)->x_loc));
1567 PetscCallCeed((*ctx)->ceed, CeedVectorDestroy(&(*ctx)->y_loc));
1568 PetscCallCeed((*ctx)->ceed, CeedVectorDestroy(&(*ctx)->coo_values_full));
1569 PetscCallCeed((*ctx)->ceed, CeedVectorDestroy(&(*ctx)->coo_values_pbd));
1570 PetscCallCeed((*ctx)->ceed, CeedOperatorDestroy(&(*ctx)->op_mult));
1571 PetscCallCeed((*ctx)->ceed, CeedOperatorDestroy(&(*ctx)->op_mult_transpose));
1572 PetscCheck(CeedDestroy(&(*ctx)->ceed) == CEED_ERROR_SUCCESS, PETSC_COMM_SELF, PETSC_ERR_LIB, "destroying libCEED context object failed");
1573
1574 // Deallocate
1575 (*ctx)->is_destroyed = PETSC_TRUE; // Flag as destroyed in case someone has stale ref
1576 PetscCall(PetscFree(*ctx));
1577 PetscFunctionReturn(PETSC_SUCCESS);
1578 }
1579
1580 /**
1581 @brief Compute the diagonal of an operator via libCEED.
1582
1583 Collective across MPI processes.
1584
1585 @param[in] A `MATCEED`
1586 @param[out] D Vector holding operator diagonal
1587
1588 @return An error code: 0 - success, otherwise - failure
1589 **/
MatGetDiagonal_Ceed(Mat A,Vec D)1590 PetscErrorCode MatGetDiagonal_Ceed(Mat A, Vec D) {
1591 PetscMemType mem_type;
1592 Vec D_loc;
1593 MatCeedContext ctx;
1594
1595 PetscFunctionBeginUser;
1596 PetscCall(MatShellGetContext(A, &ctx));
1597
1598 // Place PETSc vector in libCEED vector
1599 PetscCall(PetscLogEventBegin(MATCEED_ASSEMBLE_DIAGONAL, A, D, NULL, NULL));
1600 PetscCall(DMGetLocalVector(ctx->dm_x, &D_loc));
1601 PetscCall(VecPetscToCeed(D_loc, &mem_type, ctx->x_loc));
1602
1603 // Compute Diagonal
1604 PetscCall(PetscLogEventBegin(MATCEED_ASSEMBLE_DIAGONAL_CEEDOP, A, D, NULL, NULL));
1605 PetscCallCeed(ctx->ceed, CeedOperatorLinearAssembleDiagonal(ctx->op_mult, ctx->x_loc, CEED_REQUEST_IMMEDIATE));
1606 PetscCall(PetscLogEventEnd(MATCEED_ASSEMBLE_DIAGONAL_CEEDOP, A, D, NULL, NULL));
1607
1608 // Restore PETSc vector
1609 PetscCall(VecCeedToPetsc(ctx->x_loc, mem_type, D_loc));
1610
1611 // Local-to-Global
1612 PetscCall(VecZeroEntries(D));
1613 PetscCall(DMLocalToGlobal(ctx->dm_x, D_loc, ADD_VALUES, D));
1614 PetscCall(DMRestoreLocalVector(ctx->dm_x, &D_loc));
1615 PetscCall(PetscLogEventEnd(MATCEED_ASSEMBLE_DIAGONAL, A, D, NULL, NULL));
1616 PetscFunctionReturn(PETSC_SUCCESS);
1617 }
1618
1619 /**
1620 @brief Compute `A X = Y` for a `MATCEED`.
1621
1622 Collective across MPI processes.
1623
1624 @param[in] A `MATCEED`
1625 @param[in] X Input PETSc vector
1626 @param[out] Y Output PETSc vector
1627
1628 @return An error code: 0 - success, otherwise - failure
1629 **/
MatMult_Ceed(Mat A,Vec X,Vec Y)1630 PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) {
1631 MatCeedContext ctx;
1632
1633 PetscFunctionBeginUser;
1634 PetscCall(MatShellGetContext(A, &ctx));
1635 PetscCall(PetscLogEventBegin(ctx->log_event_mult, A, X, Y, NULL));
1636
1637 {
1638 PetscMemType x_mem_type, y_mem_type;
1639 Vec X_loc = ctx->X_loc, Y_loc;
1640
1641 // Get local vectors
1642 if (!ctx->X_loc) PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc));
1643 PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc));
1644
1645 // Global-to-local
1646 PetscCall(DMGlobalToLocal(ctx->dm_x, X, INSERT_VALUES, X_loc));
1647
1648 // Setup libCEED vectors
1649 PetscCall(VecReadPetscToCeed(X_loc, &x_mem_type, ctx->x_loc));
1650 PetscCall(VecZeroEntries(Y_loc));
1651 PetscCall(VecPetscToCeed(Y_loc, &y_mem_type, ctx->y_loc));
1652
1653 // Apply libCEED operator
1654 PetscCall(PetscLogEventBegin(ctx->log_event_ceed_mult, A, X, Y, NULL));
1655 PetscCall(PetscLogGpuTimeBegin());
1656 PetscCallCeed(ctx->ceed, CeedOperatorApplyAdd(ctx->op_mult, ctx->x_loc, ctx->y_loc, CEED_REQUEST_IMMEDIATE));
1657 // Log flops
1658 if (PetscMemTypeDevice(ctx->mem_type)) PetscCall(PetscLogGpuFlops(ctx->flops_mult));
1659 else PetscCall(PetscLogFlops(ctx->flops_mult));
1660 PetscCall(PetscLogGpuTimeEnd());
1661 PetscCall(PetscLogEventEnd(ctx->log_event_ceed_mult, A, X, Y, NULL));
1662
1663 // Restore PETSc vectors
1664 PetscCall(VecReadCeedToPetsc(ctx->x_loc, x_mem_type, X_loc));
1665 PetscCall(VecCeedToPetsc(ctx->y_loc, y_mem_type, Y_loc));
1666
1667 // Local-to-global
1668 PetscCall(VecZeroEntries(Y));
1669 PetscCall(DMLocalToGlobal(ctx->dm_y, Y_loc, ADD_VALUES, Y));
1670
1671 // Restore local vectors, as needed
1672 if (!ctx->X_loc) PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc));
1673 PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc));
1674 }
1675
1676 PetscCall(PetscLogEventEnd(ctx->log_event_mult, A, X, Y, NULL));
1677 PetscFunctionReturn(PETSC_SUCCESS);
1678 }
1679
1680 /**
1681 @brief Compute `A^T Y = X` for a `MATCEED`.
1682
1683 Collective across MPI processes.
1684
1685 @param[in] A `MATCEED`
1686 @param[in] Y Input PETSc vector
1687 @param[out] X Output PETSc vector
1688
1689 @return An error code: 0 - success, otherwise - failure
1690 **/
MatMultTranspose_Ceed(Mat A,Vec Y,Vec X)1691 PetscErrorCode MatMultTranspose_Ceed(Mat A, Vec Y, Vec X) {
1692 MatCeedContext ctx;
1693
1694 PetscFunctionBeginUser;
1695 PetscCall(MatShellGetContext(A, &ctx));
1696 PetscCall(PetscLogEventBegin(ctx->log_event_mult_transpose, A, Y, X, NULL));
1697
1698 {
1699 PetscMemType x_mem_type, y_mem_type;
1700 Vec X_loc, Y_loc = ctx->Y_loc_transpose;
1701
1702 // Get local vectors
1703 if (!ctx->Y_loc_transpose) PetscCall(DMGetLocalVector(ctx->dm_y, &Y_loc));
1704 PetscCall(DMGetLocalVector(ctx->dm_x, &X_loc));
1705
1706 // Global-to-local
1707 PetscCall(DMGlobalToLocal(ctx->dm_y, Y, INSERT_VALUES, Y_loc));
1708
1709 // Setup libCEED vectors
1710 PetscCall(VecReadPetscToCeed(Y_loc, &y_mem_type, ctx->y_loc));
1711 PetscCall(VecZeroEntries(X_loc));
1712 PetscCall(VecPetscToCeed(X_loc, &x_mem_type, ctx->x_loc));
1713
1714 // Apply libCEED operator
1715 PetscCall(PetscLogEventBegin(ctx->log_event_ceed_mult_transpose, A, Y, X, NULL));
1716 PetscCall(PetscLogGpuTimeBegin());
1717 PetscCallCeed(ctx->ceed, CeedOperatorApplyAdd(ctx->op_mult_transpose, ctx->y_loc, ctx->x_loc, CEED_REQUEST_IMMEDIATE));
1718 // Log flops
1719 if (PetscMemTypeDevice(ctx->mem_type)) PetscCall(PetscLogGpuFlops(ctx->flops_mult_transpose));
1720 else PetscCall(PetscLogFlops(ctx->flops_mult_transpose));
1721 PetscCall(PetscLogGpuTimeEnd());
1722 PetscCall(PetscLogEventEnd(ctx->log_event_ceed_mult_transpose, A, Y, X, NULL));
1723
1724 // Restore PETSc vectors
1725 PetscCall(VecReadCeedToPetsc(ctx->y_loc, y_mem_type, Y_loc));
1726 PetscCall(VecCeedToPetsc(ctx->x_loc, x_mem_type, X_loc));
1727
1728 // Local-to-global
1729 PetscCall(VecZeroEntries(X));
1730 PetscCall(DMLocalToGlobal(ctx->dm_x, X_loc, ADD_VALUES, X));
1731
1732 // Restore local vectors, as needed
1733 if (!ctx->Y_loc_transpose) PetscCall(DMRestoreLocalVector(ctx->dm_y, &Y_loc));
1734 PetscCall(DMRestoreLocalVector(ctx->dm_x, &X_loc));
1735 }
1736
1737 PetscCall(PetscLogEventEnd(ctx->log_event_mult_transpose, A, Y, X, NULL));
1738 PetscFunctionReturn(PETSC_SUCCESS);
1739 }
1740