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