xref: /libCEED/rust/libceed-sys/c-src/include/ceed/jit-source/magma/magma-basis-interp-3d.h (revision f80f4a748154eed4bc661c135f695b92b1bc45b9)
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 // macros to abstract access of shared memory and reg. file
9 #define sT(i, j) sT[(j)*P_ + (i)]
10 #define sTmp(i, j, ldw) sTmp[(j) * (ldw) + (i)]
11 
12 //////////////////////////////////////////////////////////////////////////////////////////
13 // interp basis action (3D)
14 template <typename T, int DIM_U, int DIM_V, int NCOMP_, int P_, int Q_, int rUsize, int rVsize>
15 static __device__ __inline__ void magma_interp_3d_device(const T *sT, magma_trans_t transT, T rU[DIM_U][NCOMP_][rUsize], T rV[DIM_V][NCOMP_][rVsize],
16                                                          const int tx, T rTmp[Q_], T *swork) {
17   // Assumptions
18   // 1. 1D threads of size max(P_,Q_)^2
19   // 2. input:  rU[DIM_U x NCOMP_ x rUsize] in registers (per thread)
20   // 3. output: rV[DIM_V x NCOMP_ x rVsize] in registers (per thread)
21   // 4. Three products per component
22   //  4.1 Batch P_^2 of (1xP_) matrices times (P_xQ_) matrix => Batch P_^2 of (1xQ_) matrices
23   //  4.2 Batch P_   of (Q_xP_) matrices times (P_xQ_) matrix => Batch P_   of (Q_xQ_) matrices
24   //  4.3 Batch 1   of (Q_^2xP_) matrix times (P_xQ_) matrix => (Q_^2xQ_) matrix
25   // 5. Each thread computes one row of the output of each product
26   // 6. Sync is recommended before and after the call
27 
28   for (int icomp = 0; icomp < NCOMP_; icomp++) {
29     // Batch P_^2 of (1xP_) matrices [reg] times (P_xQ_) matrix [shmem] => Batch P_^2 of (1xQ_) matrices [shmem]
30     if (tx < (P_ * P_)) {
31       const int batchid = tx;
32       const int sld     = 1;
33       T        *sTmp    = swork + batchid * (1 * Q_);
34       for (int j = 0; j < Q_; j++) {
35         rTmp[0] = 0.0;
36         for (int i = 0; i < P_; i++) {
37           rTmp[0] += rU[0][icomp][i] * sT(i, j);
38         }
39         sTmp(0, j, sld) = rTmp[0];
40       }
41     }  // end of: if (tx < P_*P_)
42     __syncthreads();
43 
44     // Batch P_ of (Q_xP_) matrices [shmem] times (P_xQ_) matrix [shmem] => Batch P_ of (Q_xQ_) matrices [reg]
45     if (tx < (P_ * Q_)) {
46       const int batchid = tx / Q_;
47       const int tx_     = tx % Q_;
48       const int sld     = Q_;
49       T        *sTmp    = swork + batchid * (Q_ * P_);  // sTmp is input
50       for (int j = 0; j < Q_; j++) {
51         rTmp[j] = 0.0;
52         for (int i = 0; i < P_; i++) {
53           rTmp[j] += sTmp(tx_, i, sld) * sT(i, j);
54         }
55       }
56     }
57     __syncthreads();
58 
59     // write rTmp[] into shmem as batch P_ of Q_xQ_ matrices
60     if (tx < (P_ * Q_)) {
61       const int batchid = tx / Q_;
62       const int tx_     = tx % Q_;
63       const int sld     = Q_;
64       T        *sTmp    = swork + batchid * (Q_ * Q_);
65       for (int j = 0; j < Q_; j++) {
66         sTmp(tx_, j, sld) = rTmp[j];
67       }
68     }
69     __syncthreads();
70 
71     // Batch 1 of (Q_^2xP_) matrices [shmem] times (P_xQ_) matrix [shmem] => Batch 1 of (Q_^2xQ_) matrices [reg]
72     if (tx < (Q_ * Q_)) {
73       // No need to declare batchid = (tx  / Q_^2) = always zero
74       // No need to declare tx_     = (tx_ % Q_^2) = always tx
75       const int sld  = Q_ * Q_;
76       T        *sTmp = swork;
77       for (int j = 0; j < Q_; j++) {
78         rTmp[0] = 0.0;
79         for (int i = 0; i < P_; i++) {
80           rTmp[0] += sTmp(tx, i, sld) * sT(i, j);
81         }
82         rV[0][icomp][j] += rTmp[0];
83       }
84     }
85     __syncthreads();
86   }
87 }
88 
89 //////////////////////////////////////////////////////////////////////////////////////////
90 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(MAXPQ *MAXPQ, MAGMA_MAXTHREADS_3D)) __global__
91     void magma_interpn_3d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV,
92                                  const int cstrdV, const int nelem) {
93   MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
94 
95   const int     tx      = threadIdx.x;
96   const int     ty      = threadIdx.y;
97   const int     elem_id = (blockIdx.x * blockDim.y) + ty;
98   magma_trans_t transT  = MagmaNoTrans;
99 
100   if (elem_id >= nelem) return;
101 
102   CeedScalar rU[1][NCOMP][P] = {0.0};  // for a non fused operator DIM is always 1
103   CeedScalar rV[1][NCOMP][Q] = {0.0};  // for a non fused operator DIM is always 1
104   CeedScalar rTmp[Q]         = {0.0};
105 
106   // shift global memory pointers by elem stride
107   dU += elem_id * estrdU;
108   dV += elem_id * estrdV;
109 
110   // assign shared memory pointers
111   CeedScalar *sT   = (CeedScalar *)(shared_data);
112   CeedScalar *sTmp = sT + P * Q;
113   sTmp += ty * (max(P * P * MAXPQ, P * Q * Q));
114 
115   // read T
116   if (ty == 0) {
117     dread_T_gm2sm<P, Q>(tx, transT, dT, sT);
118   }
119 
120   // read U (idim = 0 for dU, iDIM = 0 for rU, u_dimstride is always 0)
121   readU_3d<CeedScalar, P, 1, NCOMP, P, 0>(dU, cstrdU, rU, sTmp, tx);
122   // there is a sync at the end of this function
123 
124   magma_interp_3d_device<CeedScalar, 1, 1, NCOMP, P, Q, P, Q>(sT, transT, rU, rV, tx, rTmp, sTmp);
125   __syncthreads();
126 
127   // write V
128   writeV_3d<CeedScalar, Q, 1, NCOMP, Q, 0>(dV, cstrdV, rV, tx);
129 }
130 
131 //////////////////////////////////////////////////////////////////////////////////////////
132 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(MAXPQ *MAXPQ, MAGMA_MAXTHREADS_3D)) __global__
133     void magma_interpt_3d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV,
134                                  const int cstrdV, const int nelem) {
135   MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
136 
137   const int     tx      = threadIdx.x;
138   const int     ty      = threadIdx.y;
139   const int     elem_id = (blockIdx.x * blockDim.y) + ty;
140   magma_trans_t transT  = MagmaTrans;
141 
142   if (elem_id >= nelem) return;
143 
144   CeedScalar rU[1][NCOMP][Q] = {0.0};  // for a non fused operator DIM is always 1
145   CeedScalar rV[1][NCOMP][P] = {0.0};  // for a non fused operator DIM is always 1
146   CeedScalar rTmp[P]         = {0.0};
147 
148   // shift global memory pointers by elem stride
149   dU += elem_id * estrdU;
150   dV += elem_id * estrdV;
151 
152   // assign shared memory pointers
153   CeedScalar *sT   = (CeedScalar *)(shared_data);
154   CeedScalar *sTmp = sT + Q * P;
155   sTmp += ty * (max(Q * Q * MAXPQ, Q * P * P));
156 
157   // read T
158   if (ty == 0) {
159     dread_T_gm2sm<Q, P>(tx, transT, dT, sT);
160   }
161 
162   // read V
163   readV_3d<CeedScalar, P, 1, NCOMP, P, 0>(dV, cstrdV, rV, tx);
164 
165   // read U (idim = 0 for dU, iDIM = 0 for rU, u_dimstride is always 0)
166   readU_3d<CeedScalar, Q, 1, NCOMP, Q, 0>(dU, cstrdU, rU, sTmp, tx);
167   // there is a sync at the end of this function
168 
169   magma_interp_3d_device<CeedScalar, 1, 1, NCOMP, Q, P, Q, P>(sT, transT, rU, rV, tx, rTmp, sTmp);
170   __syncthreads();
171 
172   // write V
173   writeV_3d<CeedScalar, P, 1, NCOMP, P, 0>(dV, cstrdV, rV, tx);
174 }
175