xref: /libCEED/include/ceed/jit-source/magma/magma-basis-interp-2d.h (revision 4f69910b6e3819988a1446e35e0e85e74672bc23)
1 // Copyright (c) 2017-2025, 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 /// @file
9 /// Internal header for MAGMA tensor basis interpolation in 1D
10 #include "magma-common-tensor.h"
11 
12 // macros to abstract access of shared memory and reg. file
13 #define sT(i, j) sT[(j) * P + (i)]
14 #define sTmp(i, j, ldw) sTmp[(j) * (ldw) + (i)]
15 
16 ////////////////////////////////////////////////////////////////////////////////
17 // interp basis action (2D)
18 template <typename T, int DIM_U, int DIM_V, int NUM_COMP, int P, int Q, int rU_SIZE, int rV_SIZE>
19 static __device__ __inline__ void magma_interp_2d_device(const T *sT, T rU[DIM_U][NUM_COMP][rU_SIZE], T rV[DIM_V][NUM_COMP][rV_SIZE], const int tx,
20                                                          T rTmp, T *swork) {
21   // Assumptions
22   // 1. 1D threads of size max(P,Q)
23   // 2. input:  rU[DIM_U x NUM_COMP x rU_SIZE] in registers (per thread)
24   // 3. output: rV[DIM_V x NUM_COMP x rV_SIZE] in registers (per thread)
25   // 4. Two products per component
26   //  4.1 Batch P of (1xP) matrices times (PxQ) matrix => Batch P of (1xQ) matrices
27   //  4.2 Batch 1 of (QxP) matrix   times (PxQ) matrix => (QxQ) matrix
28   // 5. Each thread computes one row of the output of each product
29   // 6. Sync is recommended before and after the call
30 
31   for (int comp = 0; comp < NUM_COMP; comp++) {
32     // 1st product -- Batch P of (1xP) matrices [reg] x (PxQ) [shmem] => Batch P of (1xQ) matrices
33     // the batch output P x (1xQ) is written on the fly to shmem
34     if (tx < P) {
35       const int batchid = tx;
36       const int sld     = 1;
37       T        *sTmp    = swork + batchid * (1 * Q);
38       for (int j = 0; j < Q; j++) {
39         rTmp = 0.0;
40         for (int i = 0; i < P; i++) {
41           rTmp += rU[0][comp][i] * sT(i, j);
42         }
43         sTmp(0, j, sld) = rTmp;
44       }
45     }  // end of: if (tx < P)
46     __syncthreads();
47 
48     // 2nd product -- Batch 1 of a (QxP) matrix [shmem] x (PxQ) [shmem] => (QxQ) matrix [reg]
49     if (tx < Q) {
50       const int batchid = 0;
51       const int sld     = Q;
52       T        *sTmp    = swork + batchid * (Q * P);
53       for (int j = 0; j < Q; j++) {
54         rTmp = 0.0;
55         for (int i = 0; i < P; i++) {
56           rTmp += sTmp(tx, i, sld) * sT(i, j);
57         }
58         rV[0][comp][j] += rTmp;
59       }
60     }
61     __syncthreads();
62   }
63 }
64 
65 ////////////////////////////////////////////////////////////////////////////////
66 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_2D)) __global__
67     void magma_interpn_2d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV,
68                                  const int cstrdV, const int nelem) {
69   MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
70 
71   const int tx      = threadIdx.x;
72   const int ty      = threadIdx.y;
73   const int elem_id = (blockIdx.x * blockDim.y) + ty;
74 
75   if (elem_id >= nelem) return;
76 
77   CeedScalar rU[1][BASIS_NUM_COMP][BASIS_P] = {0.0};  // for a non-fused operator BASIS_DIM is always 1
78   CeedScalar rV[1][BASIS_NUM_COMP][BASIS_Q] = {0.0};  // for a non-fused operator BASIS_DIM is always 1
79   CeedScalar rTmp                           = 0.0;
80 
81   // shift global memory pointers by elem stride
82   dU += elem_id * estrdU;
83   dV += elem_id * estrdV;
84 
85   // assign shared memory pointers
86   CeedScalar *sT   = (CeedScalar *)shared_data;
87   CeedScalar *sTmp = sT + BASIS_P * BASIS_Q;
88   sTmp += ty * (BASIS_P * BASIS_MAX_P_Q);
89 
90   // read T
91   if (ty == 0) {
92     read_T_notrans_gm2sm<BASIS_P, BASIS_Q>(tx, dT, sT);
93   }
94 
95   // read U -- there is a sync at the end of this function
96   read_U_2d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dU, cstrdU, rU, sTmp, tx);
97 
98   // no sync needed here -- read_U_2d already syncs at the end
99   magma_interp_2d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_P, BASIS_Q, BASIS_P, BASIS_Q>(sT, rU, rV, tx, rTmp, sTmp);
100   __syncthreads();
101 
102   // write V
103   write_V_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dV, cstrdV, rV, tx);
104 }
105 
106 ////////////////////////////////////////////////////////////////////////////////
107 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_2D)) __global__
108     void magma_interpt_2d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV,
109                                  const int cstrdV, const int nelem) {
110   MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
111 
112   const int tx      = threadIdx.x;
113   const int ty      = threadIdx.y;
114   const int elem_id = (blockIdx.x * blockDim.y) + ty;
115 
116   if (elem_id >= nelem) return;
117 
118   CeedScalar rU[1][BASIS_NUM_COMP][BASIS_Q] = {0.0};  // for a non-fused operator BASIS_DIM is always 1
119   CeedScalar rV[1][BASIS_NUM_COMP][BASIS_P] = {0.0};  // for a non-fused operator BASIS_DIM is always 1
120   CeedScalar rTmp                           = 0.0;
121 
122   // shift global memory pointers by elem stride
123   dU += elem_id * estrdU;
124   dV += elem_id * estrdV;
125 
126   // assign shared memory pointers
127   CeedScalar *sT   = (CeedScalar *)shared_data;
128   CeedScalar *sTmp = sT + BASIS_Q * BASIS_P;
129   sTmp += ty * (BASIS_Q * BASIS_MAX_P_Q);
130 
131   // read T
132   if (ty == 0) {
133     read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dT, sT);
134   }
135 
136   // read U -- there is a sync at the end of this function
137   read_U_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU, cstrdU, rU, sTmp, tx);
138 
139   // no sync needed here -- read_U_2d already syncs at the end
140   magma_interp_2d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P>(sT, rU, rV, tx, rTmp, sTmp);
141   __syncthreads();
142 
143   // write V
144   write_V_2d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV, cstrdV, rV, tx);
145 }
146 
147 ////////////////////////////////////////////////////////////////////////////////
148 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_2D)) __global__
149     void magma_interpta_2d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV,
150                                   const int cstrdV, const int nelem) {
151   MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
152 
153   const int tx      = threadIdx.x;
154   const int ty      = threadIdx.y;
155   const int elem_id = (blockIdx.x * blockDim.y) + ty;
156 
157   if (elem_id >= nelem) return;
158 
159   CeedScalar rU[1][BASIS_NUM_COMP][BASIS_Q] = {0.0};  // for a non-fused operator BASIS_DIM is always 1
160   CeedScalar rV[1][BASIS_NUM_COMP][BASIS_P] = {0.0};  // for a non-fused operator BASIS_DIM is always 1
161   CeedScalar rTmp                           = 0.0;
162 
163   // shift global memory pointers by elem stride
164   dU += elem_id * estrdU;
165   dV += elem_id * estrdV;
166 
167   // assign shared memory pointers
168   CeedScalar *sT   = (CeedScalar *)shared_data;
169   CeedScalar *sTmp = sT + BASIS_Q * BASIS_P;
170   sTmp += ty * (BASIS_Q * BASIS_MAX_P_Q);
171 
172   // read T
173   if (ty == 0) {
174     read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dT, sT);
175   }
176 
177   // read U -- there is a sync at the end of this function
178   read_U_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU, cstrdU, rU, sTmp, tx);
179 
180   // no sync needed here -- read_U_2d already syncs at the end
181   magma_interp_2d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P>(sT, rU, rV, tx, rTmp, sTmp);
182   __syncthreads();
183 
184   // sum into V
185   sum_V_2d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV, cstrdV, rV, tx);
186 }
187