1 // Copyright (c) 2017-2026, 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 3D
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 (3D)
18 template <typename T, int DIM_U, int DIM_V, int NUM_COMP, int P, int Q, int rU_SIZE, int rV_SIZE>
magma_interp_3d_device(const T * sT,T rU[DIM_U][NUM_COMP][rU_SIZE],T rV[DIM_V][NUM_COMP][rV_SIZE],const int tx,T rTmp[Q],T * swork)19 static __device__ __inline__ void magma_interp_3d_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[Q], T *swork) {
21 // Assumptions
22 // 1. 1D threads of size max(P,Q)^2
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. Three products per component
26 // 4.1 Batch P^2 of (1xP) matrices times (PxQ) matrix => Batch P^2 of (1xQ) matrices
27 // 4.2 Batch P of (QxP) matrices times (PxQ) matrix => Batch P of (QxQ) matrices
28 // 4.3 Batch 1 of (Q^2xP_) matrix times (PxQ) matrix => (Q^2xQ_) matrix
29 // 5. Each thread computes one row of the output of each product
30 // 6. Sync is recommended before and after the call
31
32 for (int comp = 0; comp < NUM_COMP; comp++) {
33 // Batch P^2 of (1xP) matrices [reg] times (PxQ) matrix [shmem] => Batch P^2 of (1xQ) matrices [shmem]
34 if (tx < (P * 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.0;
40 for (int i = 0; i < P; i++) {
41 rTmp[0] += rU[0][comp][i] * sT(i, j);
42 }
43 sTmp(0, j, sld) = rTmp[0];
44 }
45 } // end of: if (tx < P*P)
46 __syncthreads();
47
48 // Batch P of (QxP) matrices [shmem] times (PxQ) matrix [shmem] => Batch P of (QxQ) matrices [reg]
49 if (tx < (P * Q)) {
50 const int batchid = tx / Q;
51 const int tx_ = tx % Q;
52 const int sld = Q;
53 T *sTmp = swork + batchid * (Q * P); // sTmp is input
54 for (int j = 0; j < Q; j++) {
55 rTmp[j] = 0.0;
56 for (int i = 0; i < P; i++) {
57 rTmp[j] += sTmp(tx_, i, sld) * sT(i, j);
58 }
59 }
60 }
61 __syncthreads();
62
63 // write rTmp[] into shmem as batch P of QxQ matrices
64 if (tx < (P * Q)) {
65 const int batchid = tx / Q;
66 const int tx_ = tx % Q;
67 const int sld = Q;
68 T *sTmp = swork + batchid * (Q * Q);
69 for (int j = 0; j < Q; j++) {
70 sTmp(tx_, j, sld) = rTmp[j];
71 }
72 }
73 __syncthreads();
74
75 // Batch 1 of (Q^2xP_) matrices [shmem] times (PxQ) matrix [shmem] => Batch 1 of (Q^2xQ_) matrices [reg]
76 if (tx < (Q * Q)) {
77 // No need to declare batchid = (tx / Q^2) = always zero
78 // No need to declare tx_ = (tx_ % Q^2) = always tx
79 const int sld = Q * Q;
80 T *sTmp = swork;
81 for (int j = 0; j < Q; j++) {
82 rTmp[0] = 0.0;
83 for (int i = 0; i < P; i++) {
84 rTmp[0] += sTmp(tx, i, sld) * sT(i, j);
85 }
86 rV[0][comp][j] += rTmp[0];
87 }
88 }
89 __syncthreads();
90 }
91 }
92
93 ////////////////////////////////////////////////////////////////////////////////
__launch_bounds__(MAGMA_BASIS_BOUNDS (BASIS_MAX_P_Q * BASIS_MAX_P_Q,MAGMA_MAXTHREADS_3D))94 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q *BASIS_MAX_P_Q, MAGMA_MAXTHREADS_3D)) __global__
95 void magma_interpn_3d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV,
96 const int cstrdV, const int nelem) {
97 MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
98
99 const int tx = threadIdx.x;
100 const int ty = threadIdx.y;
101 const int elem_id = (blockIdx.x * blockDim.y) + ty;
102
103 if (elem_id >= nelem) return;
104
105 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // for a non-fused operator BASIS_DIM is always 1
106 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // for a non-fused operator BASIS_DIM is always 1
107 CeedScalar rTmp[BASIS_Q] = {0.0};
108
109 // shift global memory pointers by elem stride
110 dU += elem_id * estrdU;
111 dV += elem_id * estrdV;
112
113 // assign shared memory pointers
114 CeedScalar *sT = (CeedScalar *)shared_data;
115 CeedScalar *sTmp = sT + BASIS_P * BASIS_Q;
116 sTmp += ty * (max(BASIS_P * BASIS_P * BASIS_MAX_P_Q, BASIS_P * BASIS_Q * BASIS_Q));
117
118 // read T
119 if (ty == 0) {
120 read_T_notrans_gm2sm<BASIS_P, BASIS_Q>(tx, dT, sT);
121 }
122
123 // read U (idim = 0 for dU, i_DIM = 0 for rU, u_dimstride is always 0)
124 read_U_3d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dU, cstrdU, rU, sTmp, tx);
125 // there is a sync at the end of this function
126
127 magma_interp_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_P, BASIS_Q, BASIS_P, BASIS_Q>(sT, rU, rV, tx, rTmp, sTmp);
128 __syncthreads();
129
130 // write V
131 write_V_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dV, cstrdV, rV, tx);
132 }
133
134 ////////////////////////////////////////////////////////////////////////////////
__launch_bounds__(MAGMA_BASIS_BOUNDS (BASIS_MAX_P_Q * BASIS_MAX_P_Q,MAGMA_MAXTHREADS_3D))135 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q *BASIS_MAX_P_Q, MAGMA_MAXTHREADS_3D)) __global__
136 void magma_interpt_3d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV,
137 const int cstrdV, const int nelem) {
138 MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
139
140 const int tx = threadIdx.x;
141 const int ty = threadIdx.y;
142 const int elem_id = (blockIdx.x * blockDim.y) + ty;
143
144 if (elem_id >= nelem) return;
145
146 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // for a non-fused operator BASIS_DIM is always 1
147 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // for a non-fused operator BASIS_DIM is always 1
148 CeedScalar rTmp[BASIS_P] = {0.0};
149
150 // shift global memory pointers by elem stride
151 dU += elem_id * estrdU;
152 dV += elem_id * estrdV;
153
154 // assign shared memory pointers
155 CeedScalar *sT = (CeedScalar *)shared_data;
156 CeedScalar *sTmp = sT + BASIS_Q * BASIS_P;
157 sTmp += ty * (max(BASIS_Q * BASIS_Q * BASIS_MAX_P_Q, BASIS_Q * BASIS_P * BASIS_P));
158
159 // read T
160 if (ty == 0) {
161 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dT, sT);
162 }
163
164 // read U (idim = 0 for dU, i_DIM = 0 for rU, u_dimstride is always 0)
165 read_U_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU, cstrdU, rU, sTmp, tx);
166 // there is a sync at the end of this function
167
168 magma_interp_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P>(sT, rU, rV, tx, rTmp, sTmp);
169 __syncthreads();
170
171 // write V
172 write_V_3d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV, cstrdV, rV, tx);
173 }
174
175 ////////////////////////////////////////////////////////////////////////////////
__launch_bounds__(MAGMA_BASIS_BOUNDS (BASIS_MAX_P_Q * BASIS_MAX_P_Q,MAGMA_MAXTHREADS_3D))176 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q *BASIS_MAX_P_Q, MAGMA_MAXTHREADS_3D)) __global__
177 void magma_interpta_3d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV,
178 const int cstrdV, const int nelem) {
179 MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
180
181 const int tx = threadIdx.x;
182 const int ty = threadIdx.y;
183 const int elem_id = (blockIdx.x * blockDim.y) + ty;
184
185 if (elem_id >= nelem) return;
186
187 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // for a non-fused operator BASIS_DIM is always 1
188 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // for a non-fused operator BASIS_DIM is always 1
189 CeedScalar rTmp[BASIS_P] = {0.0};
190
191 // shift global memory pointers by elem stride
192 dU += elem_id * estrdU;
193 dV += elem_id * estrdV;
194
195 // assign shared memory pointers
196 CeedScalar *sT = (CeedScalar *)shared_data;
197 CeedScalar *sTmp = sT + BASIS_Q * BASIS_P;
198 sTmp += ty * (max(BASIS_Q * BASIS_Q * BASIS_MAX_P_Q, BASIS_Q * BASIS_P * BASIS_P));
199
200 // read T
201 if (ty == 0) {
202 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dT, sT);
203 }
204
205 // read U (idim = 0 for dU, i_DIM = 0 for rU, u_dimstride is always 0)
206 read_U_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU, cstrdU, rU, sTmp, tx);
207 // there is a sync at the end of this function
208
209 magma_interp_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P>(sT, rU, rV, tx, rTmp, sTmp);
210 __syncthreads();
211
212 // sum into V
213 sum_V_3d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV, cstrdV, rV, tx);
214 }
215