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