xref: /libCEED/backends/hip-shared/ceed-hip-shared-basis.c (revision b4280a96583940f87169e4a342af92c298f7bcf5)
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 <stdbool.h>
12 #include <stddef.h>
13 #include <string.h>
14 #include <hip/hip_runtime.h>
15 
16 #include "../hip/ceed-hip-common.h"
17 #include "../hip/ceed-hip-compile.h"
18 #include "ceed-hip-shared.h"
19 
20 //------------------------------------------------------------------------------
21 // Compute a block size based on required minimum threads
22 //------------------------------------------------------------------------------
23 static CeedInt ComputeBlockSizeFromRequirement(const CeedInt required) {
24   CeedInt maxSize     = 1024;  // Max total threads per block
25   CeedInt currentSize = 64;    // Start with one group
26 
27   while (currentSize < maxSize) {
28     if (currentSize > required) break;
29     else currentSize = currentSize * 2;
30   }
31   return currentSize;
32 }
33 
34 //------------------------------------------------------------------------------
35 // Compute required thread block sizes for basis kernels given P, Q, dim, and
36 // num_comp (num_comp not currently used, but may be again in other basis
37 // parallelization options)
38 //------------------------------------------------------------------------------
39 static int ComputeBasisThreadBlockSizes(const CeedInt dim, const CeedInt P_1d, const CeedInt Q_1d, const CeedInt num_comp, CeedInt *block_sizes) {
40   // Note that this will use the same block sizes for all dimensions when compiling,
41   // but as each basis object is defined for a particular dimension, we will never
42   // call any kernels except the ones for the dimension for which we have computed the
43   // block sizes.
44   const CeedInt thread_1d = CeedIntMax(P_1d, Q_1d);
45 
46   switch (dim) {
47     case 1: {
48       // Interp kernels:
49       block_sizes[0] = 256;
50 
51       // Grad kernels:
52       block_sizes[1] = 256;
53 
54       // Weight kernels:
55       block_sizes[2] = 256;
56     } break;
57     case 2: {
58       // Interp kernels:
59       CeedInt required = thread_1d * thread_1d;
60 
61       block_sizes[0] = CeedIntMax(256, ComputeBlockSizeFromRequirement(required));
62 
63       // Grad kernels: currently use same required minimum threads
64       block_sizes[1] = CeedIntMax(256, ComputeBlockSizeFromRequirement(required));
65 
66       // Weight kernels:
67       required       = CeedIntMax(64, Q_1d * Q_1d);
68       block_sizes[2] = CeedIntMax(256, ComputeBlockSizeFromRequirement(required));
69 
70     } break;
71     case 3: {
72       // Interp kernels:
73       CeedInt required = thread_1d * thread_1d;
74 
75       block_sizes[0] = CeedIntMax(256, ComputeBlockSizeFromRequirement(required));
76 
77       // Grad kernels: currently use same required minimum threads
78       block_sizes[1] = CeedIntMax(256, ComputeBlockSizeFromRequirement(required));
79 
80       // Weight kernels:
81       required       = Q_1d * Q_1d * Q_1d;
82       block_sizes[2] = CeedIntMax(256, ComputeBlockSizeFromRequirement(required));
83     }
84   }
85   return CEED_ERROR_SUCCESS;
86 }
87 
88 //------------------------------------------------------------------------------
89 // Apply basis
90 //------------------------------------------------------------------------------
91 static int CeedBasisApplyTensorCore_Hip_shared(CeedBasis basis, bool apply_add, const CeedInt num_elem, CeedTransposeMode t_mode,
92                                                CeedEvalMode eval_mode, CeedVector u, CeedVector v) {
93   Ceed                  ceed;
94   Ceed_Hip             *ceed_Hip;
95   CeedInt               dim, num_comp;
96   const CeedScalar     *d_u;
97   CeedScalar           *d_v;
98   CeedBasis_Hip_shared *data;
99 
100   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
101   CeedCallBackend(CeedGetData(ceed, &ceed_Hip));
102   CeedCallBackend(CeedBasisGetData(basis, &data));
103   CeedCallBackend(CeedBasisGetDimension(basis, &dim));
104   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
105 
106   // Get read/write access to u, v
107   if (u != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
108   else CeedCheck(eval_mode == CEED_EVAL_WEIGHT, ceed, CEED_ERROR_BACKEND, "An input vector is required for this CeedEvalMode");
109   if (apply_add) CeedCallBackend(CeedVectorGetArray(v, CEED_MEM_DEVICE, &d_v));
110   else CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
111 
112   // Apply basis operation
113   switch (eval_mode) {
114     case CEED_EVAL_INTERP: {
115       CeedInt P_1d, Q_1d;
116       CeedInt block_size = data->block_sizes[0];
117 
118       CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
119       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
120       CeedInt thread_1d     = CeedIntMax(Q_1d, P_1d);
121       void   *interp_args[] = {(void *)&num_elem, &data->d_interp_1d, &d_u, &d_v};
122 
123       if (dim == 1) {
124         CeedInt elems_per_block = 64 * thread_1d > 256 ? 256 / thread_1d : 64;
125         elems_per_block         = elems_per_block > 0 ? elems_per_block : 1;
126         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
127         CeedInt shared_mem      = elems_per_block * thread_1d * sizeof(CeedScalar);
128 
129         if (t_mode == CEED_TRANSPOSE) {
130           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, apply_add ? data->InterpTransposeAdd : data->InterpTranspose, grid, thread_1d, 1,
131                                                      elems_per_block, shared_mem, interp_args));
132         } else {
133           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->Interp, grid, thread_1d, 1, elems_per_block, shared_mem, interp_args));
134         }
135       } else if (dim == 2) {
136         // Check if required threads is small enough to do multiple elems
137         const CeedInt elems_per_block = CeedIntMax(block_size / (thread_1d * thread_1d), 1);
138         CeedInt       grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
139         CeedInt       shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
140 
141         if (t_mode == CEED_TRANSPOSE) {
142           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, apply_add ? data->InterpTransposeAdd : data->InterpTranspose, grid, thread_1d, thread_1d,
143                                                      elems_per_block, shared_mem, interp_args));
144         } else {
145           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->Interp, grid, thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
146         }
147       } else if (dim == 3) {
148         const CeedInt elems_per_block = CeedIntMax(block_size / (thread_1d * thread_1d), 1);
149         CeedInt       grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
150         CeedInt       shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
151 
152         if (t_mode == CEED_TRANSPOSE) {
153           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, apply_add ? data->InterpTransposeAdd : data->InterpTranspose, grid, thread_1d, thread_1d,
154                                                      elems_per_block, shared_mem, interp_args));
155         } else {
156           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->Interp, grid, thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
157         }
158       }
159     } break;
160     case CEED_EVAL_GRAD: {
161       CeedInt P_1d, Q_1d;
162       CeedInt block_size = data->block_sizes[1];
163 
164       CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
165       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
166       CeedInt     thread_1d = CeedIntMax(Q_1d, P_1d);
167       CeedScalar *d_grad_1d = data->d_grad_1d;
168 
169       if (data->d_collo_grad_1d) {
170         d_grad_1d = data->d_collo_grad_1d;
171       }
172       void *grad_args[] = {(void *)&num_elem, &data->d_interp_1d, &d_grad_1d, &d_u, &d_v};
173       if (dim == 1) {
174         CeedInt elems_per_block = 64 * thread_1d > 256 ? 256 / thread_1d : 64;
175         elems_per_block         = elems_per_block > 0 ? elems_per_block : 1;
176         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
177         CeedInt shared_mem      = elems_per_block * thread_1d * sizeof(CeedScalar);
178 
179         if (t_mode == CEED_TRANSPOSE) {
180           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, apply_add ? data->GradTransposeAdd : data->GradTranspose, grid, thread_1d, 1,
181                                                      elems_per_block, shared_mem, grad_args));
182         } else {
183           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->Grad, grid, thread_1d, 1, elems_per_block, shared_mem, grad_args));
184         }
185       } else if (dim == 2) {
186         // Check if required threads is small enough to do multiple elems
187         const CeedInt elems_per_block = CeedIntMax(block_size / (thread_1d * thread_1d), 1);
188         CeedInt       grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
189         CeedInt       shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
190 
191         if (t_mode == CEED_TRANSPOSE) {
192           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, apply_add ? data->GradTransposeAdd : data->GradTranspose, grid, thread_1d, thread_1d,
193                                                      elems_per_block, shared_mem, grad_args));
194         } else {
195           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->Grad, grid, thread_1d, thread_1d, elems_per_block, shared_mem, grad_args));
196         }
197       } else if (dim == 3) {
198         const CeedInt elems_per_block = CeedIntMax(block_size / (thread_1d * thread_1d), 1);
199         CeedInt       grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
200         CeedInt       shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
201 
202         if (t_mode == CEED_TRANSPOSE) {
203           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, apply_add ? data->GradTransposeAdd : data->GradTranspose, grid, thread_1d, thread_1d,
204                                                      elems_per_block, shared_mem, grad_args));
205         } else {
206           CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->Grad, grid, thread_1d, thread_1d, elems_per_block, shared_mem, grad_args));
207         }
208       }
209     } break;
210     case CEED_EVAL_WEIGHT: {
211       CeedInt Q_1d;
212       CeedInt block_size = data->block_sizes[2];
213 
214       CeedCheck(data->d_q_weight_1d, ceed, CEED_ERROR_BACKEND, "%s not supported; q_weights_1d not set", CeedEvalModes[eval_mode]);
215       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
216       void *weight_args[] = {(void *)&num_elem, (void *)&data->d_q_weight_1d, &d_v};
217 
218       if (dim == 1) {
219         const CeedInt opt_elems       = block_size / Q_1d;
220         const CeedInt elems_per_block = opt_elems > 0 ? opt_elems : 1;
221         const CeedInt grid_size       = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
222 
223         CeedCallBackend(CeedRunKernelDim_Hip(ceed, data->Weight, grid_size, Q_1d, elems_per_block, 1, weight_args));
224       } else if (dim == 2) {
225         const CeedInt opt_elems       = block_size / (Q_1d * Q_1d);
226         const CeedInt elems_per_block = opt_elems > 0 ? opt_elems : 1;
227         const CeedInt grid_size       = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
228 
229         CeedCallBackend(CeedRunKernelDim_Hip(ceed, data->Weight, grid_size, Q_1d, Q_1d, elems_per_block, weight_args));
230       } else if (dim == 3) {
231         const CeedInt opt_elems       = block_size / (Q_1d * Q_1d);
232         const CeedInt elems_per_block = opt_elems > 0 ? opt_elems : 1;
233         const CeedInt grid_size       = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
234 
235         CeedCallBackend(CeedRunKernelDim_Hip(ceed, data->Weight, grid_size, Q_1d, Q_1d, elems_per_block, weight_args));
236       }
237     } break;
238     case CEED_EVAL_NONE: /* handled separately below */
239       break;
240     // LCOV_EXCL_START
241     case CEED_EVAL_DIV:
242     case CEED_EVAL_CURL:
243       return CeedError(ceed, CEED_ERROR_BACKEND, "%s not supported", CeedEvalModes[eval_mode]);
244       // LCOV_EXCL_STOP
245   }
246 
247   // Restore vectors, cover CEED_EVAL_NONE
248   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
249   if (eval_mode == CEED_EVAL_NONE) CeedCallBackend(CeedVectorSetArray(v, CEED_MEM_DEVICE, CEED_COPY_VALUES, (CeedScalar *)d_u));
250   if (eval_mode != CEED_EVAL_WEIGHT) CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
251   CeedCallBackend(CeedDestroy(&ceed));
252   return CEED_ERROR_SUCCESS;
253 }
254 
255 int CeedBasisApplyTensor_Hip_shared(CeedBasis basis, const CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode eval_mode, CeedVector u,
256                                     CeedVector v) {
257   CeedCallBackend(CeedBasisApplyTensorCore_Hip_shared(basis, false, num_elem, t_mode, eval_mode, u, v));
258   return CEED_ERROR_SUCCESS;
259 }
260 
261 int CeedBasisApplyAddTensor_Hip_shared(CeedBasis basis, const CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode eval_mode, CeedVector u,
262                                        CeedVector v) {
263   CeedCallBackend(CeedBasisApplyTensorCore_Hip_shared(basis, true, num_elem, t_mode, eval_mode, u, v));
264   return CEED_ERROR_SUCCESS;
265 }
266 
267 //------------------------------------------------------------------------------
268 // Basis apply - tensor AtPoints
269 //------------------------------------------------------------------------------
270 static int CeedBasisApplyAtPointsCore_Hip_shared(CeedBasis basis, bool apply_add, const CeedInt num_elem, const CeedInt *num_points,
271                                                  CeedTransposeMode t_mode, CeedEvalMode eval_mode, CeedVector x_ref, CeedVector u, CeedVector v) {
272   Ceed                  ceed;
273   CeedInt               Q_1d, dim, max_num_points = num_points[0];
274   const CeedInt         is_transpose = t_mode == CEED_TRANSPOSE;
275   const CeedScalar     *d_x, *d_u;
276   CeedScalar           *d_v;
277   CeedBasis_Hip_shared *data;
278 
279   CeedCallBackend(CeedBasisGetData(basis, &data));
280   CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
281   CeedCallBackend(CeedBasisGetDimension(basis, &dim));
282 
283   // Weight handled separately
284   if (eval_mode == CEED_EVAL_WEIGHT) {
285     CeedCallBackend(CeedVectorSetValue(v, 1.0));
286     return CEED_ERROR_SUCCESS;
287   }
288 
289   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
290 
291   // Check padded to uniform number of points per elem
292   for (CeedInt i = 1; i < num_elem; i++) max_num_points = CeedIntMax(max_num_points, num_points[i]);
293   {
294     CeedInt  num_comp, q_comp;
295     CeedSize len, len_required;
296 
297     CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
298     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, eval_mode, &q_comp));
299     CeedCallBackend(CeedVectorGetLength(is_transpose ? u : v, &len));
300     len_required = (CeedSize)num_comp * (CeedSize)q_comp * (CeedSize)num_elem * (CeedSize)max_num_points;
301     CeedCheck(len >= len_required, ceed, CEED_ERROR_BACKEND,
302               "Vector at points must be padded to the same number of points in each element for BasisApplyAtPoints on GPU backends."
303               " Found %" CeedSize_FMT ", Required %" CeedSize_FMT,
304               len, len_required);
305   }
306 
307   // Move num_points array to device
308   if (is_transpose) {
309     const CeedInt num_bytes = num_elem * sizeof(CeedInt);
310 
311     if (num_elem != data->num_elem_at_points) {
312       data->num_elem_at_points = num_elem;
313 
314       if (data->d_points_per_elem) CeedCallHip(ceed, hipFree(data->d_points_per_elem));
315       CeedCallHip(ceed, hipMalloc((void **)&data->d_points_per_elem, num_bytes));
316       CeedCallBackend(CeedFree(&data->h_points_per_elem));
317       CeedCallBackend(CeedCalloc(num_elem, &data->h_points_per_elem));
318     }
319     if (memcmp(data->h_points_per_elem, num_points, num_bytes)) {
320       memcpy(data->h_points_per_elem, num_points, num_bytes);
321       CeedCallHip(ceed, hipMemcpy(data->d_points_per_elem, num_points, num_bytes, hipMemcpyHostToDevice));
322     }
323   }
324 
325   // Build kernels if needed
326   if (data->num_points != max_num_points) {
327     CeedInt P_1d;
328 
329     CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
330     data->num_points = max_num_points;
331 
332     // -- Create interp matrix to Chebyshev coefficients
333     if (!data->d_chebyshev_interp_1d) {
334       CeedSize    interp_bytes;
335       CeedScalar *chebyshev_interp_1d;
336 
337       interp_bytes = P_1d * Q_1d * sizeof(CeedScalar);
338       CeedCallBackend(CeedCalloc(P_1d * Q_1d, &chebyshev_interp_1d));
339       CeedCallBackend(CeedBasisGetChebyshevInterp1D(basis, chebyshev_interp_1d));
340       CeedCallHip(ceed, hipMalloc((void **)&data->d_chebyshev_interp_1d, interp_bytes));
341       CeedCallHip(ceed, hipMemcpy(data->d_chebyshev_interp_1d, chebyshev_interp_1d, interp_bytes, hipMemcpyHostToDevice));
342       CeedCallBackend(CeedFree(&chebyshev_interp_1d));
343     }
344 
345     // -- Compile kernels
346     const char basis_kernel_source[] = "// AtPoints basis source\n#include <ceed/jit-source/hip/hip-shared-basis-tensor-at-points.h>\n";
347     CeedInt    num_comp;
348 
349     if (data->moduleAtPoints) CeedCallHip(ceed, hipModuleUnload(data->moduleAtPoints));
350     CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
351     CeedCallBackend(CeedCompile_Hip(ceed, basis_kernel_source, &data->moduleAtPoints, 9, "BASIS_Q_1D", Q_1d, "BASIS_P_1D", P_1d, "T_1D",
352                                     CeedIntMax(Q_1d, P_1d), "BASIS_DIM", dim, "BASIS_NUM_COMP", num_comp, "BASIS_NUM_NODES", CeedIntPow(P_1d, dim),
353                                     "BASIS_NUM_QPTS", CeedIntPow(Q_1d, dim), "BASIS_NUM_PTS", max_num_points, "BASIS_INTERP_BLOCK_SIZE",
354                                     data->block_sizes[0]));
355     CeedCallBackend(CeedGetKernel_Hip(ceed, data->moduleAtPoints, "InterpAtPoints", &data->InterpAtPoints));
356     CeedCallBackend(CeedGetKernel_Hip(ceed, data->moduleAtPoints, "InterpTransposeAtPoints", &data->InterpTransposeAtPoints));
357     CeedCallBackend(CeedGetKernel_Hip(ceed, data->moduleAtPoints, "GradAtPoints", &data->GradAtPoints));
358     CeedCallBackend(CeedGetKernel_Hip(ceed, data->moduleAtPoints, "GradTransposeAtPoints", &data->GradTransposeAtPoints));
359   }
360 
361   // Get read/write access to u, v
362   CeedCallBackend(CeedVectorGetArrayRead(x_ref, CEED_MEM_DEVICE, &d_x));
363   if (u != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
364   else CeedCheck(eval_mode == CEED_EVAL_WEIGHT, ceed, CEED_ERROR_BACKEND, "An input vector is required for this CeedEvalMode");
365   if (apply_add) CeedCallBackend(CeedVectorGetArray(v, CEED_MEM_DEVICE, &d_v));
366   else CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
367 
368   // Clear v for transpose operation
369   if (is_transpose && !apply_add) {
370     CeedInt  num_comp, q_comp, num_nodes;
371     CeedSize length;
372 
373     CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
374     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, eval_mode, &q_comp));
375     CeedCallBackend(CeedBasisGetNumNodes(basis, &num_nodes));
376     length =
377         (CeedSize)num_elem * (CeedSize)num_comp * (t_mode == CEED_TRANSPOSE ? (CeedSize)num_nodes : ((CeedSize)max_num_points * (CeedSize)q_comp));
378     CeedCallHip(ceed, hipMemset(d_v, 0, length * sizeof(CeedScalar)));
379   }
380 
381   // Basis action
382   switch (eval_mode) {
383     case CEED_EVAL_INTERP: {
384       CeedInt P_1d, Q_1d;
385       CeedInt block_size = data->block_sizes[0];
386 
387       CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
388       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
389       CeedInt thread_1d     = CeedIntMax(Q_1d, P_1d);
390       void   *interp_args[] = {(void *)&num_elem, &data->d_chebyshev_interp_1d, &data->d_points_per_elem, &d_x, &d_u, &d_v};
391 
392       if (dim == 1) {
393         CeedInt elems_per_block = 64 * thread_1d > 256 ? 256 / thread_1d : 64;
394         elems_per_block         = elems_per_block > 0 ? elems_per_block : 1;
395         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
396         CeedInt shared_mem      = elems_per_block * thread_1d * sizeof(CeedScalar);
397 
398         CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, is_transpose ? data->InterpTransposeAtPoints : data->InterpAtPoints, grid, thread_1d, 1,
399                                                    elems_per_block, shared_mem, interp_args));
400       } else if (dim == 2) {
401         // Check if required threads is small enough to do multiple elems
402         const CeedInt elems_per_block = CeedIntMax(block_size / (thread_1d * thread_1d), 1);
403         CeedInt       grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
404         CeedInt       shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
405 
406         CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, is_transpose ? data->InterpTransposeAtPoints : data->InterpAtPoints, grid, thread_1d,
407                                                    thread_1d, elems_per_block, shared_mem, interp_args));
408       } else if (dim == 3) {
409         const CeedInt elems_per_block = 1;
410         CeedInt       grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
411         CeedInt       shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
412 
413         CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, is_transpose ? data->InterpTransposeAtPoints : data->InterpAtPoints, grid, thread_1d,
414                                                    thread_1d, elems_per_block, shared_mem, interp_args));
415       }
416     } break;
417     case CEED_EVAL_GRAD: {
418       CeedInt P_1d, Q_1d;
419       CeedInt block_size = data->block_sizes[0];
420 
421       CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
422       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
423       CeedInt thread_1d   = CeedIntMax(Q_1d, P_1d);
424       void   *grad_args[] = {(void *)&num_elem, &data->d_chebyshev_interp_1d, &data->d_points_per_elem, &d_x, &d_u, &d_v};
425 
426       if (dim == 1) {
427         CeedInt elems_per_block = 64 * thread_1d > 256 ? 256 / thread_1d : 64;
428         elems_per_block         = elems_per_block > 0 ? elems_per_block : 1;
429         CeedInt grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
430         CeedInt shared_mem      = elems_per_block * thread_1d * sizeof(CeedScalar);
431 
432         CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, is_transpose ? data->GradTransposeAtPoints : data->GradAtPoints, grid, thread_1d, 1,
433                                                    elems_per_block, shared_mem, grad_args));
434       } else if (dim == 2) {
435         // Check if required threads is small enough to do multiple elems
436         const CeedInt elems_per_block = CeedIntMax(block_size / (thread_1d * thread_1d), 1);
437         CeedInt       grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
438         CeedInt       shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
439 
440         CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, is_transpose ? data->GradTransposeAtPoints : data->GradAtPoints, grid, thread_1d, thread_1d,
441                                                    elems_per_block, shared_mem, grad_args));
442       } else if (dim == 3) {
443         const CeedInt elems_per_block = 1;
444         CeedInt       grid            = num_elem / elems_per_block + (num_elem % elems_per_block > 0);
445         CeedInt       shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
446 
447         CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, is_transpose ? data->GradTransposeAtPoints : data->GradAtPoints, grid, thread_1d, thread_1d,
448                                                    elems_per_block, shared_mem, grad_args));
449       }
450     } break;
451     case CEED_EVAL_WEIGHT:
452     case CEED_EVAL_NONE: /* handled separately below */
453       break;
454     // LCOV_EXCL_START
455     case CEED_EVAL_DIV:
456     case CEED_EVAL_CURL:
457       return CeedError(ceed, CEED_ERROR_BACKEND, "%s not supported", CeedEvalModes[eval_mode]);
458       // LCOV_EXCL_STOP
459   }
460 
461   // Restore vectors, cover CEED_EVAL_NONE
462   CeedCallBackend(CeedVectorRestoreArrayRead(x_ref, &d_x));
463   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
464   if (eval_mode == CEED_EVAL_NONE) CeedCallBackend(CeedVectorSetArray(v, CEED_MEM_DEVICE, CEED_COPY_VALUES, (CeedScalar *)d_u));
465   if (eval_mode != CEED_EVAL_WEIGHT) CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
466   return CEED_ERROR_SUCCESS;
467 }
468 
469 static int CeedBasisApplyAtPoints_Hip_shared(CeedBasis basis, const CeedInt num_elem, const CeedInt *num_points, CeedTransposeMode t_mode,
470                                              CeedEvalMode eval_mode, CeedVector x_ref, CeedVector u, CeedVector v) {
471   CeedCallBackend(CeedBasisApplyAtPointsCore_Hip_shared(basis, false, num_elem, num_points, t_mode, eval_mode, x_ref, u, v));
472   return CEED_ERROR_SUCCESS;
473 }
474 
475 static int CeedBasisApplyAddAtPoints_Hip_shared(CeedBasis basis, const CeedInt num_elem, const CeedInt *num_points, CeedTransposeMode t_mode,
476                                                 CeedEvalMode eval_mode, CeedVector x_ref, CeedVector u, CeedVector v) {
477   CeedCallBackend(CeedBasisApplyAtPointsCore_Hip_shared(basis, true, num_elem, num_points, t_mode, eval_mode, x_ref, u, v));
478   return CEED_ERROR_SUCCESS;
479 }
480 
481 //------------------------------------------------------------------------------
482 // Destroy basis
483 //------------------------------------------------------------------------------
484 static int CeedBasisDestroy_Hip_shared(CeedBasis basis) {
485   Ceed                  ceed;
486   CeedBasis_Hip_shared *data;
487 
488   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
489   CeedCallBackend(CeedBasisGetData(basis, &data));
490   CeedCallHip(ceed, hipModuleUnload(data->module));
491   if (data->moduleAtPoints) CeedCallHip(ceed, hipModuleUnload(data->moduleAtPoints));
492   if (data->d_q_weight_1d) CeedCallHip(ceed, hipFree(data->d_q_weight_1d));
493   CeedCallBackend(CeedFree(&data->h_points_per_elem));
494   if (data->d_points_per_elem) CeedCallHip(ceed, hipFree(data->d_points_per_elem));
495   CeedCallHip(ceed, hipFree(data->d_interp_1d));
496   CeedCallHip(ceed, hipFree(data->d_grad_1d));
497   CeedCallHip(ceed, hipFree(data->d_collo_grad_1d));
498   CeedCallHip(ceed, hipFree(data->d_chebyshev_interp_1d));
499   CeedCallBackend(CeedFree(&data));
500   return CEED_ERROR_SUCCESS;
501 }
502 
503 //------------------------------------------------------------------------------
504 // Create tensor basis
505 //------------------------------------------------------------------------------
506 int CeedBasisCreateTensorH1_Hip_shared(CeedInt dim, CeedInt P_1d, CeedInt Q_1d, const CeedScalar *interp_1d, const CeedScalar *grad_1d,
507                                        const CeedScalar *q_ref_1d, const CeedScalar *q_weight_1d, CeedBasis basis) {
508   Ceed                  ceed;
509   CeedInt               num_comp;
510   const CeedInt         q_bytes      = Q_1d * sizeof(CeedScalar);
511   const CeedInt         interp_bytes = q_bytes * P_1d;
512   CeedBasis_Hip_shared *data;
513 
514   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
515   CeedCallBackend(CeedCalloc(1, &data));
516 
517   // Copy basis data to GPU
518   if (q_weight_1d) {
519     CeedCallHip(ceed, hipMalloc((void **)&data->d_q_weight_1d, q_bytes));
520     CeedCallHip(ceed, hipMemcpy(data->d_q_weight_1d, q_weight_1d, q_bytes, hipMemcpyHostToDevice));
521   }
522   CeedCallHip(ceed, hipMalloc((void **)&data->d_interp_1d, interp_bytes));
523   CeedCallHip(ceed, hipMemcpy(data->d_interp_1d, interp_1d, interp_bytes, hipMemcpyHostToDevice));
524   CeedCallHip(ceed, hipMalloc((void **)&data->d_grad_1d, interp_bytes));
525   CeedCallHip(ceed, hipMemcpy(data->d_grad_1d, grad_1d, interp_bytes, hipMemcpyHostToDevice));
526 
527   // Compute collocated gradient and copy to GPU
528   data->d_collo_grad_1d    = NULL;
529   bool has_collocated_grad = dim == 3 && Q_1d >= P_1d;
530 
531   if (has_collocated_grad) {
532     CeedScalar *collo_grad_1d;
533 
534     CeedCallBackend(CeedMalloc(Q_1d * Q_1d, &collo_grad_1d));
535     CeedCallBackend(CeedBasisGetCollocatedGrad(basis, collo_grad_1d));
536     CeedCallHip(ceed, hipMalloc((void **)&data->d_collo_grad_1d, q_bytes * Q_1d));
537     CeedCallHip(ceed, hipMemcpy(data->d_collo_grad_1d, collo_grad_1d, q_bytes * Q_1d, hipMemcpyHostToDevice));
538     CeedCallBackend(CeedFree(&collo_grad_1d));
539   }
540 
541   // Set number of threads per block for basis kernels
542   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
543   CeedCallBackend(ComputeBasisThreadBlockSizes(dim, P_1d, Q_1d, num_comp, data->block_sizes));
544 
545   // Compile basis kernels
546   const char basis_kernel_source[] = "// Tensor basis source\n#include <ceed/jit-source/hip/hip-shared-basis-tensor.h>\n";
547 
548   CeedCallBackend(CeedCompile_Hip(ceed, basis_kernel_source, &data->module, 11, "BASIS_Q_1D", Q_1d, "BASIS_P_1D", P_1d, "T_1D",
549                                   CeedIntMax(Q_1d, P_1d), "BASIS_DIM", dim, "BASIS_NUM_COMP", num_comp, "BASIS_NUM_NODES", CeedIntPow(P_1d, dim),
550                                   "BASIS_NUM_QPTS", CeedIntPow(Q_1d, dim), "BASIS_INTERP_BLOCK_SIZE", data->block_sizes[0], "BASIS_GRAD_BLOCK_SIZE",
551                                   data->block_sizes[1], "BASIS_WEIGHT_BLOCK_SIZE", data->block_sizes[2], "BASIS_HAS_COLLOCATED_GRAD",
552                                   has_collocated_grad));
553   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "Interp", &data->Interp));
554   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "InterpTranspose", &data->InterpTranspose));
555   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "InterpTransposeAdd", &data->InterpTransposeAdd));
556   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "Grad", &data->Grad));
557   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "GradTranspose", &data->GradTranspose));
558   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "GradTransposeAdd", &data->GradTransposeAdd));
559   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "Weight", &data->Weight));
560 
561   CeedCallBackend(CeedBasisSetData(basis, data));
562 
563   // Register backend functions
564   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyTensor_Hip_shared));
565   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "ApplyAdd", CeedBasisApplyAddTensor_Hip_shared));
566   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "ApplyAtPoints", CeedBasisApplyAtPoints_Hip_shared));
567   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "ApplyAddAtPoints", CeedBasisApplyAddAtPoints_Hip_shared));
568   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroy_Hip_shared));
569   CeedCallBackend(CeedDestroy(&ceed));
570   return CEED_ERROR_SUCCESS;
571 }
572 
573 //------------------------------------------------------------------------------
574