xref: /libCEED/backends/magma/ceed-magma-basis.c (revision 1ea55a34129e65adf46cc9aad604e9be02563d89)
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     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_path, *weight_kernel_path, *basis_kernel_source;
296       magma_int_t arch = magma_getdevice_arch();
297 
298       // Tuning parameters for NB
299       CeedCallBackend(CeedBasisGetFESpace(basis, &fe_space));
300       CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
301       switch (fe_space) {
302         case CEED_FE_SPACE_H1:
303           CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_GRAD, &q_comp_deriv));
304           break;
305         case CEED_FE_SPACE_HDIV:
306           CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_DIV, &q_comp_deriv));
307           break;
308         case CEED_FE_SPACE_HCURL:
309           CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_CURL, &q_comp_deriv));
310           break;
311       }
312       impl->NB_interp[iN]   = nontensor_rtc_get_nb(arch, 'n', q_comp_interp, P, Q, n_array[iN]);
313       impl->NB_interp_t[iN] = nontensor_rtc_get_nb(arch, 't', q_comp_interp, P, Q, n_array[iN]);
314       impl->NB_deriv[iN]    = nontensor_rtc_get_nb(arch, 'n', q_comp_deriv, P, Q, n_array[iN]);
315       impl->NB_deriv_t[iN]  = nontensor_rtc_get_nb(arch, 't', q_comp_deriv, P, Q, n_array[iN]);
316 
317       // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
318       CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
319 
320       // Compile kernels
321       CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-interp-deriv-nontensor.h", &basis_kernel_path));
322       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
323       CeedCallBackend(CeedLoadSourceToBuffer(ceed, basis_kernel_path, &basis_kernel_source));
324       if (!impl->Weight) {
325         CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
326         CeedCallBackend(CeedLoadSourceToInitializedBuffer(ceed, weight_kernel_path, &basis_kernel_source));
327       }
328       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
329       CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[iN], 8, "BASIS_Q_COMP_INTERP", q_comp_interp,
330                                        "BASIS_Q_COMP_DERIV", q_comp_deriv, "BASIS_P", P, "BASIS_Q", Q, "BASIS_NB_INTERP_N", impl->NB_interp[iN],
331                                        "BASIS_NB_INTERP_T", impl->NB_interp_t[iN], "BASIS_NB_DERIV_N", impl->NB_deriv[iN], "BASIS_NB_DERIV_T",
332                                        impl->NB_deriv_t[iN]));
333       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[iN], "magma_interp_nontensor_n", &impl->Interp[iN]));
334       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[iN], "magma_interp_nontensor_t", &impl->InterpTranspose[iN]));
335       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[iN], "magma_deriv_nontensor_n", &impl->Deriv[iN]));
336       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[iN], "magma_deriv_nontensor_t", &impl->DerivTranspose[iN]));
337       if (!impl->Weight) {
338         CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[iN], "magma_weight_nontensor", &impl->Weight));
339         CeedCallBackend(CeedFree(&weight_kernel_path));
340       }
341       CeedCallBackend(CeedFree(&basis_kernel_path));
342       CeedCallBackend(CeedFree(&basis_kernel_source));
343     }
344   }
345 
346   // Apply basis operation
347   if (e_mode != CEED_EVAL_WEIGHT) {
348     const CeedScalar *d_b = NULL;
349     CeedInt           q_comp, NB, M, K;
350     CeedMagmaFunction Kernel;
351 
352     switch (e_mode) {
353       case CEED_EVAL_INTERP:
354         d_b = impl->d_interp;
355         break;
356       case CEED_EVAL_GRAD:
357         d_b = impl->d_grad;
358         break;
359       case CEED_EVAL_DIV:
360         d_b = impl->d_div;
361         break;
362       case CEED_EVAL_CURL:
363         d_b = impl->d_curl;
364         break;
365       // LCOV_EXCL_START
366       case CEED_EVAL_WEIGHT:
367       case CEED_EVAL_NONE:
368         return CeedError(ceed, CEED_ERROR_BACKEND, "%s does not make sense in this context", CeedEvalModes[e_mode]);
369         // LCOV_EXCL_STOP
370     }
371     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, e_mode, &q_comp));
372     M = (t_mode == CEED_TRANSPOSE) ? P : Q, K = (t_mode == CEED_TRANSPOSE) ? Q : P;
373 
374     if (P <= MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P && Q <= MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
375       if (e_mode == CEED_EVAL_INTERP) {
376         if (t_mode == CEED_TRANSPOSE) {
377           Kernel = impl->InterpTranspose[iN];
378           NB     = impl->NB_interp_t[iN];
379         } else {
380           Kernel = impl->Interp[iN];
381           NB     = impl->NB_interp[iN];
382         }
383       } else {
384         if (t_mode == CEED_TRANSPOSE) {
385           Kernel = impl->DerivTranspose[iN];
386           NB     = impl->NB_deriv_t[iN];
387         } else {
388           Kernel = impl->Deriv[iN];
389           NB     = impl->NB_deriv[iN];
390         }
391       }
392       CeedInt num_t_col    = MAGMA_BASIS_NTCOL(M, MAGMA_MAXTHREADS_1D);
393       CeedInt grid         = CeedDivUpInt(N, num_t_col * NB);
394       CeedInt shared_mem_A = P * Q * sizeof(CeedScalar);
395       CeedInt shared_mem_B = num_t_col * K * NB * sizeof(CeedScalar);
396       CeedInt shared_mem   = (t_mode != CEED_TRANSPOSE && q_comp > 1) ? (shared_mem_A + shared_mem_B) : CeedIntMax(shared_mem_A, shared_mem_B);
397       void   *args[]       = {&N, &d_b, &d_u, &d_v};
398 
399       CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, Kernel, grid, M, num_t_col, 1, shared_mem, args));
400     } else {
401       for (CeedInt d = 0; d < q_comp; d++) {
402         if (t_mode == CEED_TRANSPOSE) {
403           const CeedScalar beta = (d > 0) ? 1.0 : 0.0;
404           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);
405         } else {
406           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);
407         }
408       }
409     }
410   } else {
411     CeedCheck(t_mode != CEED_TRANSPOSE, ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT incompatible with CEED_TRANSPOSE");
412     CeedInt num_t_col  = MAGMA_BASIS_NTCOL(Q, MAGMA_MAXTHREADS_1D);
413     CeedInt grid       = CeedDivUpInt(num_elem, num_t_col);
414     CeedInt shared_mem = Q * sizeof(CeedScalar) + num_t_col * Q * sizeof(CeedScalar);
415     void   *args[]     = {&num_elem, &impl->d_q_weight, &d_v};
416 
417     CeedCallBackend(CeedRunKernelDimSharedMagma(ceed, impl->Weight, grid, Q, num_t_col, 1, shared_mem, args));
418   }
419 
420   // Must sync to ensure completeness
421   ceed_magma_queue_sync(data->queue);
422 
423   // Restore vectors
424   if (e_mode != CEED_EVAL_WEIGHT) {
425     CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
426   }
427   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
428   return CEED_ERROR_SUCCESS;
429 }
430 
431 //------------------------------------------------------------------------------
432 // Destroy tensor basis
433 //------------------------------------------------------------------------------
434 static int CeedBasisDestroy_Magma(CeedBasis basis) {
435   Ceed             ceed;
436   CeedBasis_Magma *impl;
437 
438   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
439   CeedCallBackend(CeedBasisGetData(basis, &impl));
440 #ifdef CEED_MAGMA_USE_HIP
441   CeedCallHip(ceed, hipModuleUnload(impl->module));
442 #else
443   CeedCallCuda(ceed, cuModuleUnload(impl->module));
444 #endif
445   CeedCallBackend(magma_free(impl->d_interp_1d));
446   CeedCallBackend(magma_free(impl->d_grad_1d));
447   CeedCallBackend(magma_free(impl->d_q_weight_1d));
448   CeedCallBackend(CeedFree(&impl));
449   return CEED_ERROR_SUCCESS;
450 }
451 
452 //------------------------------------------------------------------------------
453 // Destroy non-tensor basis
454 //------------------------------------------------------------------------------
455 static int CeedBasisDestroyNonTensor_Magma(CeedBasis basis) {
456   Ceed                      ceed;
457   CeedBasisNonTensor_Magma *impl;
458 
459   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
460   CeedCallBackend(CeedBasisGetData(basis, &impl));
461   for (CeedInt in = 0; in < MAGMA_NONTENSOR_KERNEL_INSTANCES; in++) {
462     if (impl->module[in]) {
463 #ifdef CEED_MAGMA_USE_HIP
464       CeedCallHip(ceed, hipModuleUnload(impl->module[in]));
465 #else
466       CeedCallCuda(ceed, cuModuleUnload(impl->module[in]));
467 #endif
468     }
469   }
470   CeedCallBackend(magma_free(impl->d_interp));
471   CeedCallBackend(magma_free(impl->d_grad));
472   CeedCallBackend(magma_free(impl->d_div));
473   CeedCallBackend(magma_free(impl->d_curl));
474   CeedCallBackend(magma_free(impl->d_q_weight));
475   CeedCallBackend(CeedFree(&impl));
476   return CEED_ERROR_SUCCESS;
477 }
478 
479 //------------------------------------------------------------------------------
480 // Create tensor
481 //------------------------------------------------------------------------------
482 int CeedBasisCreateTensorH1_Magma(CeedInt dim, CeedInt P_1d, CeedInt Q_1d, const CeedScalar *interp_1d, const CeedScalar *grad_1d,
483                                   const CeedScalar *q_ref_1d, const CeedScalar *q_weight_1d, CeedBasis basis) {
484   Ceed             ceed, ceed_delegate;
485   Ceed_Magma      *data;
486   char            *interp_kernel_path, *grad_kernel_path, *weight_kernel_path, *basis_kernel_source;
487   CeedInt          num_comp;
488   CeedBasis_Magma *impl;
489 
490   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
491   CeedCallBackend(CeedGetData(ceed, &data));
492   CeedCallBackend(CeedCalloc(1, &impl));
493 
494   // Copy basis data to GPU
495   CeedCallBackend(magma_malloc((void **)&impl->d_q_weight_1d, Q_1d * sizeof(q_weight_1d[0])));
496   magma_setvector(Q_1d, sizeof(q_weight_1d[0]), q_weight_1d, 1, impl->d_q_weight_1d, 1, data->queue);
497   CeedCallBackend(magma_malloc((void **)&impl->d_interp_1d, Q_1d * P_1d * sizeof(interp_1d[0])));
498   magma_setvector(Q_1d * P_1d, sizeof(interp_1d[0]), interp_1d, 1, impl->d_interp_1d, 1, data->queue);
499   CeedCallBackend(magma_malloc((void **)&impl->d_grad_1d, Q_1d * P_1d * sizeof(grad_1d[0])));
500   magma_setvector(Q_1d * P_1d, sizeof(grad_1d[0]), grad_1d, 1, impl->d_grad_1d, 1, data->queue);
501 
502   // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
503   CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
504 
505   // Compile kernels
506   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
507   {
508     char   *interp_kernel_name_base = "ceed/jit-source/magma/magma-basis-interp";
509     CeedInt interp_kernel_name_len  = strlen(interp_kernel_name_base) + 6;
510     char    interp_kernel_name[interp_kernel_name_len];
511 
512     snprintf(interp_kernel_name, interp_kernel_name_len, "%s-%" CeedInt_FMT "d.h", interp_kernel_name_base, dim);
513     CeedCallBackend(CeedGetJitAbsolutePath(ceed, interp_kernel_name, &interp_kernel_path));
514   }
515   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
516   CeedCallBackend(CeedLoadSourceToBuffer(ceed, interp_kernel_path, &basis_kernel_source));
517   {
518     char   *grad_kernel_name_base = "ceed/jit-source/magma/magma-basis-grad";
519     CeedInt grad_kernel_name_len  = strlen(grad_kernel_name_base) + 6;
520     char    grad_kernel_name[grad_kernel_name_len];
521 
522     snprintf(grad_kernel_name, grad_kernel_name_len, "%s-%" CeedInt_FMT "d.h", grad_kernel_name_base, dim);
523     CeedCallBackend(CeedGetJitAbsolutePath(ceed, grad_kernel_name, &grad_kernel_path));
524   }
525   CeedCallBackend(CeedLoadSourceToInitializedBuffer(ceed, grad_kernel_path, &basis_kernel_source));
526   {
527     char   *weight_kernel_name_base = "ceed/jit-source/magma/magma-basis-weight";
528     CeedInt weight_kernel_name_len  = strlen(weight_kernel_name_base) + 6;
529     char    weight_kernel_name[weight_kernel_name_len];
530 
531     snprintf(weight_kernel_name, weight_kernel_name_len, "%s-%" CeedInt_FMT "d.h", weight_kernel_name_base, dim);
532     CeedCallBackend(CeedGetJitAbsolutePath(ceed, weight_kernel_name, &weight_kernel_path));
533   }
534   CeedCallBackend(CeedLoadSourceToInitializedBuffer(ceed, weight_kernel_path, &basis_kernel_source));
535   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
536   CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module, 5, "BASIS_DIM", dim, "BASIS_NUM_COMP", num_comp, "BASIS_P",
537                                    P_1d, "BASIS_Q", Q_1d, "BASIS_MAX_P_Q", CeedIntMax(P_1d, Q_1d)));
538   switch (dim) {
539     case 1:
540       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpn_1d_kernel", &impl->Interp));
541       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpt_1d_kernel", &impl->InterpTranspose));
542       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradn_1d_kernel", &impl->Grad));
543       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradt_1d_kernel", &impl->GradTranspose));
544       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_weight_1d_kernel", &impl->Weight));
545       break;
546     case 2:
547       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpn_2d_kernel", &impl->Interp));
548       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpt_2d_kernel", &impl->InterpTranspose));
549       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradn_2d_kernel", &impl->Grad));
550       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradt_2d_kernel", &impl->GradTranspose));
551       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_weight_2d_kernel", &impl->Weight));
552       break;
553     case 3:
554       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpn_3d_kernel", &impl->Interp));
555       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_interpt_3d_kernel", &impl->InterpTranspose));
556       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradn_3d_kernel", &impl->Grad));
557       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_gradt_3d_kernel", &impl->GradTranspose));
558       CeedCallBackend(CeedGetKernelMagma(ceed, impl->module, "magma_weight_3d_kernel", &impl->Weight));
559       break;
560   }
561   CeedCallBackend(CeedFree(&interp_kernel_path));
562   CeedCallBackend(CeedFree(&grad_kernel_path));
563   CeedCallBackend(CeedFree(&weight_kernel_path));
564   CeedCallBackend(CeedFree(&basis_kernel_source));
565 
566   CeedCallBackend(CeedBasisSetData(basis, impl));
567 
568   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApply_Magma));
569   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroy_Magma));
570   return CEED_ERROR_SUCCESS;
571 }
572 
573 //------------------------------------------------------------------------------
574 // Create non-tensor H^1
575 //------------------------------------------------------------------------------
576 int CeedBasisCreateH1_Magma(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp, const CeedScalar *grad,
577                             const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
578   Ceed                      ceed;
579   Ceed_Magma               *data;
580   CeedBasisNonTensor_Magma *impl;
581 
582   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
583   CeedCallBackend(CeedGetData(ceed, &data));
584   CeedCallBackend(CeedCalloc(1, &impl));
585 
586   // Copy basis data to GPU
587   CeedCallBackend(magma_malloc((void **)&impl->d_q_weight, num_qpts * sizeof(q_weight[0])));
588   magma_setvector(num_qpts, sizeof(q_weight[0]), q_weight, 1, impl->d_q_weight, 1, data->queue);
589   if (interp) {
590     CeedInt q_comp_interp;
591 
592     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
593     CeedCallBackend(magma_malloc((void **)&impl->d_interp, num_qpts * num_nodes * q_comp_interp * sizeof(interp[0])));
594     magma_setvector(num_qpts * num_nodes * q_comp_interp, sizeof(interp[0]), interp, 1, impl->d_interp, 1, data->queue);
595   }
596   if (grad) {
597     CeedInt q_comp_grad;
598 
599     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_GRAD, &q_comp_grad));
600     CeedCallBackend(magma_malloc((void **)&impl->d_grad, num_qpts * num_nodes * q_comp_grad * sizeof(grad[0])));
601     magma_setvector(num_qpts * num_nodes * q_comp_grad, sizeof(grad[0]), grad, 1, impl->d_grad, 1, data->queue);
602   }
603 
604   // Compile the weight kernel if it won't be compiled later on
605   if (num_nodes > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P || num_qpts > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
606     Ceed  ceed_delegate;
607     char *weight_kernel_path, *basis_kernel_source;
608 
609     // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
610     CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
611 
612     // Compile weight kernel (the remainder of kernel compilation happens at first call to CeedBasisApply)
613     CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
614     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
615     CeedCallBackend(CeedLoadSourceToBuffer(ceed, weight_kernel_path, &basis_kernel_source));
616     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
617     CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[0], 1, "BASIS_Q", num_qpts));
618     CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[0], "magma_weight_nontensor", &impl->Weight));
619     CeedCallBackend(CeedFree(&weight_kernel_path));
620     CeedCallBackend(CeedFree(&basis_kernel_source));
621   }
622 
623   CeedCallBackend(CeedBasisSetData(basis, impl));
624 
625   // Register backend functions
626   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Magma));
627   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyNonTensor_Magma));
628   return CEED_ERROR_SUCCESS;
629 }
630 
631 //------------------------------------------------------------------------------
632 // Create non-tensor H(div)
633 //------------------------------------------------------------------------------
634 int CeedBasisCreateHdiv_Magma(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp,
635                               const CeedScalar *div, const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
636   Ceed                      ceed;
637   Ceed_Magma               *data;
638   CeedBasisNonTensor_Magma *impl;
639 
640   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
641   CeedCallBackend(CeedGetData(ceed, &data));
642   CeedCallBackend(CeedCalloc(1, &impl));
643 
644   // Copy basis data to GPU
645   CeedCallBackend(magma_malloc((void **)&impl->d_q_weight, num_qpts * sizeof(q_weight[0])));
646   magma_setvector(num_qpts, sizeof(q_weight[0]), q_weight, 1, impl->d_q_weight, 1, data->queue);
647   if (interp) {
648     CeedInt q_comp_interp;
649 
650     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
651     CeedCallBackend(magma_malloc((void **)&impl->d_interp, num_qpts * num_nodes * q_comp_interp * sizeof(interp[0])));
652     magma_setvector(num_qpts * num_nodes * q_comp_interp, sizeof(interp[0]), interp, 1, impl->d_interp, 1, data->queue);
653   }
654   if (div) {
655     CeedInt q_comp_div;
656 
657     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_DIV, &q_comp_div));
658     CeedCallBackend(magma_malloc((void **)&impl->d_div, num_qpts * num_nodes * q_comp_div * sizeof(div[0])));
659     magma_setvector(num_qpts * num_nodes * q_comp_div, sizeof(div[0]), div, 1, impl->d_div, 1, data->queue);
660   }
661 
662   // Compile the weight kernel if it won't be compiled later on
663   if (num_nodes > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P || num_qpts > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
664     Ceed  ceed_delegate;
665     char *weight_kernel_path, *basis_kernel_source;
666 
667     // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
668     CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
669 
670     // Compile weight kernel (the remainder of kernel compilation happens at first call to CeedBasisApply)
671     CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
672     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
673     CeedCallBackend(CeedLoadSourceToBuffer(ceed, weight_kernel_path, &basis_kernel_source));
674     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
675     CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[0], 1, "BASIS_Q", num_qpts));
676     CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[0], "magma_weight_nontensor", &impl->Weight));
677     CeedCallBackend(CeedFree(&weight_kernel_path));
678     CeedCallBackend(CeedFree(&basis_kernel_source));
679   }
680 
681   CeedCallBackend(CeedBasisSetData(basis, impl));
682 
683   // Register backend functions
684   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Magma));
685   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyNonTensor_Magma));
686   return CEED_ERROR_SUCCESS;
687 }
688 
689 //------------------------------------------------------------------------------
690 // Create non-tensor H(curl)
691 //------------------------------------------------------------------------------
692 int CeedBasisCreateHcurl_Magma(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp,
693                                const CeedScalar *curl, const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
694   Ceed                      ceed;
695   Ceed_Magma               *data;
696   CeedBasisNonTensor_Magma *impl;
697 
698   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
699   CeedCallBackend(CeedGetData(ceed, &data));
700   CeedCallBackend(CeedCalloc(1, &impl));
701 
702   // Copy basis data to GPU
703   CeedCallBackend(magma_malloc((void **)&impl->d_q_weight, num_qpts * sizeof(q_weight[0])));
704   magma_setvector(num_qpts, sizeof(q_weight[0]), q_weight, 1, impl->d_q_weight, 1, data->queue);
705   if (interp) {
706     CeedInt q_comp_interp;
707 
708     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_INTERP, &q_comp_interp));
709     CeedCallBackend(magma_malloc((void **)&impl->d_interp, num_qpts * num_nodes * q_comp_interp * sizeof(interp[0])));
710     magma_setvector(num_qpts * num_nodes * q_comp_interp, sizeof(interp[0]), interp, 1, impl->d_interp, 1, data->queue);
711   }
712   if (curl) {
713     CeedInt q_comp_curl;
714 
715     CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, CEED_EVAL_CURL, &q_comp_curl));
716     CeedCallBackend(magma_malloc((void **)&impl->d_curl, num_qpts * num_nodes * q_comp_curl * sizeof(curl[0])));
717     magma_setvector(num_qpts * num_nodes * q_comp_curl, sizeof(curl[0]), curl, 1, impl->d_curl, 1, data->queue);
718   }
719 
720   // Compile the weight kernel if it won't be compiled later on
721   if (num_nodes > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_P || num_qpts > MAGMA_NONTENSOR_CUSTOM_KERNEL_MAX_Q) {
722     Ceed  ceed_delegate;
723     char *weight_kernel_path, *basis_kernel_source;
724 
725     // The RTC compilation code expects a Ceed with the common Ceed_Cuda or Ceed_Hip data
726     CeedCallBackend(CeedGetDelegate(ceed, &ceed_delegate));
727 
728     // Compile weight kernel (the remainder of kernel compilation happens at first call to CeedBasisApply)
729     CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/magma/magma-basis-weight-nontensor.h", &weight_kernel_path));
730     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
731     CeedCallBackend(CeedLoadSourceToBuffer(ceed, weight_kernel_path, &basis_kernel_source));
732     CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
733     CeedCallBackend(CeedCompileMagma(ceed_delegate, basis_kernel_source, &impl->module[0], 1, "BASIS_Q", num_qpts));
734     CeedCallBackend(CeedGetKernelMagma(ceed, impl->module[0], "magma_weight_nontensor", &impl->Weight));
735     CeedCallBackend(CeedFree(&weight_kernel_path));
736     CeedCallBackend(CeedFree(&basis_kernel_source));
737   }
738 
739   CeedCallBackend(CeedBasisSetData(basis, impl));
740 
741   // Register backend functions
742   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Magma));
743   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyNonTensor_Magma));
744   return CEED_ERROR_SUCCESS;
745 }
746 
747 //------------------------------------------------------------------------------
748