xref: /libCEED/rust/libceed-sys/c-src/backends/magma/ceed-magma-basis.c (revision 5aed82e4fa97acf4ba24a7f10a35f5303a6798e0)
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 <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     case CEED_EVAL_CURL:
233       return CeedError(ceed, CEED_ERROR_BACKEND, "%s not supported", CeedEvalModes[e_mode]);
234     case CEED_EVAL_NONE:
235       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_NONE does not make sense in this context");
236       // LCOV_EXCL_STOP
237   }
238 
239   // Must sync to ensure completeness
240   ceed_magma_queue_sync(data->queue);
241 
242   // Restore vectors
243   if (e_mode != CEED_EVAL_WEIGHT) {
244     CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
245   }
246   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
247   return CEED_ERROR_SUCCESS;
248 }
249 
250 //------------------------------------------------------------------------------
251 // Basis apply - non-tensor
252 //------------------------------------------------------------------------------
253 static int CeedBasisApplyNonTensor_Magma(CeedBasis basis, CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode e_mode, CeedVector u,
254                                          CeedVector v) {
255   Ceed                      ceed;
256   Ceed_Magma               *data;
257   CeedInt                   num_comp, num_nodes, num_qpts, P, Q, N;
258   const CeedScalar         *d_u;
259   CeedScalar               *d_v;
260   CeedBasisNonTensor_Magma *impl;
261 
262   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
263   CeedCallBackend(CeedGetData(ceed, &data));
264   CeedCallBackend(CeedBasisGetData(basis, &impl));
265   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
266   CeedCallBackend(CeedBasisGetNumNodes(basis, &num_nodes));
267   CeedCallBackend(CeedBasisGetNumQuadraturePoints(basis, &num_qpts));
268   P = num_nodes;
269   Q = num_qpts;
270   N = num_elem * num_comp;
271 
272   // Read vectors
273   if (u != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
274   else CeedCheck(e_mode == CEED_EVAL_WEIGHT, ceed, CEED_ERROR_BACKEND, "An input vector is required for this CeedEvalMode");
275   CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
276 
277   // Compile kernels for N as needed
278   CeedInt iN = 0;
279   if (P <= MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P && Q <= MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q && (e_mode != CEED_EVAL_WEIGHT || !impl->Weight)) {
280     CeedInt n_array[MAGMA_NONTENSOR_KERNEL_INSTANCES] = {MAGMA_NONTENSOR_KERNEL_N_VALUES};
281     CeedInt diff                                      = abs(n_array[iN] - N), idiff;
282 
283     for (CeedInt in = iN + 1; in < MAGMA_NONTENSOR_KERNEL_INSTANCES; in++) {
284       idiff = abs(n_array[in] - N);
285       if (idiff < diff) {
286         iN   = in;
287         diff = idiff;
288       }
289     }
290 
291     if (!impl->NB_interp[iN]) {
292       CeedFESpace fe_space;
293       CeedInt     q_comp_interp, q_comp_deriv;
294       Ceed        ceed_delegate;
295       char       *basis_kernel_source;
296       const char *basis_kernel_path, *weight_kernel_path;
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       case CEED_EVAL_NONE:
369         return CeedError(ceed, CEED_ERROR_BACKEND, "%s does not make sense in this context", CeedEvalModes[e_mode]);
370         // LCOV_EXCL_STOP
371     }
372     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, e_mode, &q_comp));
373     M = (t_mode == CEED_TRANSPOSE) ? P : Q, K = (t_mode == CEED_TRANSPOSE) ? Q : P;
374 
375     if (P <= MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P && Q <= MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
376       if (e_mode == CEED_EVAL_INTERP) {
377         if (t_mode == CEED_TRANSPOSE) {
378           Kernel = impl->InterpTranspose[iN];
379           NB     = impl->NB_interp_t[iN];
380         } else {
381           Kernel = impl->Interp[iN];
382           NB     = impl->NB_interp[iN];
383         }
384       } else {
385         if (t_mode == CEED_TRANSPOSE) {
386           Kernel = impl->DerivTranspose[iN];
387           NB     = impl->NB_deriv_t[iN];
388         } else {
389           Kernel = impl->Deriv[iN];
390           NB     = impl->NB_deriv[iN];
391         }
392       }
393       CeedInt num_t_col    = MAGMA_BASIS_NTCOL(M, MAGMA_MAXTHREADS_1D);
394       CeedInt grid         = CeedDivUpInt(N, num_t_col * NB);
395       CeedInt shared_mem_A = P * Q * sizeof(CeedScalar);
396       CeedInt shared_mem_B = num_t_col * K * NB * sizeof(CeedScalar);
397       CeedInt shared_mem   = (t_mode != CEED_TRANSPOSE && q_comp > 1) ? (shared_mem_A + shared_mem_B) : CeedIntMax(shared_mem_A, shared_mem_B);
398       void   *args[]       = {&N, &d_b, &d_u, &d_v};
399 
400       CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, Kernel, grid, M, num_t_col, 1, shared_mem, args));
401     } else {
402       for (CeedInt d = 0; d < q_comp; d++) {
403         if (t_mode == CEED_TRANSPOSE) {
404           const CeedScalar beta = (d > 0) ? 1.0 : 0.0;
405           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);
406         } else {
407           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);
408         }
409       }
410     }
411   } else {
412     CeedCheck(t_mode != CEED_TRANSPOSE, ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT incompatible with CEED_TRANSPOSE");
413     CeedInt num_t_col  = MAGMA_BASIS_NTCOL(Q, MAGMA_MAXTHREADS_1D);
414     CeedInt grid       = CeedDivUpInt(num_elem, num_t_col);
415     CeedInt shared_mem = Q * sizeof(CeedScalar) + num_t_col * Q * sizeof(CeedScalar);
416     void   *args[]     = {&num_elem, &impl->d_q_weight, &d_v};
417 
418     CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, impl->Weight, grid, Q, num_t_col, 1, shared_mem, args));
419   }
420 
421   // Must sync to ensure completeness
422   ceed_magma_queue_sync(data->queue);
423 
424   // Restore vectors
425   if (e_mode != CEED_EVAL_WEIGHT) {
426     CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
427   }
428   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
429   return CEED_ERROR_SUCCESS;
430 }
431 
432 //------------------------------------------------------------------------------
433 // Destroy tensor basis
434 //------------------------------------------------------------------------------
435 static int CeedBasisDestroy_Magma(CeedBasis basis) {
436   Ceed             ceed;
437   CeedBasis_Magma *impl;
438 
439   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
440   CeedCallBackend(CeedBasisGetData(basis, &impl));
441 #ifdef CEED_MAGMA_USE_HIP
442   CeedCallHip(ceed, hipModuleUnload(impl->module));
443 #else
444   CeedCallCuda(ceed, cuModuleUnload(impl->module));
445 #endif
446   CeedCallBackend(magma_free(impl->d_interp_1d));
447   CeedCallBackend(magma_free(impl->d_grad_1d));
448   CeedCallBackend(magma_free(impl->d_q_weight_1d));
449   CeedCallBackend(CeedFree(&impl));
450   return CEED_ERROR_SUCCESS;
451 }
452 
453 //------------------------------------------------------------------------------
454 // Destroy non-tensor basis
455 //------------------------------------------------------------------------------
456 static int CeedBasisDestroyNonTensor_Magma(CeedBasis basis) {
457   Ceed                      ceed;
458   CeedBasisNonTensor_Magma *impl;
459 
460   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
461   CeedCallBackend(CeedBasisGetData(basis, &impl));
462   for (CeedInt in = 0; in < MAGMA_NONTENSOR_KERNEL_INSTANCES; in++) {
463     if (impl->module[in]) {
464 #ifdef CEED_MAGMA_USE_HIP
465       CeedCallHip(ceed, hipModuleUnload(impl->module[in]));
466 #else
467       CeedCallCuda(ceed, cuModuleUnload(impl->module[in]));
468 #endif
469     }
470   }
471   CeedCallBackend(magma_free(impl->d_interp));
472   CeedCallBackend(magma_free(impl->d_grad));
473   CeedCallBackend(magma_free(impl->d_div));
474   CeedCallBackend(magma_free(impl->d_curl));
475   CeedCallBackend(magma_free(impl->d_q_weight));
476   CeedCallBackend(CeedFree(&impl));
477   return CEED_ERROR_SUCCESS;
478 }
479 
480 //------------------------------------------------------------------------------
481 // Create tensor
482 //------------------------------------------------------------------------------
483 int CeedBasisCreateTensorH1_Magma(CeedInt dim, CeedInt P_1d, CeedInt Q_1d, const CeedScalar *interp_1d, const CeedScalar *grad_1d,
484                                   const CeedScalar *q_ref_1d, const CeedScalar *q_weight_1d, CeedBasis basis) {
485   Ceed             ceed, ceed_delegate;
486   Ceed_Magma      *data;
487   char            *basis_kernel_source;
488   const char      *interp_kernel_path, *grad_kernel_path, *weight_kernel_path;
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       *basis_kernel_source;
610     const char *weight_kernel_path;
611 
612     // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
613     CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
614 
615     // Compile weight kernel (the remainder of kernel compilation happens at first call to CeedBasisApply)
616     CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
617     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
618     CeedCallBackend(CeedLoadSourceToBuffer(ceed, weight_kernel_path, &basis_kernel_source));
619     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
620     CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[0], 1, "BASIS_Q", num_qpts));
621     CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[0], "magma_weight_nontensor", &impl->Weight));
622     CeedCallBackend(CeedFree(&weight_kernel_path));
623     CeedCallBackend(CeedFree(&basis_kernel_source));
624   }
625 
626   CeedCallBackend(CeedBasisSetData(basis, impl));
627 
628   // Register backend functions
629   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Magma));
630   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyNonTensor_Magma));
631   return CEED_ERROR_SUCCESS;
632 }
633 
634 //------------------------------------------------------------------------------
635 // Create non-tensor H(div)
636 //------------------------------------------------------------------------------
637 int CeedBasisCreateHdiv_Magma(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp,
638                               const CeedScalar *div, const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
639   Ceed                      ceed;
640   Ceed_Magma               *data;
641   CeedBasisNonTensor_Magma *impl;
642 
643   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
644   CeedCallBackend(CeedGetData(ceed, &data));
645   CeedCallBackend(CeedCalloc(1, &impl));
646 
647   // Copy basis data to GPU
648   CeedCallBackend(magma_malloc((void **)&impl->d_q_weight, num_qpts * sizeof(q_weight[0])));
649   magma_setvector(num_qpts, sizeof(q_weight[0]), q_weight, 1, impl->d_q_weight, 1, data->queue);
650   if (interp) {
651     CeedInt q_comp_interp;
652 
653     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
654     CeedCallBackend(magma_malloc((void **)&impl->d_interp, num_qpts * num_nodes * q_comp_interp * sizeof(interp[0])));
655     magma_setvector(num_qpts * num_nodes * q_comp_interp, sizeof(interp[0]), interp, 1, impl->d_interp, 1, data->queue);
656   }
657   if (div) {
658     CeedInt q_comp_div;
659 
660     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_DIV, &q_comp_div));
661     CeedCallBackend(magma_malloc((void **)&impl->d_div, num_qpts * num_nodes * q_comp_div * sizeof(div[0])));
662     magma_setvector(num_qpts * num_nodes * q_comp_div, sizeof(div[0]), div, 1, impl->d_div, 1, data->queue);
663   }
664 
665   // Compile the weight kernel if it won't be compiled later on
666   if (num_nodes > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P || num_qpts > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
667     Ceed        ceed_delegate;
668     char       *basis_kernel_source;
669     const char *weight_kernel_path;
670 
671     // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
672     CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
673 
674     // Compile weight kernel (the remainder of kernel compilation happens at first call to CeedBasisApply)
675     CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
676     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
677     CeedCallBackend(CeedLoadSourceToBuffer(ceed, weight_kernel_path, &basis_kernel_source));
678     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
679     CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[0], 1, "BASIS_Q", num_qpts));
680     CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[0], "magma_weight_nontensor", &impl->Weight));
681     CeedCallBackend(CeedFree(&weight_kernel_path));
682     CeedCallBackend(CeedFree(&basis_kernel_source));
683   }
684 
685   CeedCallBackend(CeedBasisSetData(basis, impl));
686 
687   // Register backend functions
688   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Magma));
689   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyNonTensor_Magma));
690   return CEED_ERROR_SUCCESS;
691 }
692 
693 //------------------------------------------------------------------------------
694 // Create non-tensor H(curl)
695 //------------------------------------------------------------------------------
696 int CeedBasisCreateHcurl_Magma(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp,
697                                const CeedScalar *curl, const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
698   Ceed                      ceed;
699   Ceed_Magma               *data;
700   CeedBasisNonTensor_Magma *impl;
701 
702   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
703   CeedCallBackend(CeedGetData(ceed, &data));
704   CeedCallBackend(CeedCalloc(1, &impl));
705 
706   // Copy basis data to GPU
707   CeedCallBackend(magma_malloc((void **)&impl->d_q_weight, num_qpts * sizeof(q_weight[0])));
708   magma_setvector(num_qpts, sizeof(q_weight[0]), q_weight, 1, impl->d_q_weight, 1, data->queue);
709   if (interp) {
710     CeedInt q_comp_interp;
711 
712     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
713     CeedCallBackend(magma_malloc((void **)&impl->d_interp, num_qpts * num_nodes * q_comp_interp * sizeof(interp[0])));
714     magma_setvector(num_qpts * num_nodes * q_comp_interp, sizeof(interp[0]), interp, 1, impl->d_interp, 1, data->queue);
715   }
716   if (curl) {
717     CeedInt q_comp_curl;
718 
719     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_CURL, &q_comp_curl));
720     CeedCallBackend(magma_malloc((void **)&impl->d_curl, num_qpts * num_nodes * q_comp_curl * sizeof(curl[0])));
721     magma_setvector(num_qpts * num_nodes * q_comp_curl, sizeof(curl[0]), curl, 1, impl->d_curl, 1, data->queue);
722   }
723 
724   // Compile the weight kernel if it won't be compiled later on
725   if (num_nodes > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P || num_qpts > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
726     Ceed        ceed_delegate;
727     char       *basis_kernel_source;
728     const char *weight_kernel_path;
729 
730     // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
731     CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
732 
733     // Compile weight kernel (the remainder of kernel compilation happens at first call to CeedBasisApply)
734     CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
735     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
736     CeedCallBackend(CeedLoadSourceToBuffer(ceed, weight_kernel_path, &basis_kernel_source));
737     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
738     CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[0], 1, "BASIS_Q", num_qpts));
739     CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[0], "magma_weight_nontensor", &impl->Weight));
740     CeedCallBackend(CeedFree(&weight_kernel_path));
741     CeedCallBackend(CeedFree(&basis_kernel_source));
742   }
743 
744   CeedCallBackend(CeedBasisSetData(basis, impl));
745 
746   // Register backend functions
747   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Magma));
748   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyNonTensor_Magma));
749   return CEED_ERROR_SUCCESS;
750 }
751 
752 //------------------------------------------------------------------------------
753