xref: /libCEED/backends/magma/ceed-magma-basis.c (revision a745612c697a8f44e29c5a5127ed24c24f3757e8)
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 <string.h>
12 
13 #ifdef CEED_MAGMA_USE_HIP
14 #include "../hip/ceed-hip-common.h"
15 #include "../hip/ceed-hip-compile.h"
16 #else
17 #include "../cuda/ceed-cuda-common.h"
18 #include "../cuda/ceed-cuda-compile.h"
19 #endif
20 #include "ceed-magma-common.h"
21 #include "ceed-magma.h"
22 
23 #include "ceed-magma-gemm-nontensor.h"
24 #include "ceed-magma-gemm-selector.h"
25 
26 //------------------------------------------------------------------------------
27 // Basis apply - tensor
28 //------------------------------------------------------------------------------
29 static int CeedBasisApply_Magma(CeedBasis basis, CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode e_mode, CeedVector u, CeedVector v) {
30   Ceed              ceed;
31   Ceed_Magma       *data;
32   CeedInt           dim, num_comp, num_nodes, P_1d, Q_1d, P, Q;
33   const CeedScalar *d_u;
34   CeedScalar       *d_v;
35   CeedBasis_Magma  *impl;
36 
37   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
38   CeedCallBackend(CeedGetData(ceed, &data));
39   CeedCallBackend(CeedBasisGetData(basis, &impl));
40   CeedCallBackend(CeedBasisGetDimension(basis, &dim));
41   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
42   CeedCallBackend(CeedBasisGetNumNodes(basis, &num_nodes));
43   CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
44   CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
45   P = P_1d;
46   Q = Q_1d;
47   if (t_mode == CEED_TRANSPOSE) {
48     P = Q_1d;
49     Q = P_1d;
50   }
51 
52   // Read vectors
53   if (u != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
54   else CeedCheck(e_mode == CEED_EVAL_WEIGHT, ceed, CEED_ERROR_BACKEND, "An input vector is required for this CeedEvalMode");
55   CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
56 
57   // Apply basis operation
58   switch (e_mode) {
59     case CEED_EVAL_INTERP: {
60       // Define element sizes for dofs/quad
61       CeedInt elem_qpts_size = CeedIntPow(Q_1d, dim);
62       CeedInt elem_dofs_size = CeedIntPow(P_1d, dim);
63 
64       // E-vector ordering -------------- Q-vector ordering
65       //  component                        component
66       //    elem                             elem
67       //       node                            node
68 
69       // ---  Define strides for NOTRANSPOSE mode: ---
70       // Input (d_u) is E-vector, output (d_v) is Q-vector
71 
72       // Element strides
73       CeedInt u_elem_stride = elem_dofs_size;
74       CeedInt v_elem_stride = elem_qpts_size;
75       // Component strides
76       CeedInt u_comp_stride = num_elem * elem_dofs_size;
77       CeedInt v_comp_stride = num_elem * elem_qpts_size;
78       if (t_mode == CEED_TRANSPOSE) {
79         // Input (d_u) is Q-vector, output (d_v) is E-vector
80         // Element strides
81         v_elem_stride = elem_dofs_size;
82         u_elem_stride = elem_qpts_size;
83         // Component strides
84         v_comp_stride = num_elem * elem_dofs_size;
85         u_comp_stride = num_elem * elem_qpts_size;
86       }
87       CeedInt num_threads = 1;
88       CeedInt num_t_col   = 1;
89       CeedInt shared_mem  = 0;
90       CeedInt max_P_Q     = CeedIntMax(P, Q);
91 
92       switch (dim) {
93         case 1:
94           num_threads = max_P_Q;
95           num_t_col   = MAGMA_BASIS_NTCOL(num_threads, MAGMA_MAXTHREADS_1D);
96           shared_mem += sizeof(CeedScalar) * num_t_col * (num_comp * (1 * P + 1 * Q));
97           shared_mem += sizeof(CeedScalar) * (P * Q);
98           break;
99         case 2:
100           num_threads = max_P_Q;
101           num_t_col   = MAGMA_BASIS_NTCOL(num_threads, MAGMA_MAXTHREADS_2D);
102           shared_mem += P * Q * sizeof(CeedScalar);  // for sT
103           // for reforming rU we need P x P, and for the intermediate output we need P x Q
104           shared_mem += num_t_col * (P * max_P_Q * sizeof(CeedScalar));
105           break;
106         case 3:
107           num_threads = max_P_Q * max_P_Q;
108           num_t_col   = MAGMA_BASIS_NTCOL(num_threads, MAGMA_MAXTHREADS_3D);
109           shared_mem += sizeof(CeedScalar) * (P * Q);  // for sT
110           // rU needs P^2 x P, the intermediate output needs max(P^2 x Q, P x Q^2)
111           shared_mem += sizeof(CeedScalar) * num_t_col * (CeedIntMax(P * P * max_P_Q, P * Q * Q));
112           break;
113       }
114       CeedInt grid   = CeedDivUpInt(num_elem, num_t_col);
115       void   *args[] = {&impl->d_interp_1d, &d_u, &u_elem_stride, &u_comp_stride, &d_v, &v_elem_stride, &v_comp_stride, &num_elem};
116 
117       if (t_mode == CEED_TRANSPOSE) {
118         CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, impl->InterpTranspose, grid, num_threads, num_t_col, 1, shared_mem, args));
119       } else {
120         CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, impl->Interp, grid, num_threads, num_t_col, 1, shared_mem, args));
121       }
122     } break;
123     case CEED_EVAL_GRAD: {
124       // Define element sizes for dofs/quad
125       CeedInt elem_qpts_size = CeedIntPow(Q_1d, dim);
126       CeedInt elem_dofs_size = CeedIntPow(P_1d, dim);
127 
128       // In CEED_NOTRANSPOSE mode:
129       // d_u is (P^dim x nc), column-major layout (nc = num_comp)
130       // d_v is (Q^dim x nc x dim), column-major layout (nc = num_comp)
131       // In CEED_TRANSPOSE mode, the sizes of d_u and d_v are switched.
132 
133       // E-vector ordering -------------- Q-vector ordering
134       //                                  dim
135       //  component                        component
136       //    elem                              elem
137       //       node                            node
138 
139       // ---  Define strides for NOTRANSPOSE mode: ---
140       // Input (d_u) is E-vector, output (d_v) is Q-vector
141 
142       // Element strides
143       CeedInt u_elem_stride = elem_dofs_size;
144       CeedInt v_elem_stride = elem_qpts_size;
145       // Component strides
146       CeedInt u_comp_stride = num_elem * elem_dofs_size;
147       CeedInt v_comp_stride = num_elem * elem_qpts_size;
148       // Dimension strides
149       CeedInt u_dim_stride = 0;
150       CeedInt v_dim_stride = num_elem * elem_qpts_size * num_comp;
151       if (t_mode == CEED_TRANSPOSE) {
152         // Input (d_u) is Q-vector, output (d_v) is E-vector
153         // Element strides
154         v_elem_stride = elem_dofs_size;
155         u_elem_stride = elem_qpts_size;
156         // Component strides
157         v_comp_stride = num_elem * elem_dofs_size;
158         u_comp_stride = num_elem * elem_qpts_size;
159         // Dimension strides
160         v_dim_stride = 0;
161         u_dim_stride = num_elem * elem_qpts_size * num_comp;
162       }
163       CeedInt num_threads = 1;
164       CeedInt num_t_col   = 1;
165       CeedInt shared_mem  = 0;
166       CeedInt max_P_Q     = CeedIntMax(P, Q);
167 
168       switch (dim) {
169         case 1:
170           num_threads = max_P_Q;
171           num_t_col   = MAGMA_BASIS_NTCOL(num_threads, MAGMA_MAXTHREADS_1D);
172           shared_mem += sizeof(CeedScalar) * num_t_col * (num_comp * (1 * P + 1 * Q));
173           shared_mem += sizeof(CeedScalar) * (P * Q);
174           break;
175         case 2:
176           num_threads = max_P_Q;
177           num_t_col   = MAGMA_BASIS_NTCOL(num_threads, MAGMA_MAXTHREADS_2D);
178           shared_mem += sizeof(CeedScalar) * 2 * P * Q;  // for sTinterp and sTgrad
179           // for reforming rU we need P x P, and for the intermediate output we need P x Q
180           shared_mem += sizeof(CeedScalar) * num_t_col * (P * max_P_Q);
181           break;
182         case 3:
183           num_threads = max_P_Q * max_P_Q;
184           num_t_col   = MAGMA_BASIS_NTCOL(num_threads, MAGMA_MAXTHREADS_3D);
185           shared_mem += sizeof(CeedScalar) * 2 * P * Q;  // for sTinterp and sTgrad
186           // rU needs P^2 x P, the intermediate outputs need (P^2 x Q + P x Q^2)
187           shared_mem += sizeof(CeedScalar) * num_t_col * CeedIntMax(P * P * P, (P * P * Q) + (P * Q * Q));
188           break;
189       }
190       CeedInt grid   = CeedDivUpInt(num_elem, num_t_col);
191       void   *args[] = {&impl->d_interp_1d, &impl->d_grad_1d, &d_u,          &u_elem_stride, &u_comp_stride, &u_dim_stride, &d_v,
192                         &v_elem_stride,     &v_comp_stride,   &v_dim_stride, &num_elem};
193 
194       if (t_mode == CEED_TRANSPOSE) {
195         CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, impl->GradTranspose, grid, num_threads, num_t_col, 1, shared_mem, args));
196       } else {
197         CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, impl->Grad, grid, num_threads, num_t_col, 1, shared_mem, args));
198       }
199     } break;
200     case CEED_EVAL_WEIGHT: {
201       CeedCheck(t_mode != CEED_TRANSPOSE, ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT incompatible with CEED_TRANSPOSE");
202       CeedInt elem_dofs_size = CeedIntPow(Q, dim);
203       CeedInt num_threads    = 1;
204       CeedInt num_t_col      = 1;
205       CeedInt shared_mem     = 0;
206 
207       switch (dim) {
208         case 1:
209           num_threads = Q;
210           num_t_col   = MAGMA_BASIS_NTCOL(num_threads, MAGMA_MAXTHREADS_1D);
211           shared_mem += sizeof(CeedScalar) * Q;              // for d_q_weight_1d
212           shared_mem += sizeof(CeedScalar) * num_t_col * Q;  // for output
213           break;
214         case 2:
215           num_threads = Q;
216           num_t_col   = MAGMA_BASIS_NTCOL(num_threads, MAGMA_MAXTHREADS_2D);
217           shared_mem += sizeof(CeedScalar) * Q;  // for d_q_weight_1d
218           break;
219         case 3:
220           num_threads = Q * Q;
221           num_t_col   = MAGMA_BASIS_NTCOL(num_threads, MAGMA_MAXTHREADS_3D);
222           shared_mem += sizeof(CeedScalar) * Q;  // for d_q_weight_1d
223           break;
224       }
225       CeedInt grid   = CeedDivUpInt(num_elem, num_t_col);
226       void   *args[] = {&impl->d_q_weight_1d, &d_v, &elem_dofs_size, &num_elem};
227 
228       CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, impl->Weight, grid, num_threads, num_t_col, 1, shared_mem, args));
229     } break;
230     // LCOV_EXCL_START
231     case CEED_EVAL_DIV:
232       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_DIV not supported");
233     case CEED_EVAL_CURL:
234       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_CURL not supported");
235     case CEED_EVAL_NONE:
236       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_NONE does not make sense in this context");
237       // LCOV_EXCL_STOP
238   }
239 
240   // Must sync to ensure completeness
241   ceed_magma_queue_sync(data->queue);
242 
243   // Restore vectors
244   if (e_mode != CEED_EVAL_WEIGHT) {
245     CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
246   }
247   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
248   return CEED_ERROR_SUCCESS;
249 }
250 
251 //------------------------------------------------------------------------------
252 // Basis apply - non-tensor
253 //------------------------------------------------------------------------------
254 static int CeedBasisApplyNonTensor_Magma(CeedBasis basis, CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode e_mode, CeedVector u,
255                                          CeedVector v) {
256   Ceed                      ceed;
257   Ceed_Magma               *data;
258   CeedInt                   num_comp, num_nodes, num_qpts, P, Q, N;
259   const CeedScalar         *d_u;
260   CeedScalar               *d_v;
261   CeedBasisNonTensor_Magma *impl;
262 
263   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
264   CeedCallBackend(CeedGetData(ceed, &data));
265   CeedCallBackend(CeedBasisGetData(basis, &impl));
266   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
267   CeedCallBackend(CeedBasisGetNumNodes(basis, &num_nodes));
268   CeedCallBackend(CeedBasisGetNumQuadraturePoints(basis, &num_qpts));
269   P = num_nodes;
270   Q = num_qpts;
271   N = num_elem * num_comp;
272 
273   // Read vectors
274   if (u != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
275   else CeedCheck(e_mode == CEED_EVAL_WEIGHT, ceed, CEED_ERROR_BACKEND, "An input vector is required for this CeedEvalMode");
276   CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
277 
278   // Compile kernels for N as needed
279   CeedInt iN = 0;
280   if (P <= MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P && Q <= MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q && (e_mode != CEED_EVAL_WEIGHT || !impl->Weight)) {
281     CeedInt n_array[MAGMA_NONTENSOR_KERNEL_INSTANCES] = {MAGMA_NONTENSOR_KERNEL_N_VALUES};
282     CeedInt diff                                      = abs(n_array[iN] - N), idiff;
283 
284     for (CeedInt in = iN + 1; in < MAGMA_NONTENSOR_KERNEL_INSTANCES; in++) {
285       idiff = abs(n_array[in] - N);
286       if (idiff < diff) {
287         iN   = in;
288         diff = idiff;
289       }
290     }
291 
292     if (!impl->NB_interp[iN]) {
293       CeedFESpace fe_space;
294       CeedInt     q_comp_interp, q_comp_deriv;
295       Ceed        ceed_delegate;
296       char       *basis_kernel_path, *weight_kernel_path, *basis_kernel_source;
297       magma_int_t arch = magma_getdevice_arch();
298 
299       // Tuning parameters for NB
300       CeedCallBackend(CeedBasisGetFESpace(basis, &fe_space));
301       CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
302       switch (fe_space) {
303         case CEED_FE_SPACE_H1:
304           CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_GRAD, &q_comp_deriv));
305           break;
306         case CEED_FE_SPACE_HDIV:
307           CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_DIV, &q_comp_deriv));
308           break;
309         case CEED_FE_SPACE_HCURL:
310           CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_CURL, &q_comp_deriv));
311           break;
312       }
313       impl->NB_interp[iN]   = nontensor_rtc_get_nb(arch, 'n', q_comp_interp, P, Q, n_array[iN]);
314       impl->NB_interp_t[iN] = nontensor_rtc_get_nb(arch, 't', q_comp_interp, P, Q, n_array[iN]);
315       impl->NB_deriv[iN]    = nontensor_rtc_get_nb(arch, 'n', q_comp_deriv, P, Q, n_array[iN]);
316       impl->NB_deriv_t[iN]  = nontensor_rtc_get_nb(arch, 't', q_comp_deriv, P, Q, n_array[iN]);
317 
318       // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
319       CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
320 
321       // Compile kernels
322       CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-interp-deriv-nontensor.h", &basis_kernel_path));
323       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
324       CeedCallBackend(CeedLoadSourceToBuffer(ceed, basis_kernel_path, &basis_kernel_source));
325       if (!impl->Weight) {
326         CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
327         CeedCallBackend(CeedLoadSourceToInitializedBuffer(ceed, weight_kernel_path, &basis_kernel_source));
328       }
329       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
330       CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[iN], 8, "BASIS_Q_COMP_INTERP", q_comp_interp,
331                                        "BASIS_Q_COMP_DERIV", q_comp_deriv, "BASIS_P", P, "BASIS_Q", Q, "BASIS_NB_INTERP_N", impl->NB_interp[iN],
332                                        "BASIS_NB_INTERP_T", impl->NB_interp_t[iN], "BASIS_NB_DERIV_N", impl->NB_deriv[iN], "BASIS_NB_DERIV_T",
333                                        impl->NB_deriv_t[iN]));
334       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[iN], "magma_interp_nontensor_n", &impl->Interp[iN]));
335       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[iN], "magma_interp_nontensor_t", &impl->InterpTranspose[iN]));
336       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[iN], "magma_deriv_nontensor_n", &impl->Deriv[iN]));
337       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[iN], "magma_deriv_nontensor_t", &impl->DerivTranspose[iN]));
338       if (!impl->Weight) {
339         CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[iN], "magma_weight_nontensor", &impl->Weight));
340         CeedCallBackend(CeedFree(&weight_kernel_path));
341       }
342       CeedCallBackend(CeedFree(&basis_kernel_path));
343       CeedCallBackend(CeedFree(&basis_kernel_source));
344     }
345   }
346 
347   // Apply basis operation
348   if (e_mode != CEED_EVAL_WEIGHT) {
349     const CeedScalar *d_b = NULL;
350     CeedInt           q_comp, NB, M, K;
351     CeedMagmaFunction Kernel;
352 
353     switch (e_mode) {
354       case CEED_EVAL_INTERP:
355         d_b = impl->d_interp;
356         break;
357       case CEED_EVAL_GRAD:
358         d_b = impl->d_grad;
359         break;
360       case CEED_EVAL_DIV:
361         d_b = impl->d_div;
362         break;
363       case CEED_EVAL_CURL:
364         d_b = impl->d_curl;
365         break;
366       // LCOV_EXCL_START
367       case CEED_EVAL_WEIGHT:
368         return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT does not make sense in this context");
369       case CEED_EVAL_NONE:
370         return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_NONE does not make sense in this context");
371         // LCOV_EXCL_STOP
372     }
373     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, e_mode, &q_comp));
374     M = (t_mode == CEED_TRANSPOSE) ? P : Q, K = (t_mode == CEED_TRANSPOSE) ? Q : P;
375 
376     if (P <= MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P && Q <= MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
377       if (e_mode == CEED_EVAL_INTERP) {
378         if (t_mode == CEED_TRANSPOSE) {
379           Kernel = impl->InterpTranspose[iN];
380           NB     = impl->NB_interp_t[iN];
381         } else {
382           Kernel = impl->Interp[iN];
383           NB     = impl->NB_interp[iN];
384         }
385       } else {
386         if (t_mode == CEED_TRANSPOSE) {
387           Kernel = impl->DerivTranspose[iN];
388           NB     = impl->NB_deriv_t[iN];
389         } else {
390           Kernel = impl->Deriv[iN];
391           NB     = impl->NB_deriv[iN];
392         }
393       }
394       CeedInt num_t_col    = MAGMA_BASIS_NTCOL(M, MAGMA_MAXTHREADS_1D);
395       CeedInt grid         = CeedDivUpInt(N, num_t_col * NB);
396       CeedInt shared_mem_A = P * Q * sizeof(CeedScalar);
397       CeedInt shared_mem_B = num_t_col * K * NB * sizeof(CeedScalar);
398       CeedInt shared_mem   = (t_mode != CEED_TRANSPOSE && q_comp > 1) ? (shared_mem_A + shared_mem_B) : CeedIntMax(shared_mem_A, shared_mem_B);
399       void   *args[]       = {&N, &d_b, &d_u, &d_v};
400 
401       CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, Kernel, grid, M, num_t_col, 1, shared_mem, args));
402     } else {
403       for (CeedInt d = 0; d < q_comp; d++) {
404         if (t_mode == CEED_TRANSPOSE) {
405           const CeedScalar beta = (d > 0) ? 1.0 : 0.0;
406           magma_gemm_nontensor(MagmaNoTrans, MagmaNoTrans, P, N, Q, 1.0, d_b + d * P * Q, P, d_u + d * N * Q, Q, beta, d_v, P, data->queue);
407         } else {
408           magma_gemm_nontensor(MagmaTrans, MagmaNoTrans, Q, N, P, 1.0, d_b + d * P * Q, P, d_u, P, 0.0, d_v + d * N * Q, Q, data->queue);
409         }
410       }
411     }
412   } else {
413     CeedCheck(t_mode != CEED_TRANSPOSE, ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT incompatible with CEED_TRANSPOSE");
414     CeedInt num_t_col  = MAGMA_BASIS_NTCOL(Q, MAGMA_MAXTHREADS_1D);
415     CeedInt grid       = CeedDivUpInt(num_elem, num_t_col);
416     CeedInt shared_mem = Q * sizeof(CeedScalar) + num_t_col * Q * sizeof(CeedScalar);
417     void   *args[]     = {&num_elem, &impl->d_q_weight, &d_v};
418 
419     CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, impl->Weight, grid, Q, num_t_col, 1, shared_mem, args));
420   }
421 
422   // Must sync to ensure completeness
423   ceed_magma_queue_sync(data->queue);
424 
425   // Restore vectors
426   if (e_mode != CEED_EVAL_WEIGHT) {
427     CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
428   }
429   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
430   return CEED_ERROR_SUCCESS;
431 }
432 
433 //------------------------------------------------------------------------------
434 // Destroy tensor basis
435 //------------------------------------------------------------------------------
436 static int CeedBasisDestroy_Magma(CeedBasis basis) {
437   Ceed             ceed;
438   CeedBasis_Magma *impl;
439 
440   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
441   CeedCallBackend(CeedBasisGetData(basis, &impl));
442 #ifdef CEED_MAGMA_USE_HIP
443   CeedCallHip(ceed, hipModuleUnload(impl->module));
444 #else
445   CeedCallCuda(ceed, cuModuleUnload(impl->module));
446 #endif
447   CeedCallBackend(magma_free(impl->d_interp_1d));
448   CeedCallBackend(magma_free(impl->d_grad_1d));
449   CeedCallBackend(magma_free(impl->d_q_weight_1d));
450   CeedCallBackend(CeedFree(&impl));
451   return CEED_ERROR_SUCCESS;
452 }
453 
454 //------------------------------------------------------------------------------
455 // Destroy non-tensor basis
456 //------------------------------------------------------------------------------
457 static int CeedBasisDestroyNonTensor_Magma(CeedBasis basis) {
458   Ceed                      ceed;
459   CeedBasisNonTensor_Magma *impl;
460 
461   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
462   CeedCallBackend(CeedBasisGetData(basis, &impl));
463   for (CeedInt in = 0; in < MAGMA_NONTENSOR_KERNEL_INSTANCES; in++) {
464     if (impl->module[in]) {
465 #ifdef CEED_MAGMA_USE_HIP
466       CeedCallHip(ceed, hipModuleUnload(impl->module[in]));
467 #else
468       CeedCallCuda(ceed, cuModuleUnload(impl->module[in]));
469 #endif
470     }
471   }
472   CeedCallBackend(magma_free(impl->d_interp));
473   CeedCallBackend(magma_free(impl->d_grad));
474   CeedCallBackend(magma_free(impl->d_div));
475   CeedCallBackend(magma_free(impl->d_curl));
476   CeedCallBackend(magma_free(impl->d_q_weight));
477   CeedCallBackend(CeedFree(&impl));
478   return CEED_ERROR_SUCCESS;
479 }
480 
481 //------------------------------------------------------------------------------
482 // Create tensor
483 //------------------------------------------------------------------------------
484 int CeedBasisCreateTensorH1_Magma(CeedInt dim, CeedInt P_1d, CeedInt Q_1d, const CeedScalar *interp_1d, const CeedScalar *grad_1d,
485                                   const CeedScalar *q_ref_1d, const CeedScalar *q_weight_1d, CeedBasis basis) {
486   Ceed             ceed, ceed_delegate;
487   Ceed_Magma      *data;
488   char            *interp_kernel_path, *grad_kernel_path, *weight_kernel_path, *basis_kernel_source;
489   CeedInt          num_comp;
490   CeedBasis_Magma *impl;
491 
492   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
493   CeedCallBackend(CeedGetData(ceed, &data));
494   CeedCallBackend(CeedCalloc(1, &impl));
495 
496   // Copy basis data to GPU
497   CeedCallBackend(magma_malloc((void **)&impl->d_q_weight_1d, Q_1d * sizeof(q_weight_1d[0])));
498   magma_setvector(Q_1d, sizeof(q_weight_1d[0]), q_weight_1d, 1, impl->d_q_weight_1d, 1, data->queue);
499   CeedCallBackend(magma_malloc((void **)&impl->d_interp_1d, Q_1d * P_1d * sizeof(interp_1d[0])));
500   magma_setvector(Q_1d * P_1d, sizeof(interp_1d[0]), interp_1d, 1, impl->d_interp_1d, 1, data->queue);
501   CeedCallBackend(magma_malloc((void **)&impl->d_grad_1d, Q_1d * P_1d * sizeof(grad_1d[0])));
502   magma_setvector(Q_1d * P_1d, sizeof(grad_1d[0]), grad_1d, 1, impl->d_grad_1d, 1, data->queue);
503 
504   // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
505   CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
506 
507   // Compile kernels
508   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
509   {
510     char   *interp_kernel_name_base = "ceed/jit-source/magma/magma-basis-interp";
511     CeedInt interp_kernel_name_len  = strlen(interp_kernel_name_base) + 6;
512     char    interp_kernel_name[interp_kernel_name_len];
513 
514     snprintf(interp_kernel_name, interp_kernel_name_len, "%s-%" CeedInt_FMT "d.h", interp_kernel_name_base, dim);
515     CeedCallBackend(CeedGetJitAbsolutePath(ceed, interp_kernel_name, &interp_kernel_path));
516   }
517   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
518   CeedCallBackend(CeedLoadSourceToBuffer(ceed, interp_kernel_path, &basis_kernel_source));
519   {
520     char   *grad_kernel_name_base = "ceed/jit-source/magma/magma-basis-grad";
521     CeedInt grad_kernel_name_len  = strlen(grad_kernel_name_base) + 6;
522     char    grad_kernel_name[grad_kernel_name_len];
523 
524     snprintf(grad_kernel_name, grad_kernel_name_len, "%s-%" CeedInt_FMT "d.h", grad_kernel_name_base, dim);
525     CeedCallBackend(CeedGetJitAbsolutePath(ceed, grad_kernel_name, &grad_kernel_path));
526   }
527   CeedCallBackend(CeedLoadSourceToInitializedBuffer(ceed, grad_kernel_path, &basis_kernel_source));
528   {
529     char   *weight_kernel_name_base = "ceed/jit-source/magma/magma-basis-weight";
530     CeedInt weight_kernel_name_len  = strlen(weight_kernel_name_base) + 6;
531     char    weight_kernel_name[weight_kernel_name_len];
532 
533     snprintf(weight_kernel_name, weight_kernel_name_len, "%s-%" CeedInt_FMT "d.h", weight_kernel_name_base, dim);
534     CeedCallBackend(CeedGetJitAbsolutePath(ceed, weight_kernel_name, &weight_kernel_path));
535   }
536   CeedCallBackend(CeedLoadSourceToInitializedBuffer(ceed, weight_kernel_path, &basis_kernel_source));
537   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
538   CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module, 5, "BASIS_DIM", dim, "BASIS_NUM_COMP", num_comp, "BASIS_P",
539                                    P_1d, "BASIS_Q", Q_1d, "BASIS_MAX_P_Q", CeedIntMax(P_1d, Q_1d)));
540   switch (dim) {
541     case 1:
542       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpn_1d_kernel", &impl->Interp));
543       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpt_1d_kernel", &impl->InterpTranspose));
544       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradn_1d_kernel", &impl->Grad));
545       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradt_1d_kernel", &impl->GradTranspose));
546       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_weight_1d_kernel", &impl->Weight));
547       break;
548     case 2:
549       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpn_2d_kernel", &impl->Interp));
550       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpt_2d_kernel", &impl->InterpTranspose));
551       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradn_2d_kernel", &impl->Grad));
552       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradt_2d_kernel", &impl->GradTranspose));
553       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_weight_2d_kernel", &impl->Weight));
554       break;
555     case 3:
556       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpn_3d_kernel", &impl->Interp));
557       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpt_3d_kernel", &impl->InterpTranspose));
558       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradn_3d_kernel", &impl->Grad));
559       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradt_3d_kernel", &impl->GradTranspose));
560       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_weight_3d_kernel", &impl->Weight));
561       break;
562   }
563   CeedCallBackend(CeedFree(&interp_kernel_path));
564   CeedCallBackend(CeedFree(&grad_kernel_path));
565   CeedCallBackend(CeedFree(&weight_kernel_path));
566   CeedCallBackend(CeedFree(&basis_kernel_source));
567 
568   CeedCallBackend(CeedBasisSetData(basis, impl));
569 
570   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApply_Magma));
571   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroy_Magma));
572   return CEED_ERROR_SUCCESS;
573 }
574 
575 //------------------------------------------------------------------------------
576 // Create non-tensor H^1
577 //------------------------------------------------------------------------------
578 int CeedBasisCreateH1_Magma(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp, const CeedScalar *grad,
579                             const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
580   Ceed                      ceed;
581   Ceed_Magma               *data;
582   CeedBasisNonTensor_Magma *impl;
583 
584   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
585   CeedCallBackend(CeedGetData(ceed, &data));
586   CeedCallBackend(CeedCalloc(1, &impl));
587 
588   // Copy basis data to GPU
589   CeedCallBackend(magma_malloc((void **)&impl->d_q_weight, num_qpts * sizeof(q_weight[0])));
590   magma_setvector(num_qpts, sizeof(q_weight[0]), q_weight, 1, impl->d_q_weight, 1, data->queue);
591   if (interp) {
592     CeedInt q_comp_interp;
593 
594     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
595     CeedCallBackend(magma_malloc((void **)&impl->d_interp, num_qpts * num_nodes * q_comp_interp * sizeof(interp[0])));
596     magma_setvector(num_qpts * num_nodes * q_comp_interp, sizeof(interp[0]), interp, 1, impl->d_interp, 1, data->queue);
597   }
598   if (grad) {
599     CeedInt q_comp_grad;
600 
601     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_GRAD, &q_comp_grad));
602     CeedCallBackend(magma_malloc((void **)&impl->d_grad, num_qpts * num_nodes * q_comp_grad * sizeof(grad[0])));
603     magma_setvector(num_qpts * num_nodes * q_comp_grad, sizeof(grad[0]), grad, 1, impl->d_grad, 1, data->queue);
604   }
605 
606   // Compile the weight kernel if it won't be compiled later on
607   if (num_nodes > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P || num_qpts > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
608     Ceed  ceed_delegate;
609     char *weight_kernel_path, *basis_kernel_source;
610 
611     // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
612     CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
613 
614     // Compile weight kernel (the remainder of kernel compilation happens at first call to CeedBasisApply)
615     CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
616     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
617     CeedCallBackend(CeedLoadSourceToBuffer(ceed, weight_kernel_path, &basis_kernel_source));
618     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
619     CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[0], 1, "BASIS_Q", num_qpts));
620     CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[0], "magma_weight_nontensor", &impl->Weight));
621     CeedCallBackend(CeedFree(&weight_kernel_path));
622     CeedCallBackend(CeedFree(&basis_kernel_source));
623   }
624 
625   CeedCallBackend(CeedBasisSetData(basis, impl));
626 
627   // Register backend functions
628   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Magma));
629   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyNonTensor_Magma));
630   return CEED_ERROR_SUCCESS;
631 }
632 
633 //------------------------------------------------------------------------------
634 // Create non-tensor H(div)
635 //------------------------------------------------------------------------------
636 int CeedBasisCreateHdiv_Magma(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp,
637                               const CeedScalar *div, const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
638   Ceed                      ceed;
639   Ceed_Magma               *data;
640   CeedBasisNonTensor_Magma *impl;
641 
642   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
643   CeedCallBackend(CeedGetData(ceed, &data));
644   CeedCallBackend(CeedCalloc(1, &impl));
645 
646   // Copy basis data to GPU
647   CeedCallBackend(magma_malloc((void **)&impl->d_q_weight, num_qpts * sizeof(q_weight[0])));
648   magma_setvector(num_qpts, sizeof(q_weight[0]), q_weight, 1, impl->d_q_weight, 1, data->queue);
649   if (interp) {
650     CeedInt q_comp_interp;
651 
652     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
653     CeedCallBackend(magma_malloc((void **)&impl->d_interp, num_qpts * num_nodes * q_comp_interp * sizeof(interp[0])));
654     magma_setvector(num_qpts * num_nodes * q_comp_interp, sizeof(interp[0]), interp, 1, impl->d_interp, 1, data->queue);
655   }
656   if (div) {
657     CeedInt q_comp_div;
658 
659     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_DIV, &q_comp_div));
660     CeedCallBackend(magma_malloc((void **)&impl->d_div, num_qpts * num_nodes * q_comp_div * sizeof(div[0])));
661     magma_setvector(num_qpts * num_nodes * q_comp_div, sizeof(div[0]), div, 1, impl->d_div, 1, data->queue);
662   }
663 
664   // Compile the weight kernel if it won't be compiled later on
665   if (num_nodes > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P || num_qpts > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
666     Ceed  ceed_delegate;
667     char *weight_kernel_path, *basis_kernel_source;
668 
669     // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
670     CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
671 
672     // Compile weight kernel (the remainder of kernel compilation happens at first call to CeedBasisApply)
673     CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
674     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
675     CeedCallBackend(CeedLoadSourceToBuffer(ceed, weight_kernel_path, &basis_kernel_source));
676     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
677     CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[0], 1, "BASIS_Q", num_qpts));
678     CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[0], "magma_weight_nontensor", &impl->Weight));
679     CeedCallBackend(CeedFree(&weight_kernel_path));
680     CeedCallBackend(CeedFree(&basis_kernel_source));
681   }
682 
683   CeedCallBackend(CeedBasisSetData(basis, impl));
684 
685   // Register backend functions
686   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Magma));
687   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyNonTensor_Magma));
688   return CEED_ERROR_SUCCESS;
689 }
690 
691 //------------------------------------------------------------------------------
692 // Create non-tensor H(curl)
693 //------------------------------------------------------------------------------
694 int CeedBasisCreateHcurl_Magma(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp,
695                                const CeedScalar *curl, const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
696   Ceed                      ceed;
697   Ceed_Magma               *data;
698   CeedBasisNonTensor_Magma *impl;
699 
700   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
701   CeedCallBackend(CeedGetData(ceed, &data));
702   CeedCallBackend(CeedCalloc(1, &impl));
703 
704   // Copy basis data to GPU
705   CeedCallBackend(magma_malloc((void **)&impl->d_q_weight, num_qpts * sizeof(q_weight[0])));
706   magma_setvector(num_qpts, sizeof(q_weight[0]), q_weight, 1, impl->d_q_weight, 1, data->queue);
707   if (interp) {
708     CeedInt q_comp_interp;
709 
710     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
711     CeedCallBackend(magma_malloc((void **)&impl->d_interp, num_qpts * num_nodes * q_comp_interp * sizeof(interp[0])));
712     magma_setvector(num_qpts * num_nodes * q_comp_interp, sizeof(interp[0]), interp, 1, impl->d_interp, 1, data->queue);
713   }
714   if (curl) {
715     CeedInt q_comp_curl;
716 
717     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_CURL, &q_comp_curl));
718     CeedCallBackend(magma_malloc((void **)&impl->d_curl, num_qpts * num_nodes * q_comp_curl * sizeof(curl[0])));
719     magma_setvector(num_qpts * num_nodes * q_comp_curl, sizeof(curl[0]), curl, 1, impl->d_curl, 1, data->queue);
720   }
721 
722   // Compile the weight kernel if it won't be compiled later on
723   if (num_nodes > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P || num_qpts > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
724     Ceed  ceed_delegate;
725     char *weight_kernel_path, *basis_kernel_source;
726 
727     // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
728     CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
729 
730     // Compile weight kernel (the remainder of kernel compilation happens at first call to CeedBasisApply)
731     CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
732     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
733     CeedCallBackend(CeedLoadSourceToBuffer(ceed, weight_kernel_path, &basis_kernel_source));
734     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
735     CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[0], 1, "BASIS_Q", num_qpts));
736     CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[0], "magma_weight_nontensor", &impl->Weight));
737     CeedCallBackend(CeedFree(&weight_kernel_path));
738     CeedCallBackend(CeedFree(&basis_kernel_source));
739   }
740 
741   CeedCallBackend(CeedBasisSetData(basis, impl));
742 
743   // Register backend functions
744   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Magma));
745   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyNonTensor_Magma));
746   return CEED_ERROR_SUCCESS;
747 }
748 
749 //------------------------------------------------------------------------------
750