xref: /libCEED/rust/libceed-sys/c-src/backends/cuda-shared/ceed-cuda-shared-basis.c (revision af0e6e89dc9eb1085789d576366358a2a9a69ecc)
1 // Copyright (c) 2017-2024, 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 <cuda.h>
12 #include <cuda_runtime.h>
13 #include <stdbool.h>
14 #include <stddef.h>
15 #include <string.h>
16 
17 #include "../cuda/ceed-cuda-common.h"
18 #include "../cuda/ceed-cuda-compile.h"
19 #include "ceed-cuda-shared.h"
20 
21 //------------------------------------------------------------------------------
22 // Apply tensor basis
23 //------------------------------------------------------------------------------
24 static int CeedBasisApplyTensorCore_Cuda_shared(CeedBasis basis, bool apply_add, const CeedInt num_elem, CeedTransposeMode t_mode,
25                                                 CeedEvalMode eval_mode, CeedVector u, CeedVector v) {
26   Ceed                   ceed;
27   Ceed_Cuda             *ceed_Cuda;
28   CeedInt                dim, num_comp;
29   const CeedScalar      *d_u;
30   CeedScalar            *d_v;
31   CeedBasis_Cuda_shared *data;
32 
33   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
34   CeedCallBackend(CeedGetData(ceed, &ceed_Cuda));
35   CeedCallBackend(CeedBasisGetData(basis, &data));
36   CeedCallBackend(CeedBasisGetDimension(basis, &dim));
37   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
38 
39   // Get read/write access to u, v
40   if (u != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
41   else CeedCheck(eval_mode == CEED_EVAL_WEIGHT, ceed, CEED_ERROR_BACKEND, "An input vector is required for this CeedEvalMode");
42   if (apply_add) {
43     CeedCallBackend(CeedVectorGetArray(v, CEED_MEM_DEVICE, &d_v));
44   } else {
45     CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
46   }
47 
48   // Apply basis operation
49   switch (eval_mode) {
50     case CEED_EVAL_INTERP: {
51       CeedInt P_1d, Q_1d;
52 
53       CeedCheck(data->d_interp_1d, ceed, CEED_ERROR_BACKEND, "%s not supported; interp_1d not set", CeedEvalModes[eval_mode]);
54       CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
55       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
56       CeedInt thread_1d = CeedIntMax(Q_1d, P_1d);
57 
58       void *interp_args[] = {(void *)&num_elem, &data->d_interp_1d, &d_u, &d_v};
59 
60       if (dim == 1) {
61         // avoid >512 total threads
62         CeedInt elems_per_block = CeedIntMin(ceed_Cuda->device_prop.maxThreadsDim[2], CeedIntMax(512 / thread_1d, 1));
63         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
64         CeedInt shared_mem      = elems_per_block * thread_1d * sizeof(CeedScalar);
65 
66         if (t_mode == CEED_TRANSPOSE) {
67           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->InterpTransposeAdd : data->InterpTranspose, grid, thread_1d, 1,
68                                                       elems_per_block, shared_mem, interp_args));
69         } else {
70           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->Interp, grid, thread_1d, 1, elems_per_block, shared_mem, interp_args));
71         }
72       } else if (dim == 2) {
73         const CeedInt opt_elems[7] = {0, 32, 8, 6, 4, 2, 8};
74         // elems_per_block must be at least 1
75         CeedInt elems_per_block = CeedIntMax(thread_1d < 7 ? opt_elems[thread_1d] / num_comp : 1, 1);
76         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
77         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
78 
79         if (t_mode == CEED_TRANSPOSE) {
80           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->InterpTransposeAdd : data->InterpTranspose, grid, thread_1d, thread_1d,
81                                                       elems_per_block, shared_mem, interp_args));
82         } else {
83           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->Interp, grid, thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
84         }
85       } else if (dim == 3) {
86         CeedInt elems_per_block = 1;
87         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
88         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
89 
90         if (t_mode == CEED_TRANSPOSE) {
91           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->InterpTransposeAdd : data->InterpTranspose, grid, thread_1d, thread_1d,
92                                                       elems_per_block, shared_mem, interp_args));
93         } else {
94           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->Interp, grid, thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
95         }
96       }
97     } break;
98     case CEED_EVAL_GRAD: {
99       CeedInt P_1d, Q_1d;
100 
101       CeedCheck(data->d_grad_1d, ceed, CEED_ERROR_BACKEND, "%s not supported; grad_1d not set", CeedEvalModes[eval_mode]);
102       CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
103       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
104       CeedInt     thread_1d = CeedIntMax(Q_1d, P_1d);
105       CeedScalar *d_grad_1d = data->d_grad_1d;
106 
107       if (data->d_collo_grad_1d) {
108         d_grad_1d = data->d_collo_grad_1d;
109       }
110       void *grad_args[] = {(void *)&num_elem, &data->d_interp_1d, &d_grad_1d, &d_u, &d_v};
111 
112       if (dim == 1) {
113         // avoid >512 total threads
114         CeedInt elems_per_block = CeedIntMin(ceed_Cuda->device_prop.maxThreadsDim[2], CeedIntMax(512 / thread_1d, 1));
115         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
116         CeedInt shared_mem      = elems_per_block * thread_1d * sizeof(CeedScalar);
117 
118         if (t_mode == CEED_TRANSPOSE) {
119           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->GradTransposeAdd : data->GradTranspose, grid, thread_1d, 1,
120                                                       elems_per_block, shared_mem, grad_args));
121         } else {
122           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->Grad, grid, thread_1d, 1, elems_per_block, shared_mem, grad_args));
123         }
124       } else if (dim == 2) {
125         const CeedInt opt_elems[7] = {0, 32, 8, 6, 4, 2, 8};
126         // elems_per_block must be at least 1
127         CeedInt elems_per_block = CeedIntMax(thread_1d < 7 ? opt_elems[thread_1d] / num_comp : 1, 1);
128         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
129         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
130 
131         if (t_mode == CEED_TRANSPOSE) {
132           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->GradTransposeAdd : data->GradTranspose, grid, thread_1d, thread_1d,
133                                                       elems_per_block, shared_mem, grad_args));
134         } else {
135           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->Grad, grid, thread_1d, thread_1d, elems_per_block, shared_mem, grad_args));
136         }
137       } else if (dim == 3) {
138         CeedInt elems_per_block = 1;
139         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
140         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
141 
142         if (t_mode == CEED_TRANSPOSE) {
143           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->GradTransposeAdd : data->GradTranspose, grid, thread_1d, thread_1d,
144                                                       elems_per_block, shared_mem, grad_args));
145         } else {
146           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->Grad, grid, thread_1d, thread_1d, elems_per_block, shared_mem, grad_args));
147         }
148       }
149     } break;
150     case CEED_EVAL_WEIGHT: {
151       CeedInt Q_1d;
152       CeedInt block_size = 32;
153 
154       CeedCheck(data->d_q_weight_1d, ceed, CEED_ERROR_BACKEND, "%s not supported; q_weights_1d not set", CeedEvalModes[eval_mode]);
155       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
156       void *weight_args[] = {(void *)&num_elem, (void *)&data->d_q_weight_1d, &d_v};
157       if (dim == 1) {
158         const CeedInt elems_per_block = block_size / Q_1d;
159         const CeedInt grid_size       = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
160 
161         CeedCallBackend(CeedRunKernelDim_Cuda(ceed, data->Weight, grid_size, Q_1d, elems_per_block, 1, weight_args));
162       } else if (dim == 2) {
163         const CeedInt opt_elems       = block_size / (Q_1d * Q_1d);
164         const CeedInt elems_per_block = opt_elems > 0 ? opt_elems : 1;
165         const CeedInt grid_size       = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
166 
167         CeedCallBackend(CeedRunKernelDim_Cuda(ceed, data->Weight, grid_size, Q_1d, Q_1d, elems_per_block, weight_args));
168       } else if (dim == 3) {
169         const CeedInt opt_elems       = block_size / (Q_1d * Q_1d);
170         const CeedInt elems_per_block = opt_elems > 0 ? opt_elems : 1;
171         const CeedInt grid_size       = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
172 
173         CeedCallBackend(CeedRunKernelDim_Cuda(ceed, data->Weight, grid_size, Q_1d, Q_1d, elems_per_block, weight_args));
174       }
175     } break;
176     case CEED_EVAL_NONE: /* handled separately below */
177       break;
178     // LCOV_EXCL_START
179     case CEED_EVAL_DIV:
180     case CEED_EVAL_CURL:
181       return CeedError(ceed, CEED_ERROR_BACKEND, "%s not supported", CeedEvalModes[eval_mode]);
182       // LCOV_EXCL_STOP
183   }
184 
185   // Restore vectors, cover CEED_EVAL_NONE
186   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
187   if (eval_mode == CEED_EVAL_NONE) CeedCallBackend(CeedVectorSetArray(v, CEED_MEM_DEVICE, CEED_COPY_VALUES, (CeedScalar *)d_u));
188   if (eval_mode != CEED_EVAL_WEIGHT) CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
189   CeedCallBackend(CeedDestroy(&ceed));
190   return CEED_ERROR_SUCCESS;
191 }
192 
193 static int CeedBasisApplyTensor_Cuda_shared(CeedBasis basis, const CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode eval_mode, CeedVector u,
194                                             CeedVector v) {
195   CeedCallBackend(CeedBasisApplyTensorCore_Cuda_shared(basis, false, num_elem, t_mode, eval_mode, u, v));
196   return CEED_ERROR_SUCCESS;
197 }
198 
199 static int CeedBasisApplyAddTensor_Cuda_shared(CeedBasis basis, const CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode eval_mode,
200                                                CeedVector u, CeedVector v) {
201   CeedCallBackend(CeedBasisApplyTensorCore_Cuda_shared(basis, true, num_elem, t_mode, eval_mode, u, v));
202   return CEED_ERROR_SUCCESS;
203 }
204 
205 //------------------------------------------------------------------------------
206 // Basis apply - tensor AtPoints
207 //------------------------------------------------------------------------------
208 static int CeedBasisApplyAtPointsCore_Cuda_shared(CeedBasis basis, bool apply_add, const CeedInt num_elem, const CeedInt *num_points,
209                                                   CeedTransposeMode t_mode, CeedEvalMode eval_mode, CeedVector x_ref, CeedVector u, CeedVector v) {
210   Ceed                   ceed;
211   Ceed_Cuda             *ceed_Cuda;
212   CeedInt                Q_1d, dim, num_comp, max_num_points = num_points[0];
213   const CeedInt          is_transpose = t_mode == CEED_TRANSPOSE;
214   const CeedScalar      *d_x, *d_u;
215   CeedScalar            *d_v;
216   CeedBasis_Cuda_shared *data;
217 
218   CeedCallBackend(CeedBasisGetData(basis, &data));
219   CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
220   CeedCallBackend(CeedBasisGetDimension(basis, &dim));
221   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
222 
223   // Weight handled separately
224   if (eval_mode == CEED_EVAL_WEIGHT) {
225     CeedCallBackend(CeedVectorSetValue(v, 1.0));
226     return CEED_ERROR_SUCCESS;
227   }
228 
229   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
230   CeedCallBackend(CeedGetData(ceed, &ceed_Cuda));
231 
232   // Check padded to uniform number of points per elem
233   for (CeedInt i = 1; i < num_elem; i++) max_num_points = CeedIntMax(max_num_points, num_points[i]);
234   {
235     CeedInt  q_comp;
236     CeedSize len, len_required;
237     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, eval_mode, &q_comp));
238     CeedCallBackend(CeedVectorGetLength(is_transpose ? u : v, &len));
239     len_required = (CeedSize)num_comp * (CeedSize)q_comp * (CeedSize)num_elem * (CeedSize)max_num_points;
240     CeedCheck(len >= len_required, ceed, CEED_ERROR_BACKEND,
241               "Vector at points must be padded to the same number of points in each element for BasisApplyAtPoints on GPU backends."
242               " Found %" CeedSize_FMT ", Required %" CeedSize_FMT,
243               len, len_required);
244   }
245 
246   // Move num_points array to device
247   if (is_transpose) {
248     const CeedInt num_bytes = num_elem * sizeof(CeedInt);
249 
250     if (num_elem != data->num_elem_at_points) {
251       data->num_elem_at_points = num_elem;
252 
253       if (data->d_points_per_elem) CeedCallCuda(ceed, cudaFree(data->d_points_per_elem));
254       CeedCallCuda(ceed, cudaMalloc((void **)&data->d_points_per_elem, num_bytes));
255       CeedCallBackend(CeedFree(&data->h_points_per_elem));
256       CeedCallBackend(CeedCalloc(num_elem, &data->h_points_per_elem));
257     }
258     if (memcmp(data->h_points_per_elem, num_points, num_bytes)) {
259       memcpy(data->h_points_per_elem, num_points, num_bytes);
260       CeedCallCuda(ceed, cudaMemcpy(data->d_points_per_elem, num_points, num_bytes, cudaMemcpyHostToDevice));
261     }
262   }
263 
264   // Build kernels if needed
265   if (data->num_points != max_num_points) {
266     CeedInt P_1d;
267 
268     CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
269     data->num_points = max_num_points;
270 
271     // -- Create interp matrix to Chebyshev coefficients
272     if (!data->d_chebyshev_interp_1d) {
273       CeedSize    interp_bytes;
274       CeedScalar *chebyshev_interp_1d;
275 
276       interp_bytes = P_1d * Q_1d * sizeof(CeedScalar);
277       CeedCallBackend(CeedCalloc(P_1d * Q_1d, &chebyshev_interp_1d));
278       CeedCallBackend(CeedBasisGetChebyshevInterp1D(basis, chebyshev_interp_1d));
279       CeedCallCuda(ceed, cudaMalloc((void **)&data->d_chebyshev_interp_1d, interp_bytes));
280       CeedCallCuda(ceed, cudaMemcpy(data->d_chebyshev_interp_1d, chebyshev_interp_1d, interp_bytes, cudaMemcpyHostToDevice));
281       CeedCallBackend(CeedFree(&chebyshev_interp_1d));
282     }
283 
284     // -- Compile kernels
285     const char basis_kernel_source[] = "// AtPoints basis source\n#include <ceed/jit-source/cuda/cuda-shared-basis-tensor-at-points.h>\n";
286     CeedInt    num_comp;
287 
288     if (data->moduleAtPoints) CeedCallCuda(ceed, cuModuleUnload(data->moduleAtPoints));
289     CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
290     CeedCallBackend(CeedCompile_Cuda(ceed, basis_kernel_source, &data->moduleAtPoints, 8, "BASIS_Q_1D", Q_1d, "BASIS_P_1D", P_1d, "T_1D",
291                                      CeedIntMax(Q_1d, P_1d), "BASIS_DIM", dim, "BASIS_NUM_COMP", num_comp, "BASIS_NUM_NODES", CeedIntPow(P_1d, dim),
292                                      "BASIS_NUM_QPTS", CeedIntPow(Q_1d, dim), "BASIS_NUM_PTS", max_num_points));
293     CeedCallBackend(CeedGetKernel_Cuda(ceed, data->moduleAtPoints, "InterpAtPoints", &data->InterpAtPoints));
294     CeedCallBackend(CeedGetKernel_Cuda(ceed, data->moduleAtPoints, "InterpTransposeAtPoints", &data->InterpTransposeAtPoints));
295     CeedCallBackend(CeedGetKernel_Cuda(ceed, data->moduleAtPoints, "InterpTransposeAddAtPoints", &data->InterpTransposeAddAtPoints));
296     CeedCallBackend(CeedGetKernel_Cuda(ceed, data->moduleAtPoints, "GradAtPoints", &data->GradAtPoints));
297     CeedCallBackend(CeedGetKernel_Cuda(ceed, data->moduleAtPoints, "GradTransposeAtPoints", &data->GradTransposeAtPoints));
298     CeedCallBackend(CeedGetKernel_Cuda(ceed, data->moduleAtPoints, "GradTransposeAddAtPoints", &data->GradTransposeAddAtPoints));
299   }
300 
301   // Get read/write access to u, v
302   CeedCallBackend(CeedVectorGetArrayRead(x_ref, CEED_MEM_DEVICE, &d_x));
303   if (u != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
304   else CeedCheck(eval_mode == CEED_EVAL_WEIGHT, ceed, CEED_ERROR_BACKEND, "An input vector is required for this CeedEvalMode");
305   if (apply_add) {
306     CeedCallBackend(CeedVectorGetArray(v, CEED_MEM_DEVICE, &d_v));
307   } else {
308     CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
309   }
310 
311   // Basis action
312   switch (eval_mode) {
313     case CEED_EVAL_INTERP: {
314       CeedInt P_1d, Q_1d;
315 
316       CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
317       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
318       CeedInt thread_1d = CeedIntMax(Q_1d, P_1d);
319 
320       void *interp_args[] = {(void *)&num_elem, &data->d_chebyshev_interp_1d, &data->d_points_per_elem, &d_x, &d_u, &d_v};
321 
322       if (dim == 1) {
323         // avoid >512 total threads
324         CeedInt elems_per_block = CeedIntMin(ceed_Cuda->device_prop.maxThreadsDim[2], CeedIntMax(512 / thread_1d, 1));
325         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
326         CeedInt shared_mem      = elems_per_block * thread_1d * sizeof(CeedScalar);
327 
328         if (is_transpose) {
329           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->InterpTransposeAddAtPoints : data->InterpTransposeAtPoints, grid,
330                                                       thread_1d, 1, elems_per_block, shared_mem, interp_args));
331         } else {
332           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->InterpAtPoints, grid, thread_1d, 1, elems_per_block, shared_mem, interp_args));
333         }
334       } else if (dim == 2) {
335         const CeedInt opt_elems[7] = {0, 32, 8, 6, 4, 2, 8};
336         // elems_per_block must be at least 1
337         CeedInt elems_per_block = CeedIntMax(thread_1d < 7 ? opt_elems[thread_1d] / num_comp : 1, 1);
338         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
339         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
340 
341         if (is_transpose) {
342           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->InterpTransposeAddAtPoints : data->InterpTransposeAtPoints, grid,
343                                                       thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
344         } else {
345           CeedCallBackend(
346               CeedRunKernelDimShared_Cuda(ceed, data->InterpAtPoints, grid, thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
347         }
348       } else if (dim == 3) {
349         CeedInt elems_per_block = 1;
350         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
351         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
352 
353         if (is_transpose) {
354           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->InterpTransposeAddAtPoints : data->InterpTransposeAtPoints, grid,
355                                                       thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
356         } else {
357           CeedCallBackend(
358               CeedRunKernelDimShared_Cuda(ceed, data->InterpAtPoints, grid, thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
359         }
360       }
361     } break;
362     case CEED_EVAL_GRAD: {
363       CeedInt P_1d, Q_1d;
364 
365       CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
366       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
367       CeedInt thread_1d = CeedIntMax(Q_1d, P_1d);
368 
369       void *grad_args[] = {(void *)&num_elem, &data->d_chebyshev_interp_1d, &data->d_points_per_elem, &d_x, &d_u, &d_v};
370 
371       if (dim == 1) {
372         // avoid >512 total threads
373         CeedInt elems_per_block = CeedIntMin(ceed_Cuda->device_prop.maxThreadsDim[2], CeedIntMax(512 / thread_1d, 1));
374         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
375         CeedInt shared_mem      = elems_per_block * thread_1d * sizeof(CeedScalar);
376 
377         if (is_transpose) {
378           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->GradTransposeAddAtPoints : data->GradTransposeAtPoints, grid, thread_1d,
379                                                       1, elems_per_block, shared_mem, grad_args));
380         } else {
381           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->GradAtPoints, grid, thread_1d, 1, elems_per_block, shared_mem, grad_args));
382         }
383       } else if (dim == 2) {
384         const CeedInt opt_elems[7] = {0, 32, 8, 6, 4, 2, 8};
385         // elems_per_block must be at least 1
386         CeedInt elems_per_block = CeedIntMax(thread_1d < 7 ? opt_elems[thread_1d] / num_comp : 1, 1);
387         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
388         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
389 
390         if (is_transpose) {
391           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->GradTransposeAddAtPoints : data->GradTransposeAtPoints, grid, thread_1d,
392                                                       thread_1d, elems_per_block, shared_mem, grad_args));
393         } else {
394           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->GradAtPoints, grid, thread_1d, thread_1d, elems_per_block, shared_mem, grad_args));
395         }
396       } else if (dim == 3) {
397         CeedInt elems_per_block = 1;
398         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
399         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
400 
401         if (is_transpose) {
402           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->GradTransposeAddAtPoints : data->GradTransposeAtPoints, grid, thread_1d,
403                                                       thread_1d, elems_per_block, shared_mem, grad_args));
404         } else {
405           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->GradAtPoints, grid, thread_1d, thread_1d, elems_per_block, shared_mem, grad_args));
406         }
407       }
408     } break;
409     case CEED_EVAL_WEIGHT:
410     case CEED_EVAL_NONE: /* handled separately below */
411       break;
412     // LCOV_EXCL_START
413     case CEED_EVAL_DIV:
414     case CEED_EVAL_CURL:
415       return CeedError(ceed, CEED_ERROR_BACKEND, "%s not supported", CeedEvalModes[eval_mode]);
416       // LCOV_EXCL_STOP
417   }
418 
419   // Restore vectors, cover CEED_EVAL_NONE
420   CeedCallBackend(CeedVectorRestoreArrayRead(x_ref, &d_x));
421   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
422   if (eval_mode == CEED_EVAL_NONE) CeedCallBackend(CeedVectorSetArray(v, CEED_MEM_DEVICE, CEED_COPY_VALUES, (CeedScalar *)d_u));
423   if (eval_mode != CEED_EVAL_WEIGHT) CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
424   CeedCallBackend(CeedDestroy(&ceed));
425   return CEED_ERROR_SUCCESS;
426 }
427 
428 static int CeedBasisApplyAtPoints_Cuda_shared(CeedBasis basis, const CeedInt num_elem, const CeedInt *num_points, CeedTransposeMode t_mode,
429                                               CeedEvalMode eval_mode, CeedVector x_ref, CeedVector u, CeedVector v) {
430   CeedCallBackend(CeedBasisApplyAtPointsCore_Cuda_shared(basis, false, num_elem, num_points, t_mode, eval_mode, x_ref, u, v));
431   return CEED_ERROR_SUCCESS;
432 }
433 
434 static int CeedBasisApplyAddAtPoints_Cuda_shared(CeedBasis basis, const CeedInt num_elem, const CeedInt *num_points, CeedTransposeMode t_mode,
435                                                  CeedEvalMode eval_mode, CeedVector x_ref, CeedVector u, CeedVector v) {
436   CeedCallBackend(CeedBasisApplyAtPointsCore_Cuda_shared(basis, true, num_elem, num_points, t_mode, eval_mode, x_ref, u, v));
437   return CEED_ERROR_SUCCESS;
438 }
439 
440 //------------------------------------------------------------------------------
441 // Apply non-tensor basis
442 //------------------------------------------------------------------------------
443 static int CeedBasisApplyNonTensorCore_Cuda_shared(CeedBasis basis, bool apply_add, const CeedInt num_elem, CeedTransposeMode t_mode,
444                                                    CeedEvalMode eval_mode, CeedVector u, CeedVector v) {
445   Ceed                   ceed;
446   Ceed_Cuda             *ceed_Cuda;
447   CeedInt                dim;
448   const CeedScalar      *d_u;
449   CeedScalar            *d_v;
450   CeedBasis_Cuda_shared *data;
451 
452   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
453   CeedCallBackend(CeedGetData(ceed, &ceed_Cuda));
454   CeedCallBackend(CeedBasisGetData(basis, &data));
455   CeedCallBackend(CeedBasisGetDimension(basis, &dim));
456 
457   // Get read/write access to u, v
458   if (u != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
459   else CeedCheck(eval_mode == CEED_EVAL_WEIGHT, ceed, CEED_ERROR_BACKEND, "An input vector is required for this CeedEvalMode");
460   if (apply_add) {
461     CeedCallBackend(CeedVectorGetArray(v, CEED_MEM_DEVICE, &d_v));
462   } else {
463     CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
464   }
465 
466   // Apply basis operation
467   switch (eval_mode) {
468     case CEED_EVAL_INTERP: {
469       CeedInt P, Q;
470 
471       CeedCheck(data->d_interp_1d, ceed, CEED_ERROR_BACKEND, "%s not supported; interp not set", CeedEvalModes[eval_mode]);
472       CeedCallBackend(CeedBasisGetNumNodes(basis, &P));
473       CeedCallBackend(CeedBasisGetNumQuadraturePoints(basis, &Q));
474       CeedInt thread = CeedIntMax(Q, P);
475 
476       void *interp_args[] = {(void *)&num_elem, &data->d_interp_1d, &d_u, &d_v};
477 
478       {
479         // avoid >512 total threads
480         CeedInt elems_per_block = CeedIntMin(ceed_Cuda->device_prop.maxThreadsDim[2], CeedIntMax(512 / thread, 1));
481         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
482         CeedInt shared_mem      = elems_per_block * thread * sizeof(CeedScalar);
483 
484         if (t_mode == CEED_TRANSPOSE) {
485           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->InterpTransposeAdd : data->InterpTranspose, grid, thread, 1,
486                                                       elems_per_block, shared_mem, interp_args));
487         } else {
488           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->Interp, grid, thread, 1, elems_per_block, shared_mem, interp_args));
489         }
490       }
491     } break;
492     case CEED_EVAL_GRAD: {
493       CeedInt P, Q;
494 
495       CeedCheck(data->d_grad_1d, ceed, CEED_ERROR_BACKEND, "%s not supported; grad not set", CeedEvalModes[eval_mode]);
496       CeedCallBackend(CeedBasisGetNumNodes(basis, &P));
497       CeedCallBackend(CeedBasisGetNumQuadraturePoints(basis, &Q));
498       CeedInt thread = CeedIntMax(Q, P);
499 
500       void *grad_args[] = {(void *)&num_elem, &data->d_grad_1d, &d_u, &d_v};
501 
502       {
503         // avoid >512 total threads
504         CeedInt elems_per_block = CeedIntMin(ceed_Cuda->device_prop.maxThreadsDim[2], CeedIntMax(512 / thread, 1));
505         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
506         CeedInt shared_mem      = elems_per_block * thread * sizeof(CeedScalar);
507 
508         if (t_mode == CEED_TRANSPOSE) {
509           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, apply_add ? data->GradTransposeAdd : data->GradTranspose, grid, thread, 1,
510                                                       elems_per_block, shared_mem, grad_args));
511         } else {
512           CeedCallBackend(CeedRunKernelDimShared_Cuda(ceed, data->Grad, grid, thread, 1, elems_per_block, shared_mem, grad_args));
513         }
514       }
515     } break;
516     case CEED_EVAL_WEIGHT: {
517       CeedInt P, Q;
518 
519       CeedCheck(data->d_q_weight_1d, ceed, CEED_ERROR_BACKEND, "%s not supported; q_weights not set", CeedEvalModes[eval_mode]);
520       CeedCallBackend(CeedBasisGetNumNodes(basis, &P));
521       CeedCallBackend(CeedBasisGetNumQuadraturePoints(basis, &Q));
522       CeedInt thread = CeedIntMax(Q, P);
523 
524       void *weight_args[] = {(void *)&num_elem, (void *)&data->d_q_weight_1d, &d_v};
525 
526       {
527         // avoid >512 total threads
528         CeedInt elems_per_block = CeedIntMin(ceed_Cuda->device_prop.maxThreadsDim[2], CeedIntMax(512 / thread, 1));
529         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
530 
531         CeedCallBackend(CeedRunKernelDim_Cuda(ceed, data->Weight, grid, thread, elems_per_block, 1, weight_args));
532       }
533     } break;
534     case CEED_EVAL_NONE: /* handled separately below */
535       break;
536     // LCOV_EXCL_START
537     case CEED_EVAL_DIV:
538     case CEED_EVAL_CURL:
539       return CeedError(ceed, CEED_ERROR_BACKEND, "%s not supported", CeedEvalModes[eval_mode]);
540       // LCOV_EXCL_STOP
541   }
542 
543   // Restore vectors, cover CEED_EVAL_NONE
544   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
545   if (eval_mode == CEED_EVAL_NONE) CeedCallBackend(CeedVectorSetArray(v, CEED_MEM_DEVICE, CEED_COPY_VALUES, (CeedScalar *)d_u));
546   if (eval_mode != CEED_EVAL_WEIGHT) CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
547   CeedCallBackend(CeedDestroy(&ceed));
548   return CEED_ERROR_SUCCESS;
549 }
550 
551 static int CeedBasisApplyNonTensor_Cuda_shared(CeedBasis basis, const CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode eval_mode,
552                                                CeedVector u, CeedVector v) {
553   CeedCallBackend(CeedBasisApplyNonTensorCore_Cuda_shared(basis, false, num_elem, t_mode, eval_mode, u, v));
554   return CEED_ERROR_SUCCESS;
555 }
556 
557 static int CeedBasisApplyAddNonTensor_Cuda_shared(CeedBasis basis, const CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode eval_mode,
558                                                   CeedVector u, CeedVector v) {
559   CeedCallBackend(CeedBasisApplyNonTensorCore_Cuda_shared(basis, true, num_elem, t_mode, eval_mode, u, v));
560   return CEED_ERROR_SUCCESS;
561 }
562 
563 //------------------------------------------------------------------------------
564 // Destroy basis
565 //------------------------------------------------------------------------------
566 static int CeedBasisDestroy_Cuda_shared(CeedBasis basis) {
567   Ceed                   ceed;
568   CeedBasis_Cuda_shared *data;
569 
570   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
571   CeedCallBackend(CeedBasisGetData(basis, &data));
572   CeedCallCuda(ceed, cuModuleUnload(data->module));
573   if (data->moduleAtPoints) CeedCallCuda(ceed, cuModuleUnload(data->moduleAtPoints));
574   if (data->d_q_weight_1d) CeedCallCuda(ceed, cudaFree(data->d_q_weight_1d));
575   CeedCallBackend(CeedFree(&data->h_points_per_elem));
576   if (data->d_points_per_elem) CeedCallCuda(ceed, cudaFree(data->d_points_per_elem));
577   CeedCallCuda(ceed, cudaFree(data->d_interp_1d));
578   CeedCallCuda(ceed, cudaFree(data->d_grad_1d));
579   CeedCallCuda(ceed, cudaFree(data->d_collo_grad_1d));
580   CeedCallCuda(ceed, cudaFree(data->d_chebyshev_interp_1d));
581   CeedCallBackend(CeedFree(&data));
582   CeedCallBackend(CeedDestroy(&ceed));
583   return CEED_ERROR_SUCCESS;
584 }
585 
586 //------------------------------------------------------------------------------
587 // Create tensor basis
588 //------------------------------------------------------------------------------
589 int CeedBasisCreateTensorH1_Cuda_shared(CeedInt dim, CeedInt P_1d, CeedInt Q_1d, const CeedScalar *interp_1d, const CeedScalar *grad_1d,
590                                         const CeedScalar *q_ref_1d, const CeedScalar *q_weight_1d, CeedBasis basis) {
591   Ceed                   ceed;
592   CeedInt                num_comp;
593   const CeedInt          q_bytes      = Q_1d * sizeof(CeedScalar);
594   const CeedInt          interp_bytes = q_bytes * P_1d;
595   CeedBasis_Cuda_shared *data;
596 
597   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
598   CeedCallBackend(CeedCalloc(1, &data));
599 
600   // Copy basis data to GPU
601   if (q_weight_1d) {
602     CeedCallCuda(ceed, cudaMalloc((void **)&data->d_q_weight_1d, q_bytes));
603     CeedCallCuda(ceed, cudaMemcpy(data->d_q_weight_1d, q_weight_1d, q_bytes, cudaMemcpyHostToDevice));
604   }
605   CeedCallCuda(ceed, cudaMalloc((void **)&data->d_interp_1d, interp_bytes));
606   CeedCallCuda(ceed, cudaMemcpy(data->d_interp_1d, interp_1d, interp_bytes, cudaMemcpyHostToDevice));
607   CeedCallCuda(ceed, cudaMalloc((void **)&data->d_grad_1d, interp_bytes));
608   CeedCallCuda(ceed, cudaMemcpy(data->d_grad_1d, grad_1d, interp_bytes, cudaMemcpyHostToDevice));
609 
610   // Compute collocated gradient and copy to GPU
611   data->d_collo_grad_1d    = NULL;
612   bool has_collocated_grad = dim == 3 && Q_1d >= P_1d;
613 
614   if (has_collocated_grad) {
615     CeedScalar *collo_grad_1d;
616 
617     CeedCallBackend(CeedMalloc(Q_1d * Q_1d, &collo_grad_1d));
618     CeedCallBackend(CeedBasisGetCollocatedGrad(basis, collo_grad_1d));
619     CeedCallCuda(ceed, cudaMalloc((void **)&data->d_collo_grad_1d, q_bytes * Q_1d));
620     CeedCallCuda(ceed, cudaMemcpy(data->d_collo_grad_1d, collo_grad_1d, q_bytes * Q_1d, cudaMemcpyHostToDevice));
621     CeedCallBackend(CeedFree(&collo_grad_1d));
622   }
623 
624   // Compile basis kernels
625   const char basis_kernel_source[] = "// Tensor basis source\n#include <ceed/jit-source/cuda/cuda-shared-basis-tensor.h>\n";
626 
627   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
628   CeedCallBackend(CeedCompile_Cuda(ceed, basis_kernel_source, &data->module, 8, "BASIS_Q_1D", Q_1d, "BASIS_P_1D", P_1d, "T_1D",
629                                    CeedIntMax(Q_1d, P_1d), "BASIS_DIM", dim, "BASIS_NUM_COMP", num_comp, "BASIS_NUM_NODES", CeedIntPow(P_1d, dim),
630                                    "BASIS_NUM_QPTS", CeedIntPow(Q_1d, dim), "BASIS_HAS_COLLOCATED_GRAD", has_collocated_grad));
631   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "Interp", &data->Interp));
632   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "InterpTranspose", &data->InterpTranspose));
633   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "InterpTransposeAdd", &data->InterpTransposeAdd));
634   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "Grad", &data->Grad));
635   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "GradTranspose", &data->GradTranspose));
636   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "GradTransposeAdd", &data->GradTransposeAdd));
637   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "Weight", &data->Weight));
638 
639   CeedCallBackend(CeedBasisSetData(basis, data));
640 
641   // Register backend functions
642   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyTensor_Cuda_shared));
643   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "ApplyAdd", CeedBasisApplyAddTensor_Cuda_shared));
644   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "ApplyAtPoints", CeedBasisApplyAtPoints_Cuda_shared));
645   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "ApplyAddAtPoints", CeedBasisApplyAddAtPoints_Cuda_shared));
646   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroy_Cuda_shared));
647   CeedCallBackend(CeedDestroy(&ceed));
648   return CEED_ERROR_SUCCESS;
649 }
650 
651 //------------------------------------------------------------------------------
652 // Create non-tensor basis
653 //------------------------------------------------------------------------------
654 int CeedBasisCreateH1_Cuda_shared(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp,
655                                   const CeedScalar *grad, const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
656   Ceed                   ceed;
657   CeedInt                num_comp, q_comp_interp, q_comp_grad;
658   const CeedInt          q_bytes = num_qpts * sizeof(CeedScalar);
659   CeedBasis_Cuda_shared *data;
660 
661   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
662 
663   // Check shared memory size
664   {
665     Ceed_Cuda *cuda_data;
666 
667     CeedCallBackend(CeedGetData(ceed, &cuda_data));
668     if (((size_t)num_nodes * (size_t)num_qpts * (size_t)dim + (size_t)CeedIntMax(num_nodes, num_qpts)) * sizeof(CeedScalar) >
669         cuda_data->device_prop.sharedMemPerBlock) {
670       CeedCallBackend(CeedBasisCreateH1Fallback(ceed, topo, dim, num_nodes, num_qpts, interp, grad, q_ref, q_weight, basis));
671       CeedCallBackend(CeedDestroy(&ceed));
672       return CEED_ERROR_SUCCESS;
673     }
674   }
675 
676   CeedCallBackend(CeedCalloc(1, &data));
677 
678   // Copy basis data to GPU
679   CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
680   CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_GRAD, &q_comp_grad));
681   if (q_weight) {
682     CeedCallCuda(ceed, cudaMalloc((void **)&data->d_q_weight_1d, q_bytes));
683     CeedCallCuda(ceed, cudaMemcpy(data->d_q_weight_1d, q_weight, q_bytes, cudaMemcpyHostToDevice));
684   }
685   if (interp) {
686     const CeedInt interp_bytes = q_bytes * num_nodes * q_comp_interp;
687 
688     CeedCallCuda(ceed, cudaMalloc((void **)&data->d_interp_1d, interp_bytes));
689     CeedCallCuda(ceed, cudaMemcpy(data->d_interp_1d, interp, interp_bytes, cudaMemcpyHostToDevice));
690   }
691   if (grad) {
692     const CeedInt grad_bytes = q_bytes * num_nodes * q_comp_grad;
693 
694     CeedCallCuda(ceed, cudaMalloc((void **)&data->d_grad_1d, grad_bytes));
695     CeedCallCuda(ceed, cudaMemcpy(data->d_grad_1d, grad, grad_bytes, cudaMemcpyHostToDevice));
696   }
697 
698   // Compile basis kernels
699   const char basis_kernel_source[] = "// Non-tensor basis source\n#include <ceed/jit-source/cuda/cuda-shared-basis-nontensor.h>\n";
700 
701   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
702   CeedCallBackend(CeedCompile_Cuda(ceed, basis_kernel_source, &data->module, 5, "BASIS_Q", num_qpts, "BASIS_P", num_nodes, "T_1D",
703                                    CeedIntMax(num_qpts, num_nodes), "BASIS_DIM", dim, "BASIS_NUM_COMP", num_comp));
704   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "Interp", &data->Interp));
705   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "InterpTranspose", &data->InterpTranspose));
706   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "InterpTransposeAdd", &data->InterpTransposeAdd));
707   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "Grad", &data->Grad));
708   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "GradTranspose", &data->GradTranspose));
709   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "GradTransposeAdd", &data->GradTransposeAdd));
710   CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, "Weight", &data->Weight));
711 
712   CeedCallBackend(CeedBasisSetData(basis, data));
713 
714   // Register backend functions
715   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Cuda_shared));
716   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "ApplyAdd", CeedBasisApplyAddNonTensor_Cuda_shared));
717   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroy_Cuda_shared));
718   CeedCallBackend(CeedDestroy(&ceed));
719   return CEED_ERROR_SUCCESS;
720 }
721 
722 //------------------------------------------------------------------------------
723