xref: /libCEED/backends/blocked/ceed-blocked-operator.c (revision 0305e208bc4684eaa3113a3ca8b7ffdd1d4cb7a9)
1 // Copyright (c) 2017-2022, Lawrence Livermore National Security, LLC and other CEED contributors.
2 // All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
3 //
4 // SPDX-License-Identifier: BSD-2-Clause
5 //
6 // This file is part of CEED:  http://github.com/ceed
7 
8 #include <ceed.h>
9 #include <ceed/backend.h>
10 #include <stdbool.h>
11 #include <stddef.h>
12 #include <stdint.h>
13 
14 #include "ceed-blocked.h"
15 
16 //------------------------------------------------------------------------------
17 // Setup Input/Output Fields
18 //------------------------------------------------------------------------------
19 static int CeedOperatorSetupFields_Blocked(CeedQFunction qf, CeedOperator op, bool is_input, CeedElemRestriction *blk_restr, CeedVector *e_vecs_full,
20                                            CeedVector *e_vecs, CeedVector *q_vecs, CeedInt start_e, CeedInt num_fields, CeedInt Q) {
21   CeedInt  num_comp, size, P;
22   CeedSize e_size, q_size;
23   Ceed     ceed;
24   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
25   CeedBasis           basis;
26   CeedElemRestriction r;
27   CeedOperatorField  *op_fields;
28   CeedQFunctionField *qf_fields;
29   if (is_input) {
30     CeedCallBackend(CeedOperatorGetFields(op, NULL, &op_fields, NULL, NULL));
31     CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_fields, NULL, NULL));
32   } else {
33     CeedCallBackend(CeedOperatorGetFields(op, NULL, NULL, NULL, &op_fields));
34     CeedCallBackend(CeedQFunctionGetFields(qf, NULL, NULL, NULL, &qf_fields));
35   }
36   const CeedInt blk_size = 8;
37 
38   // Loop over fields
39   for (CeedInt i = 0; i < num_fields; i++) {
40     CeedEvalMode eval_mode;
41     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_fields[i], &eval_mode));
42 
43     if (eval_mode != CEED_EVAL_WEIGHT) {
44       CeedCallBackend(CeedOperatorFieldGetElemRestriction(op_fields[i], &r));
45       CeedCallBackend(CeedElemRestrictionGetCeed(r, &ceed));
46       CeedSize l_size;
47       CeedInt  num_elem, elem_size, comp_stride;
48       CeedCallBackend(CeedElemRestrictionGetNumElements(r, &num_elem));
49       CeedCallBackend(CeedElemRestrictionGetElementSize(r, &elem_size));
50       CeedCallBackend(CeedElemRestrictionGetLVectorSize(r, &l_size));
51       CeedCallBackend(CeedElemRestrictionGetNumComponents(r, &num_comp));
52       CeedCallBackend(CeedElemRestrictionGetCompStride(r, &comp_stride));
53 
54       CeedRestrictionType rstr_type;
55       CeedCallBackend(CeedElemRestrictionGetType(r, &rstr_type));
56       switch (rstr_type) {
57         case CEED_RESTRICTION_DEFAULT: {
58           const CeedInt *offsets = NULL;
59           CeedCallBackend(CeedElemRestrictionGetOffsets(r, CEED_MEM_HOST, &offsets));
60           CeedCallBackend(CeedElemRestrictionCreateBlocked(ceed, num_elem, elem_size, blk_size, num_comp, comp_stride, l_size, CEED_MEM_HOST,
61                                                            CEED_COPY_VALUES, offsets, &blk_restr[i + start_e]));
62           CeedCallBackend(CeedElemRestrictionRestoreOffsets(r, &offsets));
63         } break;
64         case CEED_RESTRICTION_ORIENTED: {
65           const CeedInt *offsets = NULL;
66           const bool    *orients = NULL;
67           CeedCallBackend(CeedElemRestrictionGetOffsets(r, CEED_MEM_HOST, &offsets));
68           CeedCallBackend(CeedElemRestrictionGetOrientations(r, CEED_MEM_HOST, &orients));
69           CeedCallBackend(CeedElemRestrictionCreateBlockedOriented(ceed, num_elem, elem_size, blk_size, num_comp, comp_stride, l_size, CEED_MEM_HOST,
70                                                                    CEED_COPY_VALUES, offsets, orients, &blk_restr[i + start_e]));
71           CeedCallBackend(CeedElemRestrictionRestoreOffsets(r, &offsets));
72           CeedCallBackend(CeedElemRestrictionRestoreOrientations(r, &orients));
73         } break;
74         case CEED_RESTRICTION_CURL_ORIENTED: {
75           const CeedInt *offsets      = NULL;
76           const CeedInt *curl_orients = NULL;
77           CeedCallBackend(CeedElemRestrictionGetOffsets(r, CEED_MEM_HOST, &offsets));
78           CeedCallBackend(CeedElemRestrictionGetCurlOrientations(r, CEED_MEM_HOST, &curl_orients));
79           CeedCallBackend(CeedElemRestrictionCreateBlockedCurlOriented(ceed, num_elem, elem_size, blk_size, num_comp, comp_stride, l_size,
80                                                                        CEED_MEM_HOST, CEED_COPY_VALUES, offsets, curl_orients,
81                                                                        &blk_restr[i + start_e]));
82           CeedCallBackend(CeedElemRestrictionRestoreOffsets(r, &offsets));
83           CeedCallBackend(CeedElemRestrictionRestoreCurlOrientations(r, &curl_orients));
84         } break;
85         case CEED_RESTRICTION_STRIDED: {
86           CeedInt strides[3];
87           CeedCallBackend(CeedElemRestrictionGetStrides(r, &strides));
88           CeedCallBackend(
89               CeedElemRestrictionCreateBlockedStrided(ceed, num_elem, elem_size, blk_size, num_comp, l_size, strides, &blk_restr[i + start_e]));
90         } break;
91       }
92       CeedCallBackend(CeedElemRestrictionCreateVector(blk_restr[i + start_e], NULL, &e_vecs_full[i + start_e]));
93     }
94 
95     switch (eval_mode) {
96       case CEED_EVAL_NONE:
97         CeedCallBackend(CeedQFunctionFieldGetSize(qf_fields[i], &size));
98         q_size = (CeedSize)Q * size * blk_size;
99         CeedCallBackend(CeedVectorCreate(ceed, q_size, &q_vecs[i]));
100         break;
101       case CEED_EVAL_INTERP:
102       case CEED_EVAL_GRAD:
103       case CEED_EVAL_DIV:
104       case CEED_EVAL_CURL:
105         CeedCallBackend(CeedOperatorFieldGetBasis(op_fields[i], &basis));
106         CeedCallBackend(CeedQFunctionFieldGetSize(qf_fields[i], &size));
107         CeedCallBackend(CeedBasisGetNumNodes(basis, &P));
108         CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
109         e_size = (CeedSize)P * num_comp * blk_size;
110         CeedCallBackend(CeedVectorCreate(ceed, e_size, &e_vecs[i]));
111         q_size = (CeedSize)Q * size * blk_size;
112         CeedCallBackend(CeedVectorCreate(ceed, q_size, &q_vecs[i]));
113         break;
114       case CEED_EVAL_WEIGHT:  // Only on input fields
115         CeedCallBackend(CeedOperatorFieldGetBasis(op_fields[i], &basis));
116         q_size = (CeedSize)Q * blk_size;
117         CeedCallBackend(CeedVectorCreate(ceed, q_size, &q_vecs[i]));
118         CeedCallBackend(CeedBasisApply(basis, blk_size, CEED_NOTRANSPOSE, CEED_EVAL_WEIGHT, CEED_VECTOR_NONE, q_vecs[i]));
119         break;
120     }
121   }
122   return CEED_ERROR_SUCCESS;
123 }
124 
125 //------------------------------------------------------------------------------
126 // Setup Operator
127 //------------------------------------------------------------------------------
128 static int CeedOperatorSetup_Blocked(CeedOperator op) {
129   bool is_setup_done;
130   CeedCallBackend(CeedOperatorIsSetupDone(op, &is_setup_done));
131   if (is_setup_done) return CEED_ERROR_SUCCESS;
132   Ceed ceed;
133   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
134   CeedOperator_Blocked *impl;
135   CeedCallBackend(CeedOperatorGetData(op, &impl));
136   CeedQFunction qf;
137   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
138   CeedInt Q, num_input_fields, num_output_fields;
139   CeedCallBackend(CeedOperatorGetNumQuadraturePoints(op, &Q));
140   CeedCallBackend(CeedQFunctionIsIdentity(qf, &impl->is_identity_qf));
141   CeedOperatorField *op_input_fields, *op_output_fields;
142   CeedCallBackend(CeedOperatorGetFields(op, &num_input_fields, &op_input_fields, &num_output_fields, &op_output_fields));
143   CeedQFunctionField *qf_input_fields, *qf_output_fields;
144   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_input_fields, NULL, &qf_output_fields));
145 
146   // Allocate
147   CeedCallBackend(CeedCalloc(num_input_fields + num_output_fields, &impl->blk_restr));
148   CeedCallBackend(CeedCalloc(num_input_fields + num_output_fields, &impl->e_vecs_full));
149 
150   CeedCallBackend(CeedCalloc(CEED_FIELD_MAX, &impl->input_states));
151   CeedCallBackend(CeedCalloc(CEED_FIELD_MAX, &impl->e_vecs_in));
152   CeedCallBackend(CeedCalloc(CEED_FIELD_MAX, &impl->e_vecs_out));
153   CeedCallBackend(CeedCalloc(CEED_FIELD_MAX, &impl->q_vecs_in));
154   CeedCallBackend(CeedCalloc(CEED_FIELD_MAX, &impl->q_vecs_out));
155 
156   impl->num_inputs  = num_input_fields;
157   impl->num_outputs = num_output_fields;
158 
159   // Set up infield and outfield pointer arrays
160   // Infields
161   CeedCallBackend(
162       CeedOperatorSetupFields_Blocked(qf, op, true, impl->blk_restr, impl->e_vecs_full, impl->e_vecs_in, impl->q_vecs_in, 0, num_input_fields, Q));
163   // Outfields
164   CeedCallBackend(CeedOperatorSetupFields_Blocked(qf, op, false, impl->blk_restr, impl->e_vecs_full, impl->e_vecs_out, impl->q_vecs_out,
165                                                   num_input_fields, num_output_fields, Q));
166 
167   // Identity QFunctions
168   if (impl->is_identity_qf) {
169     CeedEvalMode        in_mode, out_mode;
170     CeedQFunctionField *in_fields, *out_fields;
171     CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &in_fields, NULL, &out_fields));
172     CeedCallBackend(CeedQFunctionFieldGetEvalMode(in_fields[0], &in_mode));
173     CeedCallBackend(CeedQFunctionFieldGetEvalMode(out_fields[0], &out_mode));
174 
175     if (in_mode == CEED_EVAL_NONE && out_mode == CEED_EVAL_NONE) {
176       impl->is_identity_restr_op = true;
177     } else {
178       CeedCallBackend(CeedVectorReferenceCopy(impl->q_vecs_in[0], &impl->q_vecs_out[0]));
179     }
180   }
181 
182   CeedCallBackend(CeedOperatorSetSetupDone(op));
183 
184   return CEED_ERROR_SUCCESS;
185 }
186 
187 //------------------------------------------------------------------------------
188 // Setup Operator Inputs
189 //------------------------------------------------------------------------------
190 static inline int CeedOperatorSetupInputs_Blocked(CeedInt num_input_fields, CeedQFunctionField *qf_input_fields, CeedOperatorField *op_input_fields,
191                                                   CeedVector in_vec, bool skip_active, CeedScalar *e_data_full[2 * CEED_FIELD_MAX],
192                                                   CeedOperator_Blocked *impl, CeedRequest *request) {
193   CeedEvalMode eval_mode;
194   CeedVector   vec;
195   uint64_t     state;
196 
197   for (CeedInt i = 0; i < num_input_fields; i++) {
198     // Get input vector
199     CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
200     if (vec == CEED_VECTOR_ACTIVE) {
201       if (skip_active) continue;
202       else vec = in_vec;
203     }
204 
205     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
206     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
207     } else {
208       // Restrict
209       CeedCallBackend(CeedVectorGetState(vec, &state));
210       if (state != impl->input_states[i] || vec == in_vec) {
211         CeedCallBackend(CeedElemRestrictionApply(impl->blk_restr[i], CEED_NOTRANSPOSE, vec, impl->e_vecs_full[i], request));
212         impl->input_states[i] = state;
213       }
214       // Get evec
215       CeedCallBackend(CeedVectorGetArrayRead(impl->e_vecs_full[i], CEED_MEM_HOST, (const CeedScalar **)&e_data_full[i]));
216     }
217   }
218   return CEED_ERROR_SUCCESS;
219 }
220 
221 //------------------------------------------------------------------------------
222 // Input Basis Action
223 //------------------------------------------------------------------------------
224 static inline int CeedOperatorInputBasis_Blocked(CeedInt e, CeedInt Q, CeedQFunctionField *qf_input_fields, CeedOperatorField *op_input_fields,
225                                                  CeedInt num_input_fields, CeedInt blk_size, bool skip_active,
226                                                  CeedScalar *e_data_full[2 * CEED_FIELD_MAX], CeedOperator_Blocked *impl) {
227   CeedInt             elem_size, size, num_comp;
228   CeedElemRestriction elem_restr;
229   CeedEvalMode        eval_mode;
230   CeedBasis           basis;
231 
232   for (CeedInt i = 0; i < num_input_fields; i++) {
233     // Skip active input
234     if (skip_active) {
235       CeedVector vec;
236       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
237       if (vec == CEED_VECTOR_ACTIVE) continue;
238     }
239 
240     // Get elem_size, eval_mode, size
241     CeedCallBackend(CeedOperatorFieldGetElemRestriction(op_input_fields[i], &elem_restr));
242     CeedCallBackend(CeedElemRestrictionGetElementSize(elem_restr, &elem_size));
243     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
244     CeedCallBackend(CeedQFunctionFieldGetSize(qf_input_fields[i], &size));
245     // Basis action
246     switch (eval_mode) {
247       case CEED_EVAL_NONE:
248         CeedCallBackend(CeedVectorSetArray(impl->q_vecs_in[i], CEED_MEM_HOST, CEED_USE_POINTER, &e_data_full[i][e * Q * size]));
249         break;
250       case CEED_EVAL_INTERP:
251       case CEED_EVAL_GRAD:
252       case CEED_EVAL_DIV:
253       case CEED_EVAL_CURL:
254         CeedCallBackend(CeedOperatorFieldGetBasis(op_input_fields[i], &basis));
255         CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
256         CeedCallBackend(CeedVectorSetArray(impl->e_vecs_in[i], CEED_MEM_HOST, CEED_USE_POINTER, &e_data_full[i][e * elem_size * num_comp]));
257         CeedCallBackend(CeedBasisApply(basis, blk_size, CEED_NOTRANSPOSE, eval_mode, impl->e_vecs_in[i], impl->q_vecs_in[i]));
258         break;
259       case CEED_EVAL_WEIGHT:
260         break;  // No action
261     }
262   }
263   return CEED_ERROR_SUCCESS;
264 }
265 
266 //------------------------------------------------------------------------------
267 // Output Basis Action
268 //------------------------------------------------------------------------------
269 static inline int CeedOperatorOutputBasis_Blocked(CeedInt e, CeedInt Q, CeedQFunctionField *qf_output_fields, CeedOperatorField *op_output_fields,
270                                                   CeedInt blk_size, CeedInt num_input_fields, CeedInt num_output_fields, CeedOperator op,
271                                                   CeedScalar *e_data_full[2 * CEED_FIELD_MAX], CeedOperator_Blocked *impl) {
272   CeedInt             elem_size, num_comp;
273   CeedElemRestriction elem_restr;
274   CeedEvalMode        eval_mode;
275   CeedBasis           basis;
276 
277   for (CeedInt i = 0; i < num_output_fields; i++) {
278     // Get elem_size, eval_mode, size
279     CeedCallBackend(CeedOperatorFieldGetElemRestriction(op_output_fields[i], &elem_restr));
280     CeedCallBackend(CeedElemRestrictionGetElementSize(elem_restr, &elem_size));
281     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
282     // Basis action
283     switch (eval_mode) {
284       case CEED_EVAL_NONE:
285         break;  // No action
286       case CEED_EVAL_INTERP:
287       case CEED_EVAL_GRAD:
288       case CEED_EVAL_DIV:
289       case CEED_EVAL_CURL:
290         CeedCallBackend(CeedOperatorFieldGetBasis(op_output_fields[i], &basis));
291         CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
292         CeedCallBackend(
293             CeedVectorSetArray(impl->e_vecs_out[i], CEED_MEM_HOST, CEED_USE_POINTER, &e_data_full[i + num_input_fields][e * elem_size * num_comp]));
294         CeedCallBackend(CeedBasisApply(basis, blk_size, CEED_TRANSPOSE, eval_mode, impl->q_vecs_out[i], impl->e_vecs_out[i]));
295         break;
296       // LCOV_EXCL_START
297       case CEED_EVAL_WEIGHT: {
298         Ceed ceed;
299         CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
300         return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT cannot be an output evaluation mode");
301         // LCOV_EXCL_STOP
302       }
303     }
304   }
305   return CEED_ERROR_SUCCESS;
306 }
307 
308 //------------------------------------------------------------------------------
309 // Restore Input Vectors
310 //------------------------------------------------------------------------------
311 static inline int CeedOperatorRestoreInputs_Blocked(CeedInt num_input_fields, CeedQFunctionField *qf_input_fields, CeedOperatorField *op_input_fields,
312                                                     bool skip_active, CeedScalar *e_data_full[2 * CEED_FIELD_MAX], CeedOperator_Blocked *impl) {
313   CeedEvalMode eval_mode;
314 
315   for (CeedInt i = 0; i < num_input_fields; i++) {
316     // Skip active inputs
317     if (skip_active) {
318       CeedVector vec;
319       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
320       if (vec == CEED_VECTOR_ACTIVE) continue;
321     }
322     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
323     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
324     } else {
325       CeedCallBackend(CeedVectorRestoreArrayRead(impl->e_vecs_full[i], (const CeedScalar **)&e_data_full[i]));
326     }
327   }
328   return CEED_ERROR_SUCCESS;
329 }
330 
331 //------------------------------------------------------------------------------
332 // Operator Apply
333 //------------------------------------------------------------------------------
334 static int CeedOperatorApplyAdd_Blocked(CeedOperator op, CeedVector in_vec, CeedVector out_vec, CeedRequest *request) {
335   CeedOperator_Blocked *impl;
336   CeedCallBackend(CeedOperatorGetData(op, &impl));
337   const CeedInt blk_size = 8;
338   CeedInt       Q, num_input_fields, num_output_fields, num_elem, size;
339   CeedCallBackend(CeedOperatorGetNumElements(op, &num_elem));
340   CeedCallBackend(CeedOperatorGetNumQuadraturePoints(op, &Q));
341   CeedInt       num_blks = (num_elem / blk_size) + !!(num_elem % blk_size);
342   CeedQFunction qf;
343   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
344   CeedOperatorField *op_input_fields, *op_output_fields;
345   CeedCallBackend(CeedOperatorGetFields(op, &num_input_fields, &op_input_fields, &num_output_fields, &op_output_fields));
346   CeedQFunctionField *qf_input_fields, *qf_output_fields;
347   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_input_fields, NULL, &qf_output_fields));
348   CeedEvalMode eval_mode;
349   CeedVector   vec;
350   CeedScalar  *e_data_full[2 * CEED_FIELD_MAX] = {0};
351 
352   // Setup
353   CeedCallBackend(CeedOperatorSetup_Blocked(op));
354 
355   // Restriction only operator
356   if (impl->is_identity_restr_op) {
357     CeedCallBackend(CeedElemRestrictionApply(impl->blk_restr[0], CEED_NOTRANSPOSE, in_vec, impl->e_vecs_full[0], request));
358     CeedCallBackend(CeedElemRestrictionApply(impl->blk_restr[1], CEED_TRANSPOSE, impl->e_vecs_full[0], out_vec, request));
359     return CEED_ERROR_SUCCESS;
360   }
361 
362   // Input Evecs and Restriction
363   CeedCallBackend(CeedOperatorSetupInputs_Blocked(num_input_fields, qf_input_fields, op_input_fields, in_vec, false, e_data_full, impl, request));
364 
365   // Output Evecs
366   for (CeedInt i = 0; i < num_output_fields; i++) {
367     CeedCallBackend(CeedVectorGetArrayWrite(impl->e_vecs_full[i + impl->num_inputs], CEED_MEM_HOST, &e_data_full[i + num_input_fields]));
368   }
369 
370   // Loop through elements
371   for (CeedInt e = 0; e < num_blks * blk_size; e += blk_size) {
372     // Output pointers
373     for (CeedInt i = 0; i < num_output_fields; i++) {
374       CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
375       if (eval_mode == CEED_EVAL_NONE) {
376         CeedCallBackend(CeedQFunctionFieldGetSize(qf_output_fields[i], &size));
377         CeedCallBackend(CeedVectorSetArray(impl->q_vecs_out[i], CEED_MEM_HOST, CEED_USE_POINTER, &e_data_full[i + num_input_fields][e * Q * size]));
378       }
379     }
380 
381     // Input basis apply
382     CeedCallBackend(CeedOperatorInputBasis_Blocked(e, Q, qf_input_fields, op_input_fields, num_input_fields, blk_size, false, e_data_full, impl));
383 
384     // Q function
385     if (!impl->is_identity_qf) {
386       CeedCallBackend(CeedQFunctionApply(qf, Q * blk_size, impl->q_vecs_in, impl->q_vecs_out));
387     }
388 
389     // Output basis apply
390     CeedCallBackend(CeedOperatorOutputBasis_Blocked(e, Q, qf_output_fields, op_output_fields, blk_size, num_input_fields, num_output_fields, op,
391                                                     e_data_full, impl));
392   }
393 
394   // Output restriction
395   for (CeedInt i = 0; i < num_output_fields; i++) {
396     // Restore evec
397     CeedCallBackend(CeedVectorRestoreArray(impl->e_vecs_full[i + impl->num_inputs], &e_data_full[i + num_input_fields]));
398     // Get output vector
399     CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
400     // Active
401     if (vec == CEED_VECTOR_ACTIVE) vec = out_vec;
402     // Restrict
403     CeedCallBackend(
404         CeedElemRestrictionApply(impl->blk_restr[i + impl->num_inputs], CEED_TRANSPOSE, impl->e_vecs_full[i + impl->num_inputs], vec, request));
405   }
406 
407   // Restore input arrays
408   CeedCallBackend(CeedOperatorRestoreInputs_Blocked(num_input_fields, qf_input_fields, op_input_fields, false, e_data_full, impl));
409 
410   return CEED_ERROR_SUCCESS;
411 }
412 
413 //------------------------------------------------------------------------------
414 // Core code for assembling linear QFunction
415 //------------------------------------------------------------------------------
416 static inline int CeedOperatorLinearAssembleQFunctionCore_Blocked(CeedOperator op, bool build_objects, CeedVector *assembled,
417                                                                   CeedElemRestriction *rstr, CeedRequest *request) {
418   CeedOperator_Blocked *impl;
419   CeedCallBackend(CeedOperatorGetData(op, &impl));
420   const CeedInt blk_size = 8;
421   CeedInt       Q, num_input_fields, num_output_fields, num_elem, size;
422   CeedCallBackend(CeedOperatorGetNumElements(op, &num_elem));
423   CeedCallBackend(CeedOperatorGetNumQuadraturePoints(op, &Q));
424   CeedInt       num_blks = (num_elem / blk_size) + !!(num_elem % blk_size);
425   CeedQFunction qf;
426   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
427   CeedOperatorField *op_input_fields, *op_output_fields;
428   CeedCallBackend(CeedOperatorGetFields(op, &num_input_fields, &op_input_fields, &num_output_fields, &op_output_fields));
429   CeedQFunctionField *qf_input_fields, *qf_output_fields;
430   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_input_fields, NULL, &qf_output_fields));
431   CeedVector  vec, l_vec = impl->qf_l_vec;
432   CeedInt     num_active_in = impl->num_active_in, num_active_out = impl->num_active_out;
433   CeedSize    q_size;
434   CeedVector *active_in = impl->qf_active_in;
435   CeedScalar *a, *tmp;
436   Ceed        ceed;
437   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
438   CeedScalar *e_data_full[2 * CEED_FIELD_MAX] = {0};
439 
440   // Setup
441   CeedCallBackend(CeedOperatorSetup_Blocked(op));
442 
443   // Check for identity
444   CeedCheck(!impl->is_identity_qf, ceed, CEED_ERROR_BACKEND, "Assembling identity QFunctions not supported");
445 
446   // Input Evecs and Restriction
447   CeedCallBackend(CeedOperatorSetupInputs_Blocked(num_input_fields, qf_input_fields, op_input_fields, NULL, true, e_data_full, impl, request));
448 
449   // Count number of active input fields
450   if (!num_active_in) {
451     for (CeedInt i = 0; i < num_input_fields; i++) {
452       // Get input vector
453       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
454       // Check if active input
455       if (vec == CEED_VECTOR_ACTIVE) {
456         CeedCallBackend(CeedQFunctionFieldGetSize(qf_input_fields[i], &size));
457         CeedCallBackend(CeedVectorSetValue(impl->q_vecs_in[i], 0.0));
458         CeedCallBackend(CeedVectorGetArray(impl->q_vecs_in[i], CEED_MEM_HOST, &tmp));
459         CeedCallBackend(CeedRealloc(num_active_in + size, &active_in));
460         for (CeedInt field = 0; field < size; field++) {
461           q_size = (CeedSize)Q * blk_size;
462           CeedCallBackend(CeedVectorCreate(ceed, q_size, &active_in[num_active_in + field]));
463           CeedCallBackend(CeedVectorSetArray(active_in[num_active_in + field], CEED_MEM_HOST, CEED_USE_POINTER, &tmp[field * Q * blk_size]));
464         }
465         num_active_in += size;
466         CeedCallBackend(CeedVectorRestoreArray(impl->q_vecs_in[i], &tmp));
467       }
468     }
469     impl->num_active_in = num_active_in;
470     impl->qf_active_in  = active_in;
471   }
472 
473   // Count number of active output fields
474   if (!num_active_out) {
475     for (CeedInt i = 0; i < num_output_fields; i++) {
476       // Get output vector
477       CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
478       // Check if active output
479       if (vec == CEED_VECTOR_ACTIVE) {
480         CeedCallBackend(CeedQFunctionFieldGetSize(qf_output_fields[i], &size));
481         num_active_out += size;
482       }
483     }
484     impl->num_active_out = num_active_out;
485   }
486 
487   // Check sizes
488   CeedCheck(num_active_in > 0 && num_active_out > 0, ceed, CEED_ERROR_BACKEND, "Cannot assemble QFunction without active inputs and outputs");
489 
490   // Setup Lvec
491   if (!l_vec) {
492     CeedSize l_size = (CeedSize)num_blks * blk_size * Q * num_active_in * num_active_out;
493     CeedCallBackend(CeedVectorCreate(ceed, l_size, &l_vec));
494     impl->qf_l_vec = l_vec;
495   }
496   CeedCallBackend(CeedVectorGetArrayWrite(l_vec, CEED_MEM_HOST, &a));
497 
498   // Build objects if needed
499   CeedInt strides[3] = {1, Q, num_active_in * num_active_out * Q};
500   if (build_objects) {
501     // Create output restriction
502     CeedCallBackend(CeedElemRestrictionCreateStrided(ceed, num_elem, Q, num_active_in * num_active_out, num_active_in * num_active_out * num_elem * Q,
503                                                      strides, rstr));
504     // Create assembled vector
505     CeedSize l_size = (CeedSize)num_elem * Q * num_active_in * num_active_out;
506     CeedCallBackend(CeedVectorCreate(ceed, l_size, assembled));
507   }
508 
509   // Loop through elements
510   for (CeedInt e = 0; e < num_blks * blk_size; e += blk_size) {
511     // Input basis apply
512     CeedCallBackend(CeedOperatorInputBasis_Blocked(e, Q, qf_input_fields, op_input_fields, num_input_fields, blk_size, true, e_data_full, impl));
513 
514     // Assemble QFunction
515     for (CeedInt in = 0; in < num_active_in; in++) {
516       // Set Inputs
517       CeedCallBackend(CeedVectorSetValue(active_in[in], 1.0));
518       if (num_active_in > 1) {
519         CeedCallBackend(CeedVectorSetValue(active_in[(in + num_active_in - 1) % num_active_in], 0.0));
520       }
521       // Set Outputs
522       for (CeedInt out = 0; out < num_output_fields; out++) {
523         // Get output vector
524         CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[out], &vec));
525         // Check if active output
526         if (vec == CEED_VECTOR_ACTIVE) {
527           CeedCallBackend(CeedVectorSetArray(impl->q_vecs_out[out], CEED_MEM_HOST, CEED_USE_POINTER, a));
528           CeedCallBackend(CeedQFunctionFieldGetSize(qf_output_fields[out], &size));
529           a += size * Q * blk_size;  // Advance the pointer by the size of the output
530         }
531       }
532       // Apply QFunction
533       CeedCallBackend(CeedQFunctionApply(qf, Q * blk_size, impl->q_vecs_in, impl->q_vecs_out));
534     }
535   }
536 
537   // Un-set output Qvecs to prevent accidental overwrite of Assembled
538   for (CeedInt out = 0; out < num_output_fields; out++) {
539     // Get output vector
540     CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[out], &vec));
541     // Check if active output
542     if (vec == CEED_VECTOR_ACTIVE) {
543       CeedCallBackend(CeedVectorSetArray(impl->q_vecs_out[out], CEED_MEM_HOST, CEED_COPY_VALUES, NULL));
544     }
545   }
546 
547   // Restore input arrays
548   CeedCallBackend(CeedOperatorRestoreInputs_Blocked(num_input_fields, qf_input_fields, op_input_fields, true, e_data_full, impl));
549 
550   // Output blocked restriction
551   CeedCallBackend(CeedVectorRestoreArray(l_vec, &a));
552   CeedCallBackend(CeedVectorSetValue(*assembled, 0.0));
553   CeedElemRestriction blk_rstr = impl->qf_blk_rstr;
554   if (!impl->qf_blk_rstr) {
555     CeedCallBackend(CeedElemRestrictionCreateBlockedStrided(ceed, num_elem, Q, blk_size, num_active_in * num_active_out,
556                                                             num_active_in * num_active_out * num_elem * Q, strides, &blk_rstr));
557     impl->qf_blk_rstr = blk_rstr;
558   }
559   CeedCallBackend(CeedElemRestrictionApply(blk_rstr, CEED_TRANSPOSE, l_vec, *assembled, request));
560 
561   return CEED_ERROR_SUCCESS;
562 }
563 
564 //------------------------------------------------------------------------------
565 // Assemble Linear QFunction
566 //------------------------------------------------------------------------------
567 static int CeedOperatorLinearAssembleQFunction_Blocked(CeedOperator op, CeedVector *assembled, CeedElemRestriction *rstr, CeedRequest *request) {
568   return CeedOperatorLinearAssembleQFunctionCore_Blocked(op, true, assembled, rstr, request);
569 }
570 
571 //------------------------------------------------------------------------------
572 // Update Assembled Linear QFunction
573 //------------------------------------------------------------------------------
574 static int CeedOperatorLinearAssembleQFunctionUpdate_Blocked(CeedOperator op, CeedVector assembled, CeedElemRestriction rstr, CeedRequest *request) {
575   return CeedOperatorLinearAssembleQFunctionCore_Blocked(op, false, &assembled, &rstr, request);
576 }
577 
578 //------------------------------------------------------------------------------
579 // Operator Destroy
580 //------------------------------------------------------------------------------
581 static int CeedOperatorDestroy_Blocked(CeedOperator op) {
582   CeedOperator_Blocked *impl;
583   CeedCallBackend(CeedOperatorGetData(op, &impl));
584 
585   for (CeedInt i = 0; i < impl->num_inputs + impl->num_outputs; i++) {
586     CeedCallBackend(CeedElemRestrictionDestroy(&impl->blk_restr[i]));
587     CeedCallBackend(CeedVectorDestroy(&impl->e_vecs_full[i]));
588   }
589   CeedCallBackend(CeedFree(&impl->blk_restr));
590   CeedCallBackend(CeedFree(&impl->e_vecs_full));
591   CeedCallBackend(CeedFree(&impl->input_states));
592 
593   for (CeedInt i = 0; i < impl->num_inputs; i++) {
594     CeedCallBackend(CeedVectorDestroy(&impl->e_vecs_in[i]));
595     CeedCallBackend(CeedVectorDestroy(&impl->q_vecs_in[i]));
596   }
597   CeedCallBackend(CeedFree(&impl->e_vecs_in));
598   CeedCallBackend(CeedFree(&impl->q_vecs_in));
599 
600   for (CeedInt i = 0; i < impl->num_outputs; i++) {
601     CeedCallBackend(CeedVectorDestroy(&impl->e_vecs_out[i]));
602     CeedCallBackend(CeedVectorDestroy(&impl->q_vecs_out[i]));
603   }
604   CeedCallBackend(CeedFree(&impl->e_vecs_out));
605   CeedCallBackend(CeedFree(&impl->q_vecs_out));
606 
607   // QFunction assembly data
608   for (CeedInt i = 0; i < impl->num_active_in; i++) {
609     CeedCallBackend(CeedVectorDestroy(&impl->qf_active_in[i]));
610   }
611   CeedCallBackend(CeedFree(&impl->qf_active_in));
612   CeedCallBackend(CeedVectorDestroy(&impl->qf_l_vec));
613   CeedCallBackend(CeedElemRestrictionDestroy(&impl->qf_blk_rstr));
614 
615   CeedCallBackend(CeedFree(&impl));
616   return CEED_ERROR_SUCCESS;
617 }
618 
619 //------------------------------------------------------------------------------
620 // Operator Create
621 //------------------------------------------------------------------------------
622 int CeedOperatorCreate_Blocked(CeedOperator op) {
623   Ceed ceed;
624   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
625   CeedOperator_Blocked *impl;
626 
627   CeedCallBackend(CeedCalloc(1, &impl));
628   CeedCallBackend(CeedOperatorSetData(op, impl));
629 
630   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleQFunction", CeedOperatorLinearAssembleQFunction_Blocked));
631   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleQFunctionUpdate", CeedOperatorLinearAssembleQFunctionUpdate_Blocked));
632   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAdd", CeedOperatorApplyAdd_Blocked));
633   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "Destroy", CeedOperatorDestroy_Blocked));
634   return CEED_ERROR_SUCCESS;
635 }
636 
637 //------------------------------------------------------------------------------
638