xref: /libCEED/backends/cuda-ref/ceed-cuda-ref-operator.c (revision fd8f24fa4ef1026f43db0b90b0c68f0c3af48272)
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 <ceed/jit-tools.h>
11 #include <assert.h>
12 #include <cuda.h>
13 #include <cuda_runtime.h>
14 #include <stdbool.h>
15 #include <string.h>
16 
17 #include "../cuda/ceed-cuda-common.h"
18 #include "../cuda/ceed-cuda-compile.h"
19 #include "ceed-cuda-ref.h"
20 
21 //------------------------------------------------------------------------------
22 // Destroy operator
23 //------------------------------------------------------------------------------
24 static int CeedOperatorDestroy_Cuda(CeedOperator op) {
25   CeedOperator_Cuda *impl;
26   CeedCallBackend(CeedOperatorGetData(op, &impl));
27 
28   // Apply data
29   for (CeedInt i = 0; i < impl->numein + impl->numeout; i++) {
30     CeedCallBackend(CeedVectorDestroy(&impl->evecs[i]));
31   }
32   CeedCallBackend(CeedFree(&impl->evecs));
33 
34   for (CeedInt i = 0; i < impl->numein; i++) {
35     CeedCallBackend(CeedVectorDestroy(&impl->qvecsin[i]));
36   }
37   CeedCallBackend(CeedFree(&impl->qvecsin));
38 
39   for (CeedInt i = 0; i < impl->numeout; i++) {
40     CeedCallBackend(CeedVectorDestroy(&impl->qvecsout[i]));
41   }
42   CeedCallBackend(CeedFree(&impl->qvecsout));
43 
44   // QFunction assembly data
45   for (CeedInt i = 0; i < impl->qfnumactivein; i++) {
46     CeedCallBackend(CeedVectorDestroy(&impl->qfactivein[i]));
47   }
48   CeedCallBackend(CeedFree(&impl->qfactivein));
49 
50   // Diag data
51   if (impl->diag) {
52     Ceed ceed;
53     CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
54     CeedCallCuda(ceed, cuModuleUnload(impl->diag->module));
55     CeedCallBackend(CeedFree(&impl->diag->h_emodein));
56     CeedCallBackend(CeedFree(&impl->diag->h_emodeout));
57     CeedCallCuda(ceed, cudaFree(impl->diag->d_emodein));
58     CeedCallCuda(ceed, cudaFree(impl->diag->d_emodeout));
59     CeedCallCuda(ceed, cudaFree(impl->diag->d_identity));
60     CeedCallCuda(ceed, cudaFree(impl->diag->d_interpin));
61     CeedCallCuda(ceed, cudaFree(impl->diag->d_interpout));
62     CeedCallCuda(ceed, cudaFree(impl->diag->d_gradin));
63     CeedCallCuda(ceed, cudaFree(impl->diag->d_gradout));
64     CeedCallBackend(CeedElemRestrictionDestroy(&impl->diag->pbdiagrstr));
65     CeedCallBackend(CeedVectorDestroy(&impl->diag->elemdiag));
66     CeedCallBackend(CeedVectorDestroy(&impl->diag->pbelemdiag));
67   }
68   CeedCallBackend(CeedFree(&impl->diag));
69 
70   if (impl->asmb) {
71     Ceed ceed;
72     CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
73     CeedCallCuda(ceed, cuModuleUnload(impl->asmb->module));
74     CeedCallCuda(ceed, cudaFree(impl->asmb->d_B_in));
75     CeedCallCuda(ceed, cudaFree(impl->asmb->d_B_out));
76   }
77   CeedCallBackend(CeedFree(&impl->asmb));
78 
79   CeedCallBackend(CeedFree(&impl));
80   return CEED_ERROR_SUCCESS;
81 }
82 
83 //------------------------------------------------------------------------------
84 // Setup infields or outfields
85 //------------------------------------------------------------------------------
86 static int CeedOperatorSetupFields_Cuda(CeedQFunction qf, CeedOperator op, bool isinput, CeedVector *evecs, CeedVector *qvecs, CeedInt starte,
87                                         CeedInt numfields, CeedInt Q, CeedInt numelements) {
88   CeedInt  dim, size;
89   CeedSize q_size;
90   Ceed     ceed;
91   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
92   CeedBasis           basis;
93   CeedElemRestriction Erestrict;
94   CeedOperatorField  *opfields;
95   CeedQFunctionField *qffields;
96   CeedVector          fieldvec;
97   bool                strided;
98   bool                skiprestrict;
99 
100   if (isinput) {
101     CeedCallBackend(CeedOperatorGetFields(op, NULL, &opfields, NULL, NULL));
102     CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qffields, NULL, NULL));
103   } else {
104     CeedCallBackend(CeedOperatorGetFields(op, NULL, NULL, NULL, &opfields));
105     CeedCallBackend(CeedQFunctionGetFields(qf, NULL, NULL, NULL, &qffields));
106   }
107 
108   // Loop over fields
109   for (CeedInt i = 0; i < numfields; i++) {
110     CeedEvalMode emode;
111     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qffields[i], &emode));
112 
113     strided      = false;
114     skiprestrict = false;
115     if (emode != CEED_EVAL_WEIGHT) {
116       CeedCallBackend(CeedOperatorFieldGetElemRestriction(opfields[i], &Erestrict));
117 
118       // Check whether this field can skip the element restriction:
119       // must be passive input, with emode NONE, and have a strided restriction with CEED_STRIDES_BACKEND.
120 
121       // First, check whether the field is input or output:
122       if (isinput) {
123         // Check for passive input:
124         CeedCallBackend(CeedOperatorFieldGetVector(opfields[i], &fieldvec));
125         if (fieldvec != CEED_VECTOR_ACTIVE) {
126           // Check emode
127           if (emode == CEED_EVAL_NONE) {
128             // Check for strided restriction
129             CeedCallBackend(CeedElemRestrictionIsStrided(Erestrict, &strided));
130             if (strided) {
131               // Check if vector is already in preferred backend ordering
132               CeedCallBackend(CeedElemRestrictionHasBackendStrides(Erestrict, &skiprestrict));
133             }
134           }
135         }
136       }
137       if (skiprestrict) {
138         // We do not need an E-Vector, but will use the input field vector's data directly in the operator application.
139         evecs[i + starte] = NULL;
140       } else {
141         CeedCallBackend(CeedElemRestrictionCreateVector(Erestrict, NULL, &evecs[i + starte]));
142       }
143     }
144 
145     switch (emode) {
146       case CEED_EVAL_NONE:
147         CeedCallBackend(CeedQFunctionFieldGetSize(qffields[i], &size));
148         q_size = (CeedSize)numelements * Q * size;
149         CeedCallBackend(CeedVectorCreate(ceed, q_size, &qvecs[i]));
150         break;
151       case CEED_EVAL_INTERP:
152         CeedCallBackend(CeedQFunctionFieldGetSize(qffields[i], &size));
153         q_size = (CeedSize)numelements * Q * size;
154         CeedCallBackend(CeedVectorCreate(ceed, q_size, &qvecs[i]));
155         break;
156       case CEED_EVAL_GRAD:
157         CeedCallBackend(CeedOperatorFieldGetBasis(opfields[i], &basis));
158         CeedCallBackend(CeedQFunctionFieldGetSize(qffields[i], &size));
159         CeedCallBackend(CeedBasisGetDimension(basis, &dim));
160         q_size = (CeedSize)numelements * Q * size;
161         CeedCallBackend(CeedVectorCreate(ceed, q_size, &qvecs[i]));
162         break;
163       case CEED_EVAL_WEIGHT:  // Only on input fields
164         CeedCallBackend(CeedOperatorFieldGetBasis(opfields[i], &basis));
165         q_size = (CeedSize)numelements * Q;
166         CeedCallBackend(CeedVectorCreate(ceed, q_size, &qvecs[i]));
167         CeedCallBackend(CeedBasisApply(basis, numelements, CEED_NOTRANSPOSE, CEED_EVAL_WEIGHT, CEED_VECTOR_NONE, qvecs[i]));
168         break;
169       case CEED_EVAL_DIV:
170         break;  // TODO: Not implemented
171       case CEED_EVAL_CURL:
172         break;  // TODO: Not implemented
173     }
174   }
175   return CEED_ERROR_SUCCESS;
176 }
177 
178 //------------------------------------------------------------------------------
179 // CeedOperator needs to connect all the named fields (be they active or passive) to the named inputs and outputs of its CeedQFunction.
180 //------------------------------------------------------------------------------
181 static int CeedOperatorSetup_Cuda(CeedOperator op) {
182   bool setupdone;
183   CeedCallBackend(CeedOperatorIsSetupDone(op, &setupdone));
184   if (setupdone) return CEED_ERROR_SUCCESS;
185   Ceed ceed;
186   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
187   CeedOperator_Cuda *impl;
188   CeedCallBackend(CeedOperatorGetData(op, &impl));
189   CeedQFunction qf;
190   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
191   CeedInt Q, numelements, numinputfields, numoutputfields;
192   CeedCallBackend(CeedOperatorGetNumQuadraturePoints(op, &Q));
193   CeedCallBackend(CeedOperatorGetNumElements(op, &numelements));
194   CeedOperatorField *opinputfields, *opoutputfields;
195   CeedCallBackend(CeedOperatorGetFields(op, &numinputfields, &opinputfields, &numoutputfields, &opoutputfields));
196   CeedQFunctionField *qfinputfields, *qfoutputfields;
197   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qfinputfields, NULL, &qfoutputfields));
198 
199   // Allocate
200   CeedCallBackend(CeedCalloc(numinputfields + numoutputfields, &impl->evecs));
201 
202   CeedCallBackend(CeedCalloc(CEED_FIELD_MAX, &impl->qvecsin));
203   CeedCallBackend(CeedCalloc(CEED_FIELD_MAX, &impl->qvecsout));
204 
205   impl->numein  = numinputfields;
206   impl->numeout = numoutputfields;
207 
208   // Set up infield and outfield evecs and qvecs
209   // Infields
210   CeedCallBackend(CeedOperatorSetupFields_Cuda(qf, op, true, impl->evecs, impl->qvecsin, 0, numinputfields, Q, numelements));
211 
212   // Outfields
213   CeedCallBackend(CeedOperatorSetupFields_Cuda(qf, op, false, impl->evecs, impl->qvecsout, numinputfields, numoutputfields, Q, numelements));
214 
215   CeedCallBackend(CeedOperatorSetSetupDone(op));
216   return CEED_ERROR_SUCCESS;
217 }
218 
219 //------------------------------------------------------------------------------
220 // Setup Operator Inputs
221 //------------------------------------------------------------------------------
222 static inline int CeedOperatorSetupInputs_Cuda(CeedInt numinputfields, CeedQFunctionField *qfinputfields, CeedOperatorField *opinputfields,
223                                                CeedVector invec, const bool skipactive, CeedScalar *edata[2 * CEED_FIELD_MAX],
224                                                CeedOperator_Cuda *impl, CeedRequest *request) {
225   CeedEvalMode        emode;
226   CeedVector          vec;
227   CeedElemRestriction Erestrict;
228 
229   for (CeedInt i = 0; i < numinputfields; i++) {
230     // Get input vector
231     CeedCallBackend(CeedOperatorFieldGetVector(opinputfields[i], &vec));
232     if (vec == CEED_VECTOR_ACTIVE) {
233       if (skipactive) continue;
234       else vec = invec;
235     }
236 
237     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qfinputfields[i], &emode));
238     if (emode == CEED_EVAL_WEIGHT) {  // Skip
239     } else {
240       // Get input vector
241       CeedCallBackend(CeedOperatorFieldGetVector(opinputfields[i], &vec));
242       // Get input element restriction
243       CeedCallBackend(CeedOperatorFieldGetElemRestriction(opinputfields[i], &Erestrict));
244       if (vec == CEED_VECTOR_ACTIVE) vec = invec;
245       // Restrict, if necessary
246       if (!impl->evecs[i]) {
247         // No restriction for this field; read data directly from vec.
248         CeedCallBackend(CeedVectorGetArrayRead(vec, CEED_MEM_DEVICE, (const CeedScalar **)&edata[i]));
249       } else {
250         CeedCallBackend(CeedElemRestrictionApply(Erestrict, CEED_NOTRANSPOSE, vec, impl->evecs[i], request));
251         // Get evec
252         CeedCallBackend(CeedVectorGetArrayRead(impl->evecs[i], CEED_MEM_DEVICE, (const CeedScalar **)&edata[i]));
253       }
254     }
255   }
256   return CEED_ERROR_SUCCESS;
257 }
258 
259 //------------------------------------------------------------------------------
260 // Input Basis Action
261 //------------------------------------------------------------------------------
262 static inline int CeedOperatorInputBasis_Cuda(CeedInt numelements, CeedQFunctionField *qfinputfields, CeedOperatorField *opinputfields,
263                                               CeedInt numinputfields, const bool skipactive, CeedScalar *edata[2 * CEED_FIELD_MAX],
264                                               CeedOperator_Cuda *impl) {
265   CeedInt             elemsize, size;
266   CeedElemRestriction Erestrict;
267   CeedEvalMode        emode;
268   CeedBasis           basis;
269 
270   for (CeedInt i = 0; i < numinputfields; i++) {
271     // Skip active input
272     if (skipactive) {
273       CeedVector vec;
274       CeedCallBackend(CeedOperatorFieldGetVector(opinputfields[i], &vec));
275       if (vec == CEED_VECTOR_ACTIVE) continue;
276     }
277     // Get elemsize, emode, size
278     CeedCallBackend(CeedOperatorFieldGetElemRestriction(opinputfields[i], &Erestrict));
279     CeedCallBackend(CeedElemRestrictionGetElementSize(Erestrict, &elemsize));
280     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qfinputfields[i], &emode));
281     CeedCallBackend(CeedQFunctionFieldGetSize(qfinputfields[i], &size));
282     // Basis action
283     switch (emode) {
284       case CEED_EVAL_NONE:
285         CeedCallBackend(CeedVectorSetArray(impl->qvecsin[i], CEED_MEM_DEVICE, CEED_USE_POINTER, edata[i]));
286         break;
287       case CEED_EVAL_INTERP:
288         CeedCallBackend(CeedOperatorFieldGetBasis(opinputfields[i], &basis));
289         CeedCallBackend(CeedBasisApply(basis, numelements, CEED_NOTRANSPOSE, CEED_EVAL_INTERP, impl->evecs[i], impl->qvecsin[i]));
290         break;
291       case CEED_EVAL_GRAD:
292         CeedCallBackend(CeedOperatorFieldGetBasis(opinputfields[i], &basis));
293         CeedCallBackend(CeedBasisApply(basis, numelements, CEED_NOTRANSPOSE, CEED_EVAL_GRAD, impl->evecs[i], impl->qvecsin[i]));
294         break;
295       case CEED_EVAL_WEIGHT:
296         break;  // No action
297       case CEED_EVAL_DIV:
298         break;  // TODO: Not implemented
299       case CEED_EVAL_CURL:
300         break;  // TODO: Not implemented
301     }
302   }
303   return CEED_ERROR_SUCCESS;
304 }
305 
306 //------------------------------------------------------------------------------
307 // Restore Input Vectors
308 //------------------------------------------------------------------------------
309 static inline int CeedOperatorRestoreInputs_Cuda(CeedInt numinputfields, CeedQFunctionField *qfinputfields, CeedOperatorField *opinputfields,
310                                                  const bool skipactive, CeedScalar *edata[2 * CEED_FIELD_MAX], CeedOperator_Cuda *impl) {
311   CeedEvalMode emode;
312   CeedVector   vec;
313 
314   for (CeedInt i = 0; i < numinputfields; i++) {
315     // Skip active input
316     if (skipactive) {
317       CeedCallBackend(CeedOperatorFieldGetVector(opinputfields[i], &vec));
318       if (vec == CEED_VECTOR_ACTIVE) continue;
319     }
320     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qfinputfields[i], &emode));
321     if (emode == CEED_EVAL_WEIGHT) {  // Skip
322     } else {
323       if (!impl->evecs[i]) {  // This was a skiprestrict case
324         CeedCallBackend(CeedOperatorFieldGetVector(opinputfields[i], &vec));
325         CeedCallBackend(CeedVectorRestoreArrayRead(vec, (const CeedScalar **)&edata[i]));
326       } else {
327         CeedCallBackend(CeedVectorRestoreArrayRead(impl->evecs[i], (const CeedScalar **)&edata[i]));
328       }
329     }
330   }
331   return CEED_ERROR_SUCCESS;
332 }
333 
334 //------------------------------------------------------------------------------
335 // Apply and add to output
336 //------------------------------------------------------------------------------
337 static int CeedOperatorApplyAdd_Cuda(CeedOperator op, CeedVector invec, CeedVector outvec, CeedRequest *request) {
338   CeedOperator_Cuda *impl;
339   CeedCallBackend(CeedOperatorGetData(op, &impl));
340   CeedQFunction qf;
341   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
342   CeedInt Q, numelements, elemsize, numinputfields, numoutputfields, size;
343   CeedCallBackend(CeedOperatorGetNumQuadraturePoints(op, &Q));
344   CeedCallBackend(CeedOperatorGetNumElements(op, &numelements));
345   CeedOperatorField *opinputfields, *opoutputfields;
346   CeedCallBackend(CeedOperatorGetFields(op, &numinputfields, &opinputfields, &numoutputfields, &opoutputfields));
347   CeedQFunctionField *qfinputfields, *qfoutputfields;
348   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qfinputfields, NULL, &qfoutputfields));
349   CeedEvalMode        emode;
350   CeedVector          vec;
351   CeedBasis           basis;
352   CeedElemRestriction Erestrict;
353   CeedScalar         *edata[2 * CEED_FIELD_MAX] = {NULL};
354 
355   // Setup
356   CeedCallBackend(CeedOperatorSetup_Cuda(op));
357 
358   // Input Evecs and Restriction
359   CeedCallBackend(CeedOperatorSetupInputs_Cuda(numinputfields, qfinputfields, opinputfields, invec, false, edata, impl, request));
360 
361   // Input basis apply if needed
362   CeedCallBackend(CeedOperatorInputBasis_Cuda(numelements, qfinputfields, opinputfields, numinputfields, false, edata, impl));
363 
364   // Output pointers, as necessary
365   for (CeedInt i = 0; i < numoutputfields; i++) {
366     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qfoutputfields[i], &emode));
367     if (emode == CEED_EVAL_NONE) {
368       // Set the output Q-Vector to use the E-Vector data directly.
369       CeedCallBackend(CeedVectorGetArrayWrite(impl->evecs[i + impl->numein], CEED_MEM_DEVICE, &edata[i + numinputfields]));
370       CeedCallBackend(CeedVectorSetArray(impl->qvecsout[i], CEED_MEM_DEVICE, CEED_USE_POINTER, edata[i + numinputfields]));
371     }
372   }
373 
374   // Q function
375   CeedCallBackend(CeedQFunctionApply(qf, numelements * Q, impl->qvecsin, impl->qvecsout));
376 
377   // Output basis apply if needed
378   for (CeedInt i = 0; i < numoutputfields; i++) {
379     // Get elemsize, emode, size
380     CeedCallBackend(CeedOperatorFieldGetElemRestriction(opoutputfields[i], &Erestrict));
381     CeedCallBackend(CeedElemRestrictionGetElementSize(Erestrict, &elemsize));
382     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qfoutputfields[i], &emode));
383     CeedCallBackend(CeedQFunctionFieldGetSize(qfoutputfields[i], &size));
384     // Basis action
385     switch (emode) {
386       case CEED_EVAL_NONE:
387         break;
388       case CEED_EVAL_INTERP:
389         CeedCallBackend(CeedOperatorFieldGetBasis(opoutputfields[i], &basis));
390         CeedCallBackend(CeedBasisApply(basis, numelements, CEED_TRANSPOSE, CEED_EVAL_INTERP, impl->qvecsout[i], impl->evecs[i + impl->numein]));
391         break;
392       case CEED_EVAL_GRAD:
393         CeedCallBackend(CeedOperatorFieldGetBasis(opoutputfields[i], &basis));
394         CeedCallBackend(CeedBasisApply(basis, numelements, CEED_TRANSPOSE, CEED_EVAL_GRAD, impl->qvecsout[i], impl->evecs[i + impl->numein]));
395         break;
396       // LCOV_EXCL_START
397       case CEED_EVAL_WEIGHT: {
398         Ceed ceed;
399         CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
400         return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT cannot be an output evaluation mode");
401         break;  // Should not occur
402       }
403       case CEED_EVAL_DIV:
404         break;  // TODO: Not implemented
405       case CEED_EVAL_CURL:
406         break;  // TODO: Not implemented
407                 // LCOV_EXCL_STOP
408     }
409   }
410 
411   // Output restriction
412   for (CeedInt i = 0; i < numoutputfields; i++) {
413     // Restore evec
414     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qfoutputfields[i], &emode));
415     if (emode == CEED_EVAL_NONE) {
416       CeedCallBackend(CeedVectorRestoreArray(impl->evecs[i + impl->numein], &edata[i + numinputfields]));
417     }
418     // Get output vector
419     CeedCallBackend(CeedOperatorFieldGetVector(opoutputfields[i], &vec));
420     // Restrict
421     CeedCallBackend(CeedOperatorFieldGetElemRestriction(opoutputfields[i], &Erestrict));
422     // Active
423     if (vec == CEED_VECTOR_ACTIVE) vec = outvec;
424 
425     CeedCallBackend(CeedElemRestrictionApply(Erestrict, CEED_TRANSPOSE, impl->evecs[i + impl->numein], vec, request));
426   }
427 
428   // Restore input arrays
429   CeedCallBackend(CeedOperatorRestoreInputs_Cuda(numinputfields, qfinputfields, opinputfields, false, edata, impl));
430   return CEED_ERROR_SUCCESS;
431 }
432 
433 //------------------------------------------------------------------------------
434 // Core code for assembling linear QFunction
435 //------------------------------------------------------------------------------
436 static inline int CeedOperatorLinearAssembleQFunctionCore_Cuda(CeedOperator op, bool build_objects, CeedVector *assembled, CeedElemRestriction *rstr,
437                                                                CeedRequest *request) {
438   CeedOperator_Cuda *impl;
439   CeedCallBackend(CeedOperatorGetData(op, &impl));
440   CeedQFunction qf;
441   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
442   CeedInt  Q, numelements, numinputfields, numoutputfields, size;
443   CeedSize q_size;
444   CeedCallBackend(CeedOperatorGetNumQuadraturePoints(op, &Q));
445   CeedCallBackend(CeedOperatorGetNumElements(op, &numelements));
446   CeedOperatorField *opinputfields, *opoutputfields;
447   CeedCallBackend(CeedOperatorGetFields(op, &numinputfields, &opinputfields, &numoutputfields, &opoutputfields));
448   CeedQFunctionField *qfinputfields, *qfoutputfields;
449   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qfinputfields, NULL, &qfoutputfields));
450   CeedVector  vec;
451   CeedInt     numactivein = impl->qfnumactivein, numactiveout = impl->qfnumactiveout;
452   CeedVector *activein = impl->qfactivein;
453   CeedScalar *a, *tmp;
454   Ceed        ceed, ceedparent;
455   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
456   CeedCallBackend(CeedGetOperatorFallbackParentCeed(ceed, &ceedparent));
457   ceedparent                            = ceedparent ? ceedparent : ceed;
458   CeedScalar *edata[2 * CEED_FIELD_MAX] = {NULL};
459 
460   // Setup
461   CeedCallBackend(CeedOperatorSetup_Cuda(op));
462 
463   // Check for identity
464   bool identityqf;
465   CeedCallBackend(CeedQFunctionIsIdentity(qf, &identityqf));
466   CeedCheck(!identityqf, ceed, CEED_ERROR_BACKEND, "Assembling identity QFunctions not supported");
467 
468   // Input Evecs and Restriction
469   CeedCallBackend(CeedOperatorSetupInputs_Cuda(numinputfields, qfinputfields, opinputfields, NULL, true, edata, impl, request));
470 
471   // Count number of active input fields
472   if (!numactivein) {
473     for (CeedInt i = 0; i < numinputfields; i++) {
474       // Get input vector
475       CeedCallBackend(CeedOperatorFieldGetVector(opinputfields[i], &vec));
476       // Check if active input
477       if (vec == CEED_VECTOR_ACTIVE) {
478         CeedCallBackend(CeedQFunctionFieldGetSize(qfinputfields[i], &size));
479         CeedCallBackend(CeedVectorSetValue(impl->qvecsin[i], 0.0));
480         CeedCallBackend(CeedVectorGetArray(impl->qvecsin[i], CEED_MEM_DEVICE, &tmp));
481         CeedCallBackend(CeedRealloc(numactivein + size, &activein));
482         for (CeedInt field = 0; field < size; field++) {
483           q_size = (CeedSize)Q * numelements;
484           CeedCallBackend(CeedVectorCreate(ceed, q_size, &activein[numactivein + field]));
485           CeedCallBackend(CeedVectorSetArray(activein[numactivein + field], CEED_MEM_DEVICE, CEED_USE_POINTER, &tmp[field * Q * numelements]));
486         }
487         numactivein += size;
488         CeedCallBackend(CeedVectorRestoreArray(impl->qvecsin[i], &tmp));
489       }
490     }
491     impl->qfnumactivein = numactivein;
492     impl->qfactivein    = activein;
493   }
494 
495   // Count number of active output fields
496   if (!numactiveout) {
497     for (CeedInt i = 0; i < numoutputfields; i++) {
498       // Get output vector
499       CeedCallBackend(CeedOperatorFieldGetVector(opoutputfields[i], &vec));
500       // Check if active output
501       if (vec == CEED_VECTOR_ACTIVE) {
502         CeedCallBackend(CeedQFunctionFieldGetSize(qfoutputfields[i], &size));
503         numactiveout += size;
504       }
505     }
506     impl->qfnumactiveout = numactiveout;
507   }
508 
509   // Check sizes
510   CeedCheck(numactivein > 0 && numactiveout > 0, ceed, CEED_ERROR_BACKEND, "Cannot assemble QFunction without active inputs and outputs");
511 
512   // Build objects if needed
513   if (build_objects) {
514     // Create output restriction
515     CeedInt strides[3] = {1, numelements * Q, Q}; /* *NOPAD* */
516     CeedCallBackend(CeedElemRestrictionCreateStrided(ceedparent, numelements, Q, numactivein * numactiveout,
517                                                      numactivein * numactiveout * numelements * Q, strides, rstr));
518     // Create assembled vector
519     CeedSize l_size = (CeedSize)numelements * Q * numactivein * numactiveout;
520     CeedCallBackend(CeedVectorCreate(ceedparent, l_size, assembled));
521   }
522   CeedCallBackend(CeedVectorSetValue(*assembled, 0.0));
523   CeedCallBackend(CeedVectorGetArray(*assembled, CEED_MEM_DEVICE, &a));
524 
525   // Input basis apply
526   CeedCallBackend(CeedOperatorInputBasis_Cuda(numelements, qfinputfields, opinputfields, numinputfields, true, edata, impl));
527 
528   // Assemble QFunction
529   for (CeedInt in = 0; in < numactivein; in++) {
530     // Set Inputs
531     CeedCallBackend(CeedVectorSetValue(activein[in], 1.0));
532     if (numactivein > 1) {
533       CeedCallBackend(CeedVectorSetValue(activein[(in + numactivein - 1) % numactivein], 0.0));
534     }
535     // Set Outputs
536     for (CeedInt out = 0; out < numoutputfields; out++) {
537       // Get output vector
538       CeedCallBackend(CeedOperatorFieldGetVector(opoutputfields[out], &vec));
539       // Check if active output
540       if (vec == CEED_VECTOR_ACTIVE) {
541         CeedCallBackend(CeedVectorSetArray(impl->qvecsout[out], CEED_MEM_DEVICE, CEED_USE_POINTER, a));
542         CeedCallBackend(CeedQFunctionFieldGetSize(qfoutputfields[out], &size));
543         a += size * Q * numelements;  // Advance the pointer by the size of the output
544       }
545     }
546     // Apply QFunction
547     CeedCallBackend(CeedQFunctionApply(qf, Q * numelements, impl->qvecsin, impl->qvecsout));
548   }
549 
550   // Un-set output Qvecs to prevent accidental overwrite of Assembled
551   for (CeedInt out = 0; out < numoutputfields; out++) {
552     // Get output vector
553     CeedCallBackend(CeedOperatorFieldGetVector(opoutputfields[out], &vec));
554     // Check if active output
555     if (vec == CEED_VECTOR_ACTIVE) {
556       CeedCallBackend(CeedVectorTakeArray(impl->qvecsout[out], CEED_MEM_DEVICE, NULL));
557     }
558   }
559 
560   // Restore input arrays
561   CeedCallBackend(CeedOperatorRestoreInputs_Cuda(numinputfields, qfinputfields, opinputfields, true, edata, impl));
562 
563   // Restore output
564   CeedCallBackend(CeedVectorRestoreArray(*assembled, &a));
565 
566   return CEED_ERROR_SUCCESS;
567 }
568 
569 //------------------------------------------------------------------------------
570 // Assemble Linear QFunction
571 //------------------------------------------------------------------------------
572 static int CeedOperatorLinearAssembleQFunction_Cuda(CeedOperator op, CeedVector *assembled, CeedElemRestriction *rstr, CeedRequest *request) {
573   return CeedOperatorLinearAssembleQFunctionCore_Cuda(op, true, assembled, rstr, request);
574 }
575 
576 //------------------------------------------------------------------------------
577 // Update Assembled Linear QFunction
578 //------------------------------------------------------------------------------
579 static int CeedOperatorLinearAssembleQFunctionUpdate_Cuda(CeedOperator op, CeedVector assembled, CeedElemRestriction rstr, CeedRequest *request) {
580   return CeedOperatorLinearAssembleQFunctionCore_Cuda(op, false, &assembled, &rstr, request);
581 }
582 
583 //------------------------------------------------------------------------------
584 // Create point block restriction
585 //------------------------------------------------------------------------------
586 static int CreatePBRestriction(CeedElemRestriction rstr, CeedElemRestriction *pbRstr) {
587   Ceed ceed;
588   CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed));
589   const CeedInt *offsets;
590   CeedCallBackend(CeedElemRestrictionGetOffsets(rstr, CEED_MEM_HOST, &offsets));
591 
592   // Expand offsets
593   CeedInt  nelem, ncomp, elemsize, compstride, *pbOffsets;
594   CeedSize l_size;
595   CeedCallBackend(CeedElemRestrictionGetNumElements(rstr, &nelem));
596   CeedCallBackend(CeedElemRestrictionGetNumComponents(rstr, &ncomp));
597   CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elemsize));
598   CeedCallBackend(CeedElemRestrictionGetCompStride(rstr, &compstride));
599   CeedCallBackend(CeedElemRestrictionGetLVectorSize(rstr, &l_size));
600   CeedInt shift = ncomp;
601   if (compstride != 1) shift *= ncomp;
602   CeedCallBackend(CeedCalloc(nelem * elemsize, &pbOffsets));
603   for (CeedInt i = 0; i < nelem * elemsize; i++) {
604     pbOffsets[i] = offsets[i] * shift;
605   }
606 
607   // Create new restriction
608   CeedCallBackend(
609       CeedElemRestrictionCreate(ceed, nelem, elemsize, ncomp * ncomp, 1, l_size * ncomp, CEED_MEM_HOST, CEED_OWN_POINTER, pbOffsets, pbRstr));
610 
611   // Cleanup
612   CeedCallBackend(CeedElemRestrictionRestoreOffsets(rstr, &offsets));
613 
614   return CEED_ERROR_SUCCESS;
615 }
616 
617 //------------------------------------------------------------------------------
618 // Assemble diagonal setup
619 //------------------------------------------------------------------------------
620 static inline int CeedOperatorAssembleDiagonalSetup_Cuda(CeedOperator op, const bool pointBlock, CeedInt use_ceedsize_idx) {
621   Ceed ceed;
622   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
623   CeedQFunction qf;
624   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
625   CeedInt numinputfields, numoutputfields;
626   CeedCallBackend(CeedQFunctionGetNumArgs(qf, &numinputfields, &numoutputfields));
627 
628   // Determine active input basis
629   CeedOperatorField  *opfields;
630   CeedQFunctionField *qffields;
631   CeedCallBackend(CeedOperatorGetFields(op, NULL, &opfields, NULL, NULL));
632   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qffields, NULL, NULL));
633   CeedInt             numemodein = 0, ncomp = 0, dim = 1;
634   CeedEvalMode       *emodein = NULL;
635   CeedBasis           basisin = NULL;
636   CeedElemRestriction rstrin  = NULL;
637   for (CeedInt i = 0; i < numinputfields; i++) {
638     CeedVector vec;
639     CeedCallBackend(CeedOperatorFieldGetVector(opfields[i], &vec));
640     if (vec == CEED_VECTOR_ACTIVE) {
641       CeedElemRestriction rstr;
642       CeedCallBackend(CeedOperatorFieldGetBasis(opfields[i], &basisin));
643       CeedCallBackend(CeedBasisGetNumComponents(basisin, &ncomp));
644       CeedCallBackend(CeedBasisGetDimension(basisin, &dim));
645       CeedCallBackend(CeedOperatorFieldGetElemRestriction(opfields[i], &rstr));
646       CeedCheck(!rstrin || rstrin == rstr, ceed, CEED_ERROR_BACKEND,
647                 "Backend does not implement multi-field non-composite operator diagonal assembly");
648       rstrin = rstr;
649       CeedEvalMode emode;
650       CeedCallBackend(CeedQFunctionFieldGetEvalMode(qffields[i], &emode));
651       switch (emode) {
652         case CEED_EVAL_NONE:
653         case CEED_EVAL_INTERP:
654           CeedCallBackend(CeedRealloc(numemodein + 1, &emodein));
655           emodein[numemodein] = emode;
656           numemodein += 1;
657           break;
658         case CEED_EVAL_GRAD:
659           CeedCallBackend(CeedRealloc(numemodein + dim, &emodein));
660           for (CeedInt d = 0; d < dim; d++) emodein[numemodein + d] = emode;
661           numemodein += dim;
662           break;
663         case CEED_EVAL_WEIGHT:
664         case CEED_EVAL_DIV:
665         case CEED_EVAL_CURL:
666           break;  // Caught by QF Assembly
667       }
668     }
669   }
670 
671   // Determine active output basis
672   CeedCallBackend(CeedOperatorGetFields(op, NULL, NULL, NULL, &opfields));
673   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, NULL, NULL, &qffields));
674   CeedInt             numemodeout = 0;
675   CeedEvalMode       *emodeout    = NULL;
676   CeedBasis           basisout    = NULL;
677   CeedElemRestriction rstrout     = NULL;
678   for (CeedInt i = 0; i < numoutputfields; i++) {
679     CeedVector vec;
680     CeedCallBackend(CeedOperatorFieldGetVector(opfields[i], &vec));
681     if (vec == CEED_VECTOR_ACTIVE) {
682       CeedElemRestriction rstr;
683       CeedCallBackend(CeedOperatorFieldGetBasis(opfields[i], &basisout));
684       CeedCallBackend(CeedOperatorFieldGetElemRestriction(opfields[i], &rstr));
685       CeedCheck(!rstrout || rstrout == rstr, ceed, CEED_ERROR_BACKEND,
686                 "Backend does not implement multi-field non-composite operator diagonal assembly");
687       rstrout = rstr;
688       CeedEvalMode emode;
689       CeedCallBackend(CeedQFunctionFieldGetEvalMode(qffields[i], &emode));
690       switch (emode) {
691         case CEED_EVAL_NONE:
692         case CEED_EVAL_INTERP:
693           CeedCallBackend(CeedRealloc(numemodeout + 1, &emodeout));
694           emodeout[numemodeout] = emode;
695           numemodeout += 1;
696           break;
697         case CEED_EVAL_GRAD:
698           CeedCallBackend(CeedRealloc(numemodeout + dim, &emodeout));
699           for (CeedInt d = 0; d < dim; d++) emodeout[numemodeout + d] = emode;
700           numemodeout += dim;
701           break;
702         case CEED_EVAL_WEIGHT:
703         case CEED_EVAL_DIV:
704         case CEED_EVAL_CURL:
705           break;  // Caught by QF Assembly
706       }
707     }
708   }
709 
710   // Operator data struct
711   CeedOperator_Cuda *impl;
712   CeedCallBackend(CeedOperatorGetData(op, &impl));
713   CeedCallBackend(CeedCalloc(1, &impl->diag));
714   CeedOperatorDiag_Cuda *diag = impl->diag;
715   diag->basisin               = basisin;
716   diag->basisout              = basisout;
717   diag->h_emodein             = emodein;
718   diag->h_emodeout            = emodeout;
719   diag->numemodein            = numemodein;
720   diag->numemodeout           = numemodeout;
721 
722   // Assemble kernel
723   char *diagonal_kernel_path, *diagonal_kernel_source;
724   CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/cuda/cuda-ref-operator-assemble-diagonal.h", &diagonal_kernel_path));
725   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Diagonal Assembly Kernel Source -----\n");
726   CeedCallBackend(CeedLoadSourceToBuffer(ceed, diagonal_kernel_path, &diagonal_kernel_source));
727   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Diagonal Assembly Source Complete! -----\n");
728   CeedInt nnodes, nqpts;
729   CeedCallBackend(CeedBasisGetNumNodes(basisin, &nnodes));
730   CeedCallBackend(CeedBasisGetNumQuadraturePoints(basisin, &nqpts));
731   diag->nnodes = nnodes;
732   CeedCallCuda(ceed, CeedCompile_Cuda(ceed, diagonal_kernel_source, &diag->module, 6, "NUMEMODEIN", numemodein, "NUMEMODEOUT", numemodeout, "NNODES",
733                                       nnodes, "NQPTS", nqpts, "NCOMP", ncomp, "CEEDSIZE", use_ceedsize_idx));
734   CeedCallCuda(ceed, CeedGetKernel_Cuda(ceed, diag->module, "linearDiagonal", &diag->linearDiagonal));
735   CeedCallCuda(ceed, CeedGetKernel_Cuda(ceed, diag->module, "linearPointBlockDiagonal", &diag->linearPointBlock));
736   CeedCallBackend(CeedFree(&diagonal_kernel_path));
737   CeedCallBackend(CeedFree(&diagonal_kernel_source));
738 
739   // Basis matrices
740   const CeedInt     qBytes = nqpts * sizeof(CeedScalar);
741   const CeedInt     iBytes = qBytes * nnodes;
742   const CeedInt     gBytes = qBytes * nnodes * dim;
743   const CeedInt     eBytes = sizeof(CeedEvalMode);
744   const CeedScalar *interpin, *interpout, *gradin, *gradout;
745 
746   // CEED_EVAL_NONE
747   CeedScalar *identity = NULL;
748   bool        evalNone = false;
749   for (CeedInt i = 0; i < numemodein; i++) evalNone = evalNone || (emodein[i] == CEED_EVAL_NONE);
750   for (CeedInt i = 0; i < numemodeout; i++) evalNone = evalNone || (emodeout[i] == CEED_EVAL_NONE);
751   if (evalNone) {
752     CeedCallBackend(CeedCalloc(nqpts * nnodes, &identity));
753     for (CeedInt i = 0; i < (nnodes < nqpts ? nnodes : nqpts); i++) identity[i * nnodes + i] = 1.0;
754     CeedCallCuda(ceed, cudaMalloc((void **)&diag->d_identity, iBytes));
755     CeedCallCuda(ceed, cudaMemcpy(diag->d_identity, identity, iBytes, cudaMemcpyHostToDevice));
756   }
757 
758   // CEED_EVAL_INTERP
759   CeedCallBackend(CeedBasisGetInterp(basisin, &interpin));
760   CeedCallCuda(ceed, cudaMalloc((void **)&diag->d_interpin, iBytes));
761   CeedCallCuda(ceed, cudaMemcpy(diag->d_interpin, interpin, iBytes, cudaMemcpyHostToDevice));
762   CeedCallBackend(CeedBasisGetInterp(basisout, &interpout));
763   CeedCallCuda(ceed, cudaMalloc((void **)&diag->d_interpout, iBytes));
764   CeedCallCuda(ceed, cudaMemcpy(diag->d_interpout, interpout, iBytes, cudaMemcpyHostToDevice));
765 
766   // CEED_EVAL_GRAD
767   CeedCallBackend(CeedBasisGetGrad(basisin, &gradin));
768   CeedCallCuda(ceed, cudaMalloc((void **)&diag->d_gradin, gBytes));
769   CeedCallCuda(ceed, cudaMemcpy(diag->d_gradin, gradin, gBytes, cudaMemcpyHostToDevice));
770   CeedCallBackend(CeedBasisGetGrad(basisout, &gradout));
771   CeedCallCuda(ceed, cudaMalloc((void **)&diag->d_gradout, gBytes));
772   CeedCallCuda(ceed, cudaMemcpy(diag->d_gradout, gradout, gBytes, cudaMemcpyHostToDevice));
773 
774   // Arrays of emodes
775   CeedCallCuda(ceed, cudaMalloc((void **)&diag->d_emodein, numemodein * eBytes));
776   CeedCallCuda(ceed, cudaMemcpy(diag->d_emodein, emodein, numemodein * eBytes, cudaMemcpyHostToDevice));
777   CeedCallCuda(ceed, cudaMalloc((void **)&diag->d_emodeout, numemodeout * eBytes));
778   CeedCallCuda(ceed, cudaMemcpy(diag->d_emodeout, emodeout, numemodeout * eBytes, cudaMemcpyHostToDevice));
779 
780   // Restriction
781   diag->diagrstr = rstrout;
782 
783   return CEED_ERROR_SUCCESS;
784 }
785 
786 //------------------------------------------------------------------------------
787 // Assemble diagonal common code
788 //------------------------------------------------------------------------------
789 static inline int CeedOperatorAssembleDiagonalCore_Cuda(CeedOperator op, CeedVector assembled, CeedRequest *request, const bool pointBlock) {
790   Ceed ceed;
791   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
792   CeedOperator_Cuda *impl;
793   CeedCallBackend(CeedOperatorGetData(op, &impl));
794 
795   // Assemble QFunction
796   CeedVector          assembledqf = NULL;
797   CeedElemRestriction rstr        = NULL;
798   CeedCallBackend(CeedOperatorLinearAssembleQFunctionBuildOrUpdate(op, &assembledqf, &rstr, request));
799   CeedCallBackend(CeedElemRestrictionDestroy(&rstr));
800 
801   CeedSize assembled_length = 0, assembledqf_length = 0;
802   CeedCallBackend(CeedVectorGetLength(assembled, &assembled_length));
803   CeedCallBackend(CeedVectorGetLength(assembledqf, &assembledqf_length));
804   CeedInt use_ceedsize_idx = 0;
805   if ((assembled_length > INT_MAX) || (assembledqf_length > INT_MAX)) use_ceedsize_idx = 1;
806 
807   // Setup
808   if (!impl->diag) {
809     CeedCallBackend(CeedOperatorAssembleDiagonalSetup_Cuda(op, pointBlock, use_ceedsize_idx));
810   }
811   CeedOperatorDiag_Cuda *diag = impl->diag;
812   assert(diag != NULL);
813 
814   // Restriction
815   if (pointBlock && !diag->pbdiagrstr) {
816     CeedElemRestriction pbdiagrstr;
817     CeedCallBackend(CreatePBRestriction(diag->diagrstr, &pbdiagrstr));
818     diag->pbdiagrstr = pbdiagrstr;
819   }
820   CeedElemRestriction diagrstr = pointBlock ? diag->pbdiagrstr : diag->diagrstr;
821 
822   // Create diagonal vector
823   CeedVector elemdiag = pointBlock ? diag->pbelemdiag : diag->elemdiag;
824   if (!elemdiag) {
825     CeedCallBackend(CeedElemRestrictionCreateVector(diagrstr, NULL, &elemdiag));
826     if (pointBlock) diag->pbelemdiag = elemdiag;
827     else diag->elemdiag = elemdiag;
828   }
829   CeedCallBackend(CeedVectorSetValue(elemdiag, 0.0));
830 
831   // Assemble element operator diagonals
832   CeedScalar       *elemdiagarray;
833   const CeedScalar *assembledqfarray;
834   CeedCallBackend(CeedVectorGetArray(elemdiag, CEED_MEM_DEVICE, &elemdiagarray));
835   CeedCallBackend(CeedVectorGetArrayRead(assembledqf, CEED_MEM_DEVICE, &assembledqfarray));
836   CeedInt nelem;
837   CeedCallBackend(CeedElemRestrictionGetNumElements(diagrstr, &nelem));
838 
839   // Compute the diagonal of B^T D B
840   int   elemsPerBlock = 1;
841   int   grid          = nelem / elemsPerBlock + ((nelem / elemsPerBlock * elemsPerBlock < nelem) ? 1 : 0);
842   void *args[]        = {(void *)&nelem,   &diag->d_identity, &diag->d_interpin, &diag->d_gradin,   &diag->d_interpout,
843                          &diag->d_gradout, &diag->d_emodein,  &diag->d_emodeout, &assembledqfarray, &elemdiagarray};
844   if (pointBlock) {
845     CeedCallBackend(CeedRunKernelDim_Cuda(ceed, diag->linearPointBlock, grid, diag->nnodes, 1, elemsPerBlock, args));
846   } else {
847     CeedCallBackend(CeedRunKernelDim_Cuda(ceed, diag->linearDiagonal, grid, diag->nnodes, 1, elemsPerBlock, args));
848   }
849 
850   // Restore arrays
851   CeedCallBackend(CeedVectorRestoreArray(elemdiag, &elemdiagarray));
852   CeedCallBackend(CeedVectorRestoreArrayRead(assembledqf, &assembledqfarray));
853 
854   // Assemble local operator diagonal
855   CeedCallBackend(CeedElemRestrictionApply(diagrstr, CEED_TRANSPOSE, elemdiag, assembled, request));
856 
857   // Cleanup
858   CeedCallBackend(CeedVectorDestroy(&assembledqf));
859 
860   return CEED_ERROR_SUCCESS;
861 }
862 
863 //------------------------------------------------------------------------------
864 // Assemble Linear Diagonal
865 //------------------------------------------------------------------------------
866 static int CeedOperatorLinearAssembleAddDiagonal_Cuda(CeedOperator op, CeedVector assembled, CeedRequest *request) {
867   CeedCallBackend(CeedOperatorAssembleDiagonalCore_Cuda(op, assembled, request, false));
868   return CEED_ERROR_SUCCESS;
869 }
870 
871 //------------------------------------------------------------------------------
872 // Assemble Linear Point Block Diagonal
873 //------------------------------------------------------------------------------
874 static int CeedOperatorLinearAssembleAddPointBlockDiagonal_Cuda(CeedOperator op, CeedVector assembled, CeedRequest *request) {
875   CeedCallBackend(CeedOperatorAssembleDiagonalCore_Cuda(op, assembled, request, true));
876   return CEED_ERROR_SUCCESS;
877 }
878 
879 //------------------------------------------------------------------------------
880 // Single operator assembly setup
881 //------------------------------------------------------------------------------
882 static int CeedSingleOperatorAssembleSetup_Cuda(CeedOperator op, CeedInt use_ceedsize_idx) {
883   Ceed ceed;
884   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
885   CeedOperator_Cuda *impl;
886   CeedCallBackend(CeedOperatorGetData(op, &impl));
887 
888   // Get intput and output fields
889   CeedInt            num_input_fields, num_output_fields;
890   CeedOperatorField *input_fields;
891   CeedOperatorField *output_fields;
892   CeedCallBackend(CeedOperatorGetFields(op, &num_input_fields, &input_fields, &num_output_fields, &output_fields));
893 
894   // Determine active input basis eval mode
895   CeedQFunction qf;
896   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
897   CeedQFunctionField *qf_fields;
898   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_fields, NULL, NULL));
899   // Note that the kernel will treat each dimension of a gradient action separately;
900   // i.e., when an active input has a CEED_EVAL_GRAD mode, num_emode_in will increment by dim.
901   // However, for the purposes of loading the B matrices, it will be treated as one mode, and we will load/copy the entire gradient matrix at once, so
902   // num_B_in_mats_to_load will be incremented by 1.
903   CeedInt             num_emode_in = 0, dim = 1, num_B_in_mats_to_load = 0, size_B_in = 0;
904   CeedEvalMode       *eval_mode_in = NULL;  // will be of size num_B_in_mats_load
905   CeedBasis           basis_in     = NULL;
906   CeedInt             nqpts = 0, esize = 0;
907   CeedElemRestriction rstr_in = NULL;
908   for (CeedInt i = 0; i < num_input_fields; i++) {
909     CeedVector vec;
910     CeedCallBackend(CeedOperatorFieldGetVector(input_fields[i], &vec));
911     if (vec == CEED_VECTOR_ACTIVE) {
912       CeedCallBackend(CeedOperatorFieldGetBasis(input_fields[i], &basis_in));
913       CeedCallBackend(CeedBasisGetDimension(basis_in, &dim));
914       CeedCallBackend(CeedBasisGetNumQuadraturePoints(basis_in, &nqpts));
915       CeedCallBackend(CeedOperatorFieldGetElemRestriction(input_fields[i], &rstr_in));
916       CeedCallBackend(CeedElemRestrictionGetElementSize(rstr_in, &esize));
917       CeedEvalMode eval_mode;
918       CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_fields[i], &eval_mode));
919       if (eval_mode != CEED_EVAL_NONE) {
920         CeedCallBackend(CeedRealloc(num_B_in_mats_to_load + 1, &eval_mode_in));
921         eval_mode_in[num_B_in_mats_to_load] = eval_mode;
922         num_B_in_mats_to_load += 1;
923         if (eval_mode == CEED_EVAL_GRAD) {
924           num_emode_in += dim;
925           size_B_in += dim * esize * nqpts;
926         } else {
927           num_emode_in += 1;
928           size_B_in += esize * nqpts;
929         }
930       }
931     }
932   }
933 
934   // Determine active output basis; basis_out and rstr_out only used if same as input, TODO
935   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, NULL, NULL, &qf_fields));
936   CeedInt             num_emode_out = 0, num_B_out_mats_to_load = 0, size_B_out = 0;
937   CeedEvalMode       *eval_mode_out = NULL;
938   CeedBasis           basis_out     = NULL;
939   CeedElemRestriction rstr_out      = NULL;
940   for (CeedInt i = 0; i < num_output_fields; i++) {
941     CeedVector vec;
942     CeedCallBackend(CeedOperatorFieldGetVector(output_fields[i], &vec));
943     if (vec == CEED_VECTOR_ACTIVE) {
944       CeedCallBackend(CeedOperatorFieldGetBasis(output_fields[i], &basis_out));
945       CeedCallBackend(CeedOperatorFieldGetElemRestriction(output_fields[i], &rstr_out));
946       CeedCheck(!rstr_out || rstr_out == rstr_in, ceed, CEED_ERROR_BACKEND, "Backend does not implement multi-field non-composite operator assembly");
947       CeedEvalMode eval_mode;
948       CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_fields[i], &eval_mode));
949       if (eval_mode != CEED_EVAL_NONE) {
950         CeedCallBackend(CeedRealloc(num_B_out_mats_to_load + 1, &eval_mode_out));
951         eval_mode_out[num_B_out_mats_to_load] = eval_mode;
952         num_B_out_mats_to_load += 1;
953         if (eval_mode == CEED_EVAL_GRAD) {
954           num_emode_out += dim;
955           size_B_out += dim * esize * nqpts;
956         } else {
957           num_emode_out += 1;
958           size_B_out += esize * nqpts;
959         }
960       }
961     }
962   }
963   CeedCheck(num_emode_in > 0 && num_emode_out > 0, ceed, CEED_ERROR_UNSUPPORTED, "Cannot assemble operator without inputs/outputs");
964 
965   CeedInt nelem, ncomp;
966   CeedCallBackend(CeedElemRestrictionGetNumElements(rstr_in, &nelem));
967   CeedCallBackend(CeedElemRestrictionGetNumComponents(rstr_in, &ncomp));
968 
969   CeedCallBackend(CeedCalloc(1, &impl->asmb));
970   CeedOperatorAssemble_Cuda *asmb = impl->asmb;
971   asmb->nelem                     = nelem;
972 
973   // Compile kernels
974   int elemsPerBlock     = 1;
975   asmb->elemsPerBlock   = elemsPerBlock;
976   CeedInt    block_size = esize * esize * elemsPerBlock;
977   Ceed_Cuda *cuda_data;
978   CeedCallBackend(CeedGetData(ceed, &cuda_data));
979   char *assembly_kernel_path, *assembly_kernel_source;
980   CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/cuda/cuda-ref-operator-assemble.h", &assembly_kernel_path));
981   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Assembly Kernel Source -----\n");
982   CeedCallBackend(CeedLoadSourceToBuffer(ceed, assembly_kernel_path, &assembly_kernel_source));
983   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Assembly Source Complete! -----\n");
984   bool fallback = block_size > cuda_data->device_prop.maxThreadsPerBlock;
985   if (fallback) {
986     // Use fallback kernel with 1D threadblock
987     block_size         = esize * elemsPerBlock;
988     asmb->block_size_x = esize;
989     asmb->block_size_y = 1;
990   } else {  // Use kernel with 2D threadblock
991     asmb->block_size_x = esize;
992     asmb->block_size_y = esize;
993   }
994   CeedCallCuda(
995       ceed, CeedCompile_Cuda(ceed, assembly_kernel_source, &asmb->module, 8, "NELEM", nelem, "NUMEMODEIN", num_emode_in, "NUMEMODEOUT", num_emode_out,
996                              "NQPTS", nqpts, "NNODES", esize, "BLOCK_SIZE", block_size, "NCOMP", ncomp, "CEEDSIZE", use_ceedsize_idx));
997   CeedCallCuda(ceed, CeedGetKernel_Cuda(ceed, asmb->module, fallback ? "linearAssembleFallback" : "linearAssemble", &asmb->linearAssemble));
998   CeedCallBackend(CeedFree(&assembly_kernel_path));
999   CeedCallBackend(CeedFree(&assembly_kernel_source));
1000 
1001   // Build 'full' B matrices (not 1D arrays used for tensor-product matrices)
1002   const CeedScalar *interp_in, *grad_in;
1003   CeedCallBackend(CeedBasisGetInterp(basis_in, &interp_in));
1004   CeedCallBackend(CeedBasisGetGrad(basis_in, &grad_in));
1005 
1006   // Load into B_in, in order that they will be used in eval_mode
1007   const CeedInt inBytes   = size_B_in * sizeof(CeedScalar);
1008   CeedInt       mat_start = 0;
1009   CeedCallCuda(ceed, cudaMalloc((void **)&asmb->d_B_in, inBytes));
1010   for (int i = 0; i < num_B_in_mats_to_load; i++) {
1011     CeedEvalMode eval_mode = eval_mode_in[i];
1012     if (eval_mode == CEED_EVAL_INTERP) {
1013       CeedCallCuda(ceed, cudaMemcpy(&asmb->d_B_in[mat_start], interp_in, esize * nqpts * sizeof(CeedScalar), cudaMemcpyHostToDevice));
1014       mat_start += esize * nqpts;
1015     } else if (eval_mode == CEED_EVAL_GRAD) {
1016       CeedCallCuda(ceed, cudaMemcpy(&asmb->d_B_in[mat_start], grad_in, dim * esize * nqpts * sizeof(CeedScalar), cudaMemcpyHostToDevice));
1017       mat_start += dim * esize * nqpts;
1018     }
1019   }
1020 
1021   const CeedScalar *interp_out, *grad_out;
1022   // Note that this function currently assumes 1 basis, so this should always be true
1023   // for now
1024   if (basis_out == basis_in) {
1025     interp_out = interp_in;
1026     grad_out   = grad_in;
1027   } else {
1028     CeedCallBackend(CeedBasisGetInterp(basis_out, &interp_out));
1029     CeedCallBackend(CeedBasisGetGrad(basis_out, &grad_out));
1030   }
1031 
1032   // Load into B_out, in order that they will be used in eval_mode
1033   const CeedInt outBytes = size_B_out * sizeof(CeedScalar);
1034   mat_start              = 0;
1035   CeedCallCuda(ceed, cudaMalloc((void **)&asmb->d_B_out, outBytes));
1036   for (int i = 0; i < num_B_out_mats_to_load; i++) {
1037     CeedEvalMode eval_mode = eval_mode_out[i];
1038     if (eval_mode == CEED_EVAL_INTERP) {
1039       CeedCallCuda(ceed, cudaMemcpy(&asmb->d_B_out[mat_start], interp_out, esize * nqpts * sizeof(CeedScalar), cudaMemcpyHostToDevice));
1040       mat_start += esize * nqpts;
1041     } else if (eval_mode == CEED_EVAL_GRAD) {
1042       CeedCallCuda(ceed, cudaMemcpy(&asmb->d_B_out[mat_start], grad_out, dim * esize * nqpts * sizeof(CeedScalar), cudaMemcpyHostToDevice));
1043       mat_start += dim * esize * nqpts;
1044     }
1045   }
1046   return CEED_ERROR_SUCCESS;
1047 }
1048 
1049 //------------------------------------------------------------------------------
1050 // Assemble matrix data for COO matrix of assembled operator.
1051 // The sparsity pattern is set by CeedOperatorLinearAssembleSymbolic.
1052 //
1053 // Note that this (and other assembly routines) currently assume only one active input restriction/basis per operator (could have multiple basis eval
1054 // modes).
1055 // TODO: allow multiple active input restrictions/basis objects
1056 //------------------------------------------------------------------------------
1057 static int CeedSingleOperatorAssemble_Cuda(CeedOperator op, CeedInt offset, CeedVector values) {
1058   Ceed ceed;
1059   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
1060   CeedOperator_Cuda *impl;
1061   CeedCallBackend(CeedOperatorGetData(op, &impl));
1062 
1063   // Assemble QFunction
1064   CeedVector          assembled_qf = NULL;
1065   CeedElemRestriction rstr_q       = NULL;
1066   CeedCallBackend(CeedOperatorLinearAssembleQFunctionBuildOrUpdate(op, &assembled_qf, &rstr_q, CEED_REQUEST_IMMEDIATE));
1067   CeedCallBackend(CeedElemRestrictionDestroy(&rstr_q));
1068   CeedScalar *values_array;
1069   CeedCallBackend(CeedVectorGetArray(values, CEED_MEM_DEVICE, &values_array));
1070   values_array += offset;
1071   const CeedScalar *qf_array;
1072   CeedCallBackend(CeedVectorGetArrayRead(assembled_qf, CEED_MEM_DEVICE, &qf_array));
1073 
1074   CeedSize values_length = 0, assembled_qf_length = 0;
1075   CeedCallBackend(CeedVectorGetLength(values, &values_length));
1076   CeedCallBackend(CeedVectorGetLength(assembled_qf, &assembled_qf_length));
1077   CeedInt use_ceedsize_idx = 0;
1078   if ((values_length > INT_MAX) || (assembled_qf_length > INT_MAX)) use_ceedsize_idx = 1;
1079   // Setup
1080   if (!impl->asmb) {
1081     CeedCallBackend(CeedSingleOperatorAssembleSetup_Cuda(op, use_ceedsize_idx));
1082     assert(impl->asmb != NULL);
1083   }
1084 
1085   // Compute B^T D B
1086   const CeedInt nelem         = impl->asmb->nelem;
1087   const CeedInt elemsPerBlock = impl->asmb->elemsPerBlock;
1088   const CeedInt grid          = nelem / elemsPerBlock + ((nelem / elemsPerBlock * elemsPerBlock < nelem) ? 1 : 0);
1089   void         *args[]        = {&impl->asmb->d_B_in, &impl->asmb->d_B_out, &qf_array, &values_array};
1090   CeedCallBackend(
1091       CeedRunKernelDim_Cuda(ceed, impl->asmb->linearAssemble, grid, impl->asmb->block_size_x, impl->asmb->block_size_y, elemsPerBlock, args));
1092 
1093   // Restore arrays
1094   CeedCallBackend(CeedVectorRestoreArray(values, &values_array));
1095   CeedCallBackend(CeedVectorRestoreArrayRead(assembled_qf, &qf_array));
1096 
1097   // Cleanup
1098   CeedCallBackend(CeedVectorDestroy(&assembled_qf));
1099 
1100   return CEED_ERROR_SUCCESS;
1101 }
1102 
1103 //------------------------------------------------------------------------------
1104 // Create operator
1105 //------------------------------------------------------------------------------
1106 int CeedOperatorCreate_Cuda(CeedOperator op) {
1107   Ceed ceed;
1108   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
1109   CeedOperator_Cuda *impl;
1110 
1111   CeedCallBackend(CeedCalloc(1, &impl));
1112   CeedCallBackend(CeedOperatorSetData(op, impl));
1113 
1114   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleQFunction", CeedOperatorLinearAssembleQFunction_Cuda));
1115   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleQFunctionUpdate", CeedOperatorLinearAssembleQFunctionUpdate_Cuda));
1116   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleAddDiagonal", CeedOperatorLinearAssembleAddDiagonal_Cuda));
1117   CeedCallBackend(
1118       CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleAddPointBlockDiagonal", CeedOperatorLinearAssembleAddPointBlockDiagonal_Cuda));
1119   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "LinearAssembleSingle", CeedSingleOperatorAssemble_Cuda));
1120   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAdd", CeedOperatorApplyAdd_Cuda));
1121   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "Destroy", CeedOperatorDestroy_Cuda));
1122   return CEED_ERROR_SUCCESS;
1123 }
1124 
1125 //------------------------------------------------------------------------------
1126