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