1 // Copyright (c) 2017-2024, 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 gradient in 3D 10 11 #include "magma-common-tensor.h" 12 13 // macros to abstract access of shared memory and reg. file 14 #define sT(i, j) sT[(j) * P + (i)] 15 #define sTmp(i, j, ldw) sTmp[(j) * (ldw) + (i)] 16 #define sTmp2(i, j, ldw) sTmp2[(j) * (ldw) + (i)] 17 18 //////////////////////////////////////////////////////////////////////////////// 19 // Helper function to add or set into V 20 template <typename T, bool Add> 21 struct magma_grad_3d_device_accumulate; 22 23 template <typename T> 24 struct magma_grad_3d_device_accumulate<T, true> { 25 static __device__ __inline__ void op(T &rV, const T &rTmp) { rV += rTmp; } 26 }; 27 28 template <typename T> 29 struct magma_grad_3d_device_accumulate<T, false> { 30 static __device__ __inline__ void op(T &rV, const T &rTmp) { rV = rTmp; } 31 }; 32 33 //////////////////////////////////////////////////////////////////////////////// 34 // grad basis action (3D) 35 // This function is called three times at a higher level for 3D 36 // DIM_U -- for the size of rU[DIM_U * NUM_COMP * MAX_P_Q] 37 // DIM_V -- for the size of rV[DIM_V * NUM_COMP * MAX_P_Q] 38 // i_DIM -- the index of the outermost loop over dimensions in grad 39 // i_DIM_U -- which dim index of rU is accessed (always 0 for notrans, 0, 1, or 2 for trans) 40 // i_DIM_V -- which dim index of rV is accessed (0, 1, or 2 for notrans, always 0 for trans) 41 template <typename T, int DIM_U, int DIM_V, int NUM_COMP, int P, int Q, int rU_SIZE, int rV_SIZE, int i_DIM, int i_DIM_U, int i_DIM_V, bool ADD> 42 static __device__ __inline__ void magma_grad_3d_device(const T *sTinterp, const T *sTgrad, T rU[DIM_U][NUM_COMP][rU_SIZE], 43 T rV[DIM_V][NUM_COMP][rV_SIZE], const int tx, T rTmp, T *swork) { 44 // Assumptions 45 // 0. This device routine applies grad for one dim only (i_DIM), so it should be thrice for 3D 46 // 1. 1D threads of size max(P,Q)^2 47 // 2. input: rU[DIM_U x NUM_COMP x rU_SIZE] in registers (per thread) 48 // 3. output: rV[DIM_V x NUM_COMP x rV_SIZE] in registers (per thread) 49 // 4. Three products per each (dim,component) pair 50 // 4.1 Batch P^2 of (1xP) matrices times (PxQ) matrix => Batch P^2 of (1xQ) matrices 51 // 4.2 Batch P of (QxP) matrices times (PxQ) matrix => Batch P of (QxQ) matrices 52 // 4.3 Batch 1 of (Q^2xP_) matrix times (PxQ) matrix => (Q^2xQ_) matrix 53 // 6. Each thread computes one row of the output of each product 54 // 7. Sync is recommended before and after the call 55 56 T *sW1 = swork; 57 T *sW2 = sW1 + P * P * Q; 58 for (int comp = 0; comp < NUM_COMP; comp++) { 59 // Batch P^2 of (1xP) matrices [reg] times (PxQ) matrix [shmem] => Batch P^2 of (1xQ) matrices [shmem] 60 if (tx < (P * P)) { 61 const int batchid = tx; 62 const int sld = 1; 63 const T *sT = (i_DIM == 0) ? sTgrad : sTinterp; 64 T *sTmp = sW1 + batchid * (1 * Q); 65 for (int j = 0; j < Q; j++) { 66 rTmp = 0.0; 67 for (int i = 0; i < P; i++) { 68 rTmp += rU[i_DIM_U][comp][i] * sT(i, j); 69 } 70 sTmp(0, j, sld) = rTmp; 71 } 72 } // end of: if (tx < P*P) 73 __syncthreads(); 74 75 // Batch P of (QxP) matrices [shmem] times (PxQ) matrix [shmem] => Batch P of (QxQ) matrices [reg] 76 if (tx < (P * Q)) { 77 const int batchid = tx / Q; 78 const int tx_ = tx % Q; 79 const int sld = Q; 80 const T *sT = (i_DIM == 1) ? sTgrad : sTinterp; 81 T *sTmp = sW1 + batchid * (Q * P); // sTmp is input 82 T *sTmp2 = sW2 + batchid * (Q * Q); // sTmp2 is output 83 for (int j = 0; j < Q; j++) { 84 rTmp = 0.0; 85 for (int i = 0; i < P; i++) { 86 rTmp += sTmp(tx_, i, sld) * sT(i, j); 87 } 88 sTmp2(tx_, j, sld) = rTmp; 89 } 90 } 91 __syncthreads(); 92 93 // Batch 1 of (Q^2xP_) matrices [shmem] times (PxQ) matrix [shmem] => Batch 1 of (Q^2xQ_) matrices [reg] 94 if (tx < (Q * Q)) { 95 // No need to declare batchid = (tx / Q^2) = always zero 96 // No need to declare tx_ = (tx_ % Q^2) = always tx 97 const int sld = Q * Q; 98 const T *sT = (i_DIM == 2) ? sTgrad : sTinterp; 99 T *sTmp = sW2; // sTmp is input 100 for (int j = 0; j < Q; j++) { 101 rTmp = 0.0; 102 for (int i = 0; i < P; i++) { 103 rTmp += sTmp(tx, i, sld) * sT(i, j); 104 } 105 magma_grad_3d_device_accumulate<T, ADD>::op(rV[i_DIM_V][comp][j], rTmp); 106 } 107 } 108 __syncthreads(); 109 } // loop over NUM_COMP 110 } 111 112 //////////////////////////////////////////////////////////////////////////////// 113 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q *BASIS_MAX_P_Q, MAGMA_MAXTHREADS_3D)) __global__ 114 void magma_gradn_3d_kernel(const CeedScalar *dinterp1d, const CeedScalar *dgrad1d, const CeedScalar *dU, const int estrdU, const int cstrdU, 115 const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) { 116 MAGMA_DEVICE_SHARED(CeedScalar, shared_data) 117 118 const int tx = threadIdx.x; 119 const int ty = threadIdx.y; 120 const int elem_id = (blockIdx.x * blockDim.y) + ty; 121 122 if (elem_id >= nelem) return; 123 124 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // here DIM_U = 1, but might be different for a fused operator 125 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // here DIM_V = 1, but might be different for a fused operator 126 CeedScalar rTmp = 0.0; 127 128 // shift global memory pointers by elem stride 129 dU += elem_id * estrdU; 130 dV += elem_id * estrdV; 131 132 // assign shared memory pointers 133 CeedScalar *sTinterp = (CeedScalar *)shared_data; 134 CeedScalar *sTgrad = sTinterp + BASIS_P * BASIS_Q; 135 CeedScalar *sTmp = sTgrad + BASIS_P * BASIS_Q; 136 sTmp += ty * (max(BASIS_P * BASIS_P * BASIS_P, (BASIS_P * BASIS_P * BASIS_Q) + (BASIS_P * BASIS_Q * BASIS_Q))); 137 138 // read T 139 if (ty == 0) { 140 read_T_notrans_gm2sm<BASIS_P, BASIS_Q>(tx, dinterp1d, sTinterp); 141 read_T_notrans_gm2sm<BASIS_P, BASIS_Q>(tx, dgrad1d, sTgrad); 142 } 143 __syncthreads(); 144 145 /* read U (idim = 0 for dU, i_DIM = 0 for rU) -- 146 there is a sync at the end of this function */ 147 read_U_3d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dU + (0 * dstrdU), cstrdU, rU, sTmp, tx); 148 149 /* first call (i_DIM = 0, i_DIM_U = 0, i_DIM_V = 0) -- 150 output from rV[0][][] into dV (idim = 0) */ 151 magma_grad_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_P, BASIS_Q, BASIS_P, BASIS_Q, 0, 0, 0, false>(sTinterp, sTgrad, rU, rV, tx, rTmp, 152 sTmp); 153 /* there is a sync at the end of magma_grad_3d_device */ 154 write_V_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dV + (0 * dstrdV), cstrdV, rV, tx); 155 156 /* second call (i_DIM = 1, i_DIM_U = 0, i_DIM_V = 0) -- 157 output from rV[0][][] into dV (idim = 1) */ 158 magma_grad_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_P, BASIS_Q, BASIS_P, BASIS_Q, 1, 0, 0, false>(sTinterp, sTgrad, rU, rV, tx, rTmp, 159 sTmp); 160 /* there is a sync at the end of magma_grad_3d_device */ 161 write_V_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dV + (1 * dstrdV), cstrdV, rV, tx); 162 163 /* third call (i_DIM = 2, i_DIM_U = 0, i_DIM_V = 0) -- 164 output from rV[0][][] into dV (idim = 2) */ 165 magma_grad_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_P, BASIS_Q, BASIS_P, BASIS_Q, 2, 0, 0, false>(sTinterp, sTgrad, rU, rV, tx, rTmp, 166 sTmp); 167 /* there is a sync at the end of magma_grad_3d_device */ 168 write_V_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dV + (2 * dstrdV), cstrdV, rV, tx); 169 } 170 171 //////////////////////////////////////////////////////////////////////////////// 172 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q *BASIS_MAX_P_Q, MAGMA_MAXTHREADS_3D)) __global__ 173 void magma_gradt_3d_kernel(const CeedScalar *dinterp1d, const CeedScalar *dgrad1d, const CeedScalar *dU, const int estrdU, const int cstrdU, 174 const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) { 175 MAGMA_DEVICE_SHARED(CeedScalar, shared_data) 176 177 const int tx = threadIdx.x; 178 const int ty = threadIdx.y; 179 const int elem_id = (blockIdx.x * blockDim.y) + ty; 180 181 if (elem_id >= nelem) return; 182 183 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // here DIM_U = 1, but might be different for a fused operator 184 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // here DIM_V = 1, but might be different for a fused operator 185 CeedScalar rTmp = 0.0; 186 187 // shift global memory pointers by elem stride 188 dU += elem_id * estrdU; 189 dV += elem_id * estrdV; 190 191 // assign shared memory pointers 192 CeedScalar *sTinterp = (CeedScalar *)shared_data; 193 CeedScalar *sTgrad = sTinterp + BASIS_Q * BASIS_P; 194 CeedScalar *sTmp = sTgrad + BASIS_Q * BASIS_P; 195 sTmp += ty * (max(BASIS_Q * BASIS_Q * BASIS_Q, (BASIS_Q * BASIS_Q * BASIS_P) + (BASIS_Q * BASIS_P * BASIS_P))); 196 197 // read T 198 if (ty == 0) { 199 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dinterp1d, sTinterp); 200 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dgrad1d, sTgrad); 201 } 202 __syncthreads(); 203 204 /* read U (idim = 0 for dU, i_DIM = 0 for rU) -- 205 there is a sync at the end of this function */ 206 read_U_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (0 * dstrdU), cstrdU, rU, sTmp, tx); 207 /* then first call (i_DIM = 0, i_DIM_U = 0, i_DIM_V = 0) */ 208 magma_grad_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P, 0, 0, 0, true>(sTinterp, sTgrad, rU, rV, tx, rTmp, sTmp); 209 /* there is a sync at the end of magma_grad_3d_device */ 210 211 /* read U (idim = 1 for dU, i_DIM = 0 for rU) -- 212 there is a sync at the end of this function */ 213 read_U_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (1 * dstrdU), cstrdU, rU, sTmp, tx); 214 /* then second call (i_DIM = 1, i_DIM_U = 0, i_DIM_V = 0) */ 215 magma_grad_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P, 1, 0, 0, true>(sTinterp, sTgrad, rU, rV, tx, rTmp, sTmp); 216 /* there is a sync at the end of magma_grad_3d_device */ 217 218 /* read U (idim = 2 for dU, i_DIM = 0 for rU) -- 219 there is a sync at the end of this function */ 220 read_U_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (2 * dstrdU), cstrdU, rU, sTmp, tx); 221 /* then third call (i_DIM = 2, i_DIM_U = 0, i_DIM_V = 0) */ 222 magma_grad_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P, 2, 0, 0, true>(sTinterp, sTgrad, rU, rV, tx, rTmp, sTmp); 223 /* there is a sync at the end of magma_grad_3d_device */ 224 225 // write V 226 write_V_3d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV + (0 * dstrdV), cstrdV, rV, tx); 227 } 228 229 //////////////////////////////////////////////////////////////////////////////// 230 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q *BASIS_MAX_P_Q, MAGMA_MAXTHREADS_3D)) __global__ 231 void magma_gradta_3d_kernel(const CeedScalar *dinterp1d, const CeedScalar *dgrad1d, const CeedScalar *dU, const int estrdU, const int cstrdU, 232 const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) { 233 MAGMA_DEVICE_SHARED(CeedScalar, shared_data) 234 235 const int tx = threadIdx.x; 236 const int ty = threadIdx.y; 237 const int elem_id = (blockIdx.x * blockDim.y) + ty; 238 239 if (elem_id >= nelem) return; 240 241 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // here DIM_U = 1, but might be different for a fused operator 242 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // here DIM_V = 1, but might be different for a fused operator 243 CeedScalar rTmp = 0.0; 244 245 // shift global memory pointers by elem stride 246 dU += elem_id * estrdU; 247 dV += elem_id * estrdV; 248 249 // assign shared memory pointers 250 CeedScalar *sTinterp = (CeedScalar *)shared_data; 251 CeedScalar *sTgrad = sTinterp + BASIS_Q * BASIS_P; 252 CeedScalar *sTmp = sTgrad + BASIS_Q * BASIS_P; 253 sTmp += ty * (max(BASIS_Q * BASIS_Q * BASIS_Q, (BASIS_Q * BASIS_Q * BASIS_P) + (BASIS_Q * BASIS_P * BASIS_P))); 254 255 // read T 256 if (ty == 0) { 257 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dinterp1d, sTinterp); 258 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dgrad1d, sTgrad); 259 } 260 __syncthreads(); 261 262 /* read U (idim = 0 for dU, i_DIM = 0 for rU) -- 263 there is a sync at the end of this function */ 264 read_U_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (0 * dstrdU), cstrdU, rU, sTmp, tx); 265 /* then first call (i_DIM = 0, i_DIM_U = 0, i_DIM_V = 0) */ 266 magma_grad_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P, 0, 0, 0, true>(sTinterp, sTgrad, rU, rV, tx, rTmp, sTmp); 267 /* there is a sync at the end of magma_grad_3d_device */ 268 269 /* read U (idim = 1 for dU, i_DIM = 0 for rU) -- 270 there is a sync at the end of this function */ 271 read_U_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (1 * dstrdU), cstrdU, rU, sTmp, tx); 272 /* then second call (i_DIM = 1, i_DIM_U = 0, i_DIM_V = 0) */ 273 magma_grad_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P, 1, 0, 0, true>(sTinterp, sTgrad, rU, rV, tx, rTmp, sTmp); 274 /* there is a sync at the end of magma_grad_3d_device */ 275 276 /* read U (idim = 2 for dU, i_DIM = 0 for rU) -- 277 there is a sync at the end of this function */ 278 read_U_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (2 * dstrdU), cstrdU, rU, sTmp, tx); 279 /* then third call (i_DIM = 2, i_DIM_U = 0, i_DIM_V = 0) */ 280 magma_grad_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P, 2, 0, 0, true>(sTinterp, sTgrad, rU, rV, tx, rTmp, sTmp); 281 /* there is a sync at the end of magma_grad_3d_device */ 282 283 // sum into V 284 sum_V_3d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV + (0 * dstrdV), cstrdV, rV, tx); 285 } 286