xref: /libCEED/examples/petsc/src/matops.c (revision 6c88e6a2c47a72cd5c3efacc26a364b9cb7d40db)
1 #include "../include/matops.h"
2 #include "../include/petscutils.h"
3 
4 // -----------------------------------------------------------------------------
5 // Setup apply operator context data
6 // -----------------------------------------------------------------------------
7 PetscErrorCode SetupApplyOperatorCtx(MPI_Comm comm, DM dm, Ceed ceed,
8                                      CeedData ceed_data, Vec X_loc,
9                                      OperatorApplyContext op_apply_ctx) {
10   PetscErrorCode ierr;
11   PetscFunctionBeginUser;
12 
13   op_apply_ctx->comm = comm;
14   op_apply_ctx->dm = dm;
15   op_apply_ctx->X_loc = X_loc;
16   ierr = VecDuplicate(X_loc, &op_apply_ctx->Y_loc); CHKERRQ(ierr);
17   op_apply_ctx->x_ceed = ceed_data->x_ceed;
18   op_apply_ctx->y_ceed = ceed_data->y_ceed;
19   op_apply_ctx->op = ceed_data->op_apply;
20   op_apply_ctx->ceed = ceed;
21   PetscFunctionReturn(0);
22 }
23 
24 // -----------------------------------------------------------------------------
25 // Setup error operator context data
26 // -----------------------------------------------------------------------------
27 PetscErrorCode SetupErrorOperatorCtx(MPI_Comm comm, DM dm, Ceed ceed,
28                                      CeedData ceed_data, Vec X_loc, CeedOperator op_error,
29                                      OperatorApplyContext op_error_ctx) {
30   PetscErrorCode ierr;
31   PetscFunctionBeginUser;
32 
33   op_error_ctx->comm = comm;
34   op_error_ctx->dm = dm;
35   op_error_ctx->X_loc = X_loc;
36   ierr = VecDuplicate(X_loc, &op_error_ctx->Y_loc); CHKERRQ(ierr);
37   op_error_ctx->x_ceed = ceed_data->x_ceed;
38   op_error_ctx->y_ceed = ceed_data->y_ceed;
39   op_error_ctx->op = op_error;
40   op_error_ctx->ceed = ceed;
41   PetscFunctionReturn(0);
42 }
43 
44 // -----------------------------------------------------------------------------
45 // This function returns the computed diagonal of the operator
46 // -----------------------------------------------------------------------------
47 PetscErrorCode MatGetDiag(Mat A, Vec D) {
48   PetscErrorCode ierr;
49   OperatorApplyContext op_apply_ctx;
50 
51   PetscFunctionBeginUser;
52 
53   ierr = MatShellGetContext(A, &op_apply_ctx); CHKERRQ(ierr);
54 
55   // Compute Diagonal via libCEED
56   PetscScalar *y;
57   PetscMemType mem_type;
58 
59   // -- Place PETSc vector in libCEED vector
60   ierr = VecGetArrayAndMemType(op_apply_ctx->Y_loc, &y, &mem_type); CHKERRQ(ierr);
61   CeedVectorSetArray(op_apply_ctx->y_ceed, MemTypeP2C(mem_type),
62                      CEED_USE_POINTER, y);
63 
64   // -- Compute Diagonal
65   CeedOperatorLinearAssembleDiagonal(op_apply_ctx->op, op_apply_ctx->y_ceed,
66                                      CEED_REQUEST_IMMEDIATE);
67 
68   // -- Local-to-Global
69   CeedVectorTakeArray(op_apply_ctx->y_ceed, MemTypeP2C(mem_type), NULL);
70   ierr = VecRestoreArrayAndMemType(op_apply_ctx->Y_loc, &y); CHKERRQ(ierr);
71   ierr = VecZeroEntries(D); CHKERRQ(ierr);
72   ierr = DMLocalToGlobal(op_apply_ctx->dm, op_apply_ctx->Y_loc, ADD_VALUES, D);
73   CHKERRQ(ierr);
74 
75   PetscFunctionReturn(0);
76 };
77 
78 // -----------------------------------------------------------------------------
79 // This function uses libCEED to compute the action of the Laplacian with
80 // Dirichlet boundary conditions
81 // -----------------------------------------------------------------------------
82 PetscErrorCode ApplyLocal_Ceed(Vec X, Vec Y,
83                                OperatorApplyContext op_apply_ctx) {
84   PetscErrorCode ierr;
85   PetscScalar *x, *y;
86   PetscMemType x_mem_type, y_mem_type;
87 
88   PetscFunctionBeginUser;
89 
90   // Global-to-local
91   ierr = DMGlobalToLocal(op_apply_ctx->dm, X, INSERT_VALUES, op_apply_ctx->X_loc);
92   CHKERRQ(ierr);
93 
94   // Setup libCEED vectors
95   ierr = VecGetArrayReadAndMemType(op_apply_ctx->X_loc, (const PetscScalar **)&x,
96                                    &x_mem_type); CHKERRQ(ierr);
97   ierr = VecGetArrayAndMemType(op_apply_ctx->Y_loc, &y, &y_mem_type);
98   CHKERRQ(ierr);
99   CeedVectorSetArray(op_apply_ctx->x_ceed, MemTypeP2C(x_mem_type),
100                      CEED_USE_POINTER, x);
101   CeedVectorSetArray(op_apply_ctx->y_ceed, MemTypeP2C(y_mem_type),
102                      CEED_USE_POINTER, y);
103 
104   // Apply libCEED operator
105   CeedOperatorApply(op_apply_ctx->op, op_apply_ctx->x_ceed, op_apply_ctx->y_ceed,
106                     CEED_REQUEST_IMMEDIATE);
107 
108   // Restore PETSc vectors
109   CeedVectorTakeArray(op_apply_ctx->x_ceed, MemTypeP2C(x_mem_type), NULL);
110   CeedVectorTakeArray(op_apply_ctx->y_ceed, MemTypeP2C(y_mem_type), NULL);
111   ierr = VecRestoreArrayReadAndMemType(op_apply_ctx->X_loc,
112                                        (const PetscScalar **)&x);
113   CHKERRQ(ierr);
114   ierr = VecRestoreArrayAndMemType(op_apply_ctx->Y_loc, &y); CHKERRQ(ierr);
115 
116   // Local-to-global
117   ierr = VecZeroEntries(Y); CHKERRQ(ierr);
118   ierr = DMLocalToGlobal(op_apply_ctx->dm, op_apply_ctx->Y_loc, ADD_VALUES, Y);
119   CHKERRQ(ierr);
120 
121   PetscFunctionReturn(0);
122 };
123 
124 // -----------------------------------------------------------------------------
125 // This function wraps the libCEED operator for a MatShell
126 // -----------------------------------------------------------------------------
127 PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) {
128   PetscErrorCode ierr;
129   OperatorApplyContext op_apply_ctx;
130 
131   PetscFunctionBeginUser;
132 
133   ierr = MatShellGetContext(A, &op_apply_ctx); CHKERRQ(ierr);
134 
135   // libCEED for local action of residual evaluator
136   ierr = ApplyLocal_Ceed(X, Y, op_apply_ctx); CHKERRQ(ierr);
137 
138   PetscFunctionReturn(0);
139 };
140 
141 // -----------------------------------------------------------------------------
142 // This function wraps the libCEED operator for a SNES residual evaluation
143 // -----------------------------------------------------------------------------
144 PetscErrorCode FormResidual_Ceed(SNES snes, Vec X, Vec Y, void *ctx) {
145   PetscErrorCode ierr;
146   OperatorApplyContext op_apply_ctx = (OperatorApplyContext)ctx;
147 
148   PetscFunctionBeginUser;
149 
150   // libCEED for local action of residual evaluator
151   ierr = ApplyLocal_Ceed(X, Y, op_apply_ctx); CHKERRQ(ierr);
152 
153   PetscFunctionReturn(0);
154 };
155 
156 // -----------------------------------------------------------------------------
157 // This function uses libCEED to compute the action of the prolongation operator
158 // -----------------------------------------------------------------------------
159 PetscErrorCode MatMult_Prolong(Mat A, Vec X, Vec Y) {
160   PetscErrorCode ierr;
161   ProlongRestrContext pr_restr_ctx;
162   PetscScalar *c, *f;
163   PetscMemType c_mem_type, f_mem_type;
164 
165   PetscFunctionBeginUser;
166 
167   ierr = MatShellGetContext(A, &pr_restr_ctx); CHKERRQ(ierr);
168 
169   // Global-to-local
170   ierr = VecZeroEntries(pr_restr_ctx->loc_vec_c); CHKERRQ(ierr);
171   ierr = DMGlobalToLocal(pr_restr_ctx->dmc, X, INSERT_VALUES,
172                          pr_restr_ctx->loc_vec_c);
173   CHKERRQ(ierr);
174 
175   // Setup libCEED vectors
176   ierr = VecGetArrayReadAndMemType(pr_restr_ctx->loc_vec_c,
177                                    (const PetscScalar **)&c,
178                                    &c_mem_type); CHKERRQ(ierr);
179   ierr = VecGetArrayAndMemType(pr_restr_ctx->loc_vec_f, &f, &f_mem_type);
180   CHKERRQ(ierr);
181   CeedVectorSetArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type),
182                      CEED_USE_POINTER, c);
183   CeedVectorSetArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type),
184                      CEED_USE_POINTER, f);
185 
186   // Apply libCEED operator
187   CeedOperatorApply(pr_restr_ctx->op_prolong, pr_restr_ctx->ceed_vec_c,
188                     pr_restr_ctx->ceed_vec_f, CEED_REQUEST_IMMEDIATE);
189 
190   // Restore PETSc vectors
191   CeedVectorTakeArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type), NULL);
192   CeedVectorTakeArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type), NULL);
193   ierr = VecRestoreArrayReadAndMemType(pr_restr_ctx->loc_vec_c,
194                                        (const PetscScalar **)&c);
195   CHKERRQ(ierr);
196   ierr = VecRestoreArrayAndMemType(pr_restr_ctx->loc_vec_f, &f); CHKERRQ(ierr);
197 
198   // Multiplicity
199   ierr = VecPointwiseMult(pr_restr_ctx->loc_vec_f, pr_restr_ctx->loc_vec_f,
200                           pr_restr_ctx->mult_vec);
201 
202   // Local-to-global
203   ierr = VecZeroEntries(Y); CHKERRQ(ierr);
204   ierr = DMLocalToGlobal(pr_restr_ctx->dmf, pr_restr_ctx->loc_vec_f, ADD_VALUES,
205                          Y); CHKERRQ(ierr);
206 
207   PetscFunctionReturn(0);
208 };
209 
210 // -----------------------------------------------------------------------------
211 // This function uses libCEED to compute the action of the restriction operator
212 // -----------------------------------------------------------------------------
213 PetscErrorCode MatMult_Restrict(Mat A, Vec X, Vec Y) {
214   PetscErrorCode ierr;
215   ProlongRestrContext pr_restr_ctx;
216   PetscScalar *c, *f;
217   PetscMemType c_mem_type, f_mem_type;
218 
219   PetscFunctionBeginUser;
220 
221   ierr = MatShellGetContext(A, &pr_restr_ctx); CHKERRQ(ierr);
222 
223   // Global-to-local
224   ierr = VecZeroEntries(pr_restr_ctx->loc_vec_f); CHKERRQ(ierr);
225   ierr = DMGlobalToLocal(pr_restr_ctx->dmf, X, INSERT_VALUES,
226                          pr_restr_ctx->loc_vec_f);
227   CHKERRQ(ierr);
228 
229   // Multiplicity
230   ierr = VecPointwiseMult(pr_restr_ctx->loc_vec_f, pr_restr_ctx->loc_vec_f,
231                           pr_restr_ctx->mult_vec);
232   CHKERRQ(ierr);
233 
234   // Setup libCEED vectors
235   ierr = VecGetArrayReadAndMemType(pr_restr_ctx->loc_vec_f,
236                                    (const PetscScalar **)&f,
237                                    &f_mem_type); CHKERRQ(ierr);
238   ierr = VecGetArrayAndMemType(pr_restr_ctx->loc_vec_c, &c, &c_mem_type);
239   CHKERRQ(ierr);
240   CeedVectorSetArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type),
241                      CEED_USE_POINTER, f);
242   CeedVectorSetArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type),
243                      CEED_USE_POINTER, c);
244 
245   // Apply CEED operator
246   CeedOperatorApply(pr_restr_ctx->op_restrict, pr_restr_ctx->ceed_vec_f,
247                     pr_restr_ctx->ceed_vec_c, CEED_REQUEST_IMMEDIATE);
248 
249   // Restore PETSc vectors
250   CeedVectorTakeArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type), NULL);
251   CeedVectorTakeArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type), NULL);
252   ierr = VecRestoreArrayReadAndMemType(pr_restr_ctx->loc_vec_f,
253                                        (const PetscScalar **)&f); CHKERRQ(ierr);
254   ierr = VecRestoreArrayAndMemType(pr_restr_ctx->loc_vec_c, &c); CHKERRQ(ierr);
255 
256   // Local-to-global
257   ierr = VecZeroEntries(Y); CHKERRQ(ierr);
258   ierr = DMLocalToGlobal(pr_restr_ctx->dmc, pr_restr_ctx->loc_vec_c, ADD_VALUES,
259                          Y);
260   CHKERRQ(ierr);
261 
262   PetscFunctionReturn(0);
263 };
264 
265 // -----------------------------------------------------------------------------
266 // This function calculates the error in the final solution
267 // -----------------------------------------------------------------------------
268 PetscErrorCode ComputeErrorMax(OperatorApplyContext op_error_ctx,
269                                Vec X, CeedVector target,
270                                PetscScalar *max_error) {
271   PetscErrorCode ierr;
272   PetscScalar *x;
273   PetscMemType mem_type;
274   CeedVector collocated_error;
275   CeedSize length;
276 
277   PetscFunctionBeginUser;
278   CeedVectorGetLength(target, &length);
279   CeedVectorCreate(op_error_ctx->ceed, length, &collocated_error);
280 
281   // Global-to-local
282   ierr = DMGlobalToLocal(op_error_ctx->dm, X, INSERT_VALUES, op_error_ctx->X_loc);
283   CHKERRQ(ierr);
284 
285   // Setup CEED vector
286   ierr = VecGetArrayAndMemType(op_error_ctx->X_loc, &x, &mem_type); CHKERRQ(ierr);
287   CeedVectorSetArray(op_error_ctx->x_ceed, MemTypeP2C(mem_type),
288                      CEED_USE_POINTER, x);
289 
290   // Apply CEED operator
291   CeedOperatorApply(op_error_ctx->op, op_error_ctx->x_ceed, collocated_error,
292                     CEED_REQUEST_IMMEDIATE);
293 
294   // Restore PETSc vector
295   CeedVectorTakeArray(op_error_ctx->x_ceed, MemTypeP2C(mem_type), NULL);
296   ierr = VecRestoreArrayReadAndMemType(op_error_ctx->X_loc,
297                                        (const PetscScalar **)&x);
298   CHKERRQ(ierr);
299 
300   // Reduce max error
301   *max_error = 0;
302   const CeedScalar *e;
303   CeedVectorGetArrayRead(collocated_error, CEED_MEM_HOST, &e);
304   for (CeedInt i=0; i<length; i++) {
305     *max_error = PetscMax(*max_error, PetscAbsScalar(e[i]));
306   }
307   CeedVectorRestoreArrayRead(collocated_error, &e);
308   ierr = MPI_Allreduce(MPI_IN_PLACE, max_error, 1, MPIU_REAL, MPIU_MAX,
309                        op_error_ctx->comm);
310   CHKERRQ(ierr);
311 
312   // Cleanup
313   CeedVectorDestroy(&collocated_error);
314 
315   PetscFunctionReturn(0);
316 };
317 
318 // -----------------------------------------------------------------------------
319