xref: /libCEED/examples/petsc/src/matops.c (revision ab1ff315f74846d3fb318e101473fbae03c37843)
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 uses libCEED to compute the action of the prolongation operator
143 // -----------------------------------------------------------------------------
144 PetscErrorCode MatMult_Prolong(Mat A, Vec X, Vec Y) {
145   PetscErrorCode ierr;
146   ProlongRestrContext pr_restr_ctx;
147   PetscScalar *c, *f;
148   PetscMemType c_mem_type, f_mem_type;
149 
150   PetscFunctionBeginUser;
151 
152   ierr = MatShellGetContext(A, &pr_restr_ctx); CHKERRQ(ierr);
153 
154   // Global-to-local
155   ierr = VecZeroEntries(pr_restr_ctx->loc_vec_c); CHKERRQ(ierr);
156   ierr = DMGlobalToLocal(pr_restr_ctx->dmc, X, INSERT_VALUES,
157                          pr_restr_ctx->loc_vec_c);
158   CHKERRQ(ierr);
159 
160   // Setup libCEED vectors
161   ierr = VecGetArrayReadAndMemType(pr_restr_ctx->loc_vec_c,
162                                    (const PetscScalar **)&c,
163                                    &c_mem_type); CHKERRQ(ierr);
164   ierr = VecGetArrayAndMemType(pr_restr_ctx->loc_vec_f, &f, &f_mem_type);
165   CHKERRQ(ierr);
166   CeedVectorSetArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type),
167                      CEED_USE_POINTER, c);
168   CeedVectorSetArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type),
169                      CEED_USE_POINTER, f);
170 
171   // Apply libCEED operator
172   CeedOperatorApply(pr_restr_ctx->op_prolong, pr_restr_ctx->ceed_vec_c,
173                     pr_restr_ctx->ceed_vec_f, CEED_REQUEST_IMMEDIATE);
174 
175   // Restore PETSc vectors
176   CeedVectorTakeArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type), NULL);
177   CeedVectorTakeArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type), NULL);
178   ierr = VecRestoreArrayReadAndMemType(pr_restr_ctx->loc_vec_c,
179                                        (const PetscScalar **)&c);
180   CHKERRQ(ierr);
181   ierr = VecRestoreArrayAndMemType(pr_restr_ctx->loc_vec_f, &f); CHKERRQ(ierr);
182 
183   // Multiplicity
184   ierr = VecPointwiseMult(pr_restr_ctx->loc_vec_f, pr_restr_ctx->loc_vec_f,
185                           pr_restr_ctx->mult_vec);
186 
187   // Local-to-global
188   ierr = VecZeroEntries(Y); CHKERRQ(ierr);
189   ierr = DMLocalToGlobal(pr_restr_ctx->dmf, pr_restr_ctx->loc_vec_f, ADD_VALUES,
190                          Y); CHKERRQ(ierr);
191 
192   PetscFunctionReturn(0);
193 };
194 
195 // -----------------------------------------------------------------------------
196 // This function uses libCEED to compute the action of the restriction operator
197 // -----------------------------------------------------------------------------
198 PetscErrorCode MatMult_Restrict(Mat A, Vec X, Vec Y) {
199   PetscErrorCode ierr;
200   ProlongRestrContext pr_restr_ctx;
201   PetscScalar *c, *f;
202   PetscMemType c_mem_type, f_mem_type;
203 
204   PetscFunctionBeginUser;
205 
206   ierr = MatShellGetContext(A, &pr_restr_ctx); CHKERRQ(ierr);
207 
208   // Global-to-local
209   ierr = VecZeroEntries(pr_restr_ctx->loc_vec_f); CHKERRQ(ierr);
210   ierr = DMGlobalToLocal(pr_restr_ctx->dmf, X, INSERT_VALUES,
211                          pr_restr_ctx->loc_vec_f);
212   CHKERRQ(ierr);
213 
214   // Multiplicity
215   ierr = VecPointwiseMult(pr_restr_ctx->loc_vec_f, pr_restr_ctx->loc_vec_f,
216                           pr_restr_ctx->mult_vec);
217   CHKERRQ(ierr);
218 
219   // Setup libCEED vectors
220   ierr = VecGetArrayReadAndMemType(pr_restr_ctx->loc_vec_f,
221                                    (const PetscScalar **)&f,
222                                    &f_mem_type); CHKERRQ(ierr);
223   ierr = VecGetArrayAndMemType(pr_restr_ctx->loc_vec_c, &c, &c_mem_type);
224   CHKERRQ(ierr);
225   CeedVectorSetArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type),
226                      CEED_USE_POINTER, f);
227   CeedVectorSetArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type),
228                      CEED_USE_POINTER, c);
229 
230   // Apply CEED operator
231   CeedOperatorApply(pr_restr_ctx->op_restrict, pr_restr_ctx->ceed_vec_f,
232                     pr_restr_ctx->ceed_vec_c, CEED_REQUEST_IMMEDIATE);
233 
234   // Restore PETSc vectors
235   CeedVectorTakeArray(pr_restr_ctx->ceed_vec_c, MemTypeP2C(c_mem_type), NULL);
236   CeedVectorTakeArray(pr_restr_ctx->ceed_vec_f, MemTypeP2C(f_mem_type), NULL);
237   ierr = VecRestoreArrayReadAndMemType(pr_restr_ctx->loc_vec_f,
238                                        (const PetscScalar **)&f); CHKERRQ(ierr);
239   ierr = VecRestoreArrayAndMemType(pr_restr_ctx->loc_vec_c, &c); CHKERRQ(ierr);
240 
241   // Local-to-global
242   ierr = VecZeroEntries(Y); CHKERRQ(ierr);
243   ierr = DMLocalToGlobal(pr_restr_ctx->dmc, pr_restr_ctx->loc_vec_c, ADD_VALUES,
244                          Y);
245   CHKERRQ(ierr);
246 
247   PetscFunctionReturn(0);
248 };
249 
250 // -----------------------------------------------------------------------------
251 // This function calculates the error in the final solution
252 // -----------------------------------------------------------------------------
253 PetscErrorCode ComputeL2Error(Vec X, PetscScalar *l2_error,
254                               OperatorApplyContext op_error_ctx) {
255 
256   Vec        E;
257   PetscFunctionBeginUser;
258   PetscCall(VecDuplicate(X, &E));
259   PetscCall(ApplyLocal_Ceed(X, E, op_error_ctx) );
260   PetscScalar error_sq = 1.0;
261   PetscCall(VecSum(E, &error_sq));
262   *l2_error = sqrt(error_sq);
263 
264   PetscFunctionReturn(0);
265 };
266 
267 // -----------------------------------------------------------------------------
268