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 2D 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 17 //////////////////////////////////////////////////////////////////////////////// 18 // Helper function to add or set into V 19 template <typename T, bool Add> 20 struct magma_grad_2d_device_accumulate; 21 22 template <typename T> 23 struct magma_grad_2d_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_2d_device_accumulate<T, false> { 29 static __device__ __inline__ void op(T &rV, const T &rTmp) { rV = rTmp; } 30 }; 31 32 //////////////////////////////////////////////////////////////////////////////// 33 // grad basis action (2D) 34 // This function is called two times at a higher level for 2D 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 or 1 for trans) 39 // i_DIM_V -- which dim index of rV is accessed (0 or 1 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_2d_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 called twice for 2D 45 // 1. 1D threads of size max(P,Q) 46 // 2. input: rU[DIM_U x NUM_COMP x P] in registers (per thread) 47 // 3. output: rV[DIM_V x NUM_COMP x Q] in registers (per thread) 48 // 4. Two products per each (dim,component) pair 49 // 4.1 Batch P of (1xP) matrices times (PxQ) matrix => Batch P of (1xQ) matrices 50 // 4.2 Batch 1 of (QxP) matrix times (PxQ) matrix => (QxQ) matrix 51 // 6. Each thread computes one row of the output of each product 52 // 7. Sync is recommended before and after the call 53 54 for (int comp = 0; comp < NUM_COMP; comp++) { 55 // 1st product -- Batch P of (1xP) matrices [reg] x (PxQ) [shmem] => Batch P of (1xQ) matrices 56 // the batch output P x (1xQ) is written on the fly to shmem 57 if (tx < P) { 58 const int batchid = tx; 59 const int sld = 1; 60 const T *sT = (i_DIM == 0) ? sTgrad : sTinterp; 61 T *sTmp = swork + batchid * (1 * Q); 62 for (int j = 0; j < Q; j++) { 63 rTmp = 0.0; 64 for (int i = 0; i < P; i++) { 65 rTmp += rU[i_DIM_U][comp][i] * sT(i, j); 66 } 67 sTmp(0, j, sld) = rTmp; 68 } 69 } // end of: if (tx < P) 70 __syncthreads(); 71 72 // 2nd product -- Batch 1 of a (QxP) matrix [shmem] x (PxQ) [shmem] => (QxQ) matrix [reg] 73 if (tx < Q) { 74 const int batchid = 0; 75 const int sld = Q; 76 const T *sT = (i_DIM == 1) ? sTgrad : sTinterp; 77 T *sTmp = swork + batchid * (Q * P); 78 for (int j = 0; j < Q; j++) { 79 rTmp = 0.0; 80 for (int i = 0; i < P; i++) { 81 rTmp += sTmp(tx, i, sld) * sT(i, j); 82 } 83 magma_grad_2d_device_accumulate<T, ADD>::op(rV[i_DIM_V][comp][j], rTmp); 84 } 85 } 86 __syncthreads(); 87 } // loop over NUM_COMP 88 } 89 90 //////////////////////////////////////////////////////////////////////////////// 91 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_2D)) __global__ 92 void magma_gradn_2d_kernel(const CeedScalar *dinterp1d, const CeedScalar *dgrad1d, const CeedScalar *dU, const int estrdU, const int cstrdU, 93 const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) { 94 MAGMA_DEVICE_SHARED(CeedScalar, shared_data) 95 96 const int tx = threadIdx.x; 97 const int ty = threadIdx.y; 98 const int elem_id = (blockIdx.x * blockDim.y) + ty; 99 100 if (elem_id >= nelem) return; 101 102 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // here DIM_U = 1, but might be different for a fused operator 103 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // here DIM_V = 1, but might be different for a fused operator 104 CeedScalar rTmp = 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 *sTinterp = (CeedScalar *)shared_data; 112 CeedScalar *sTgrad = sTinterp + BASIS_P * BASIS_Q; 113 CeedScalar *sTmp = sTgrad + BASIS_P * BASIS_Q; 114 sTmp += ty * (BASIS_P * BASIS_MAX_P_Q); 115 116 // read T 117 if (ty == 0) { 118 read_T_notrans_gm2sm<BASIS_P, BASIS_Q>(tx, dinterp1d, sTinterp); 119 read_T_notrans_gm2sm<BASIS_P, BASIS_Q>(tx, dgrad1d, sTgrad); 120 } 121 122 /* read U (idim = 0 for dU, i_DIM = 0 for rU) -- 123 there is a sync at the end of this function */ 124 read_U_2d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dU + (0 * dstrdU), cstrdU, rU, sTmp, tx); 125 126 /* first call (i_DIM = 0, i_DIM_U = 0, i_DIM_V = 0) -- 127 output from rV[0][][] into dV (idim = 0) */ 128 magma_grad_2d_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, 129 sTmp); 130 /* there is a sync at the end of magma_grad_2d_device */ 131 write_V_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dV + (0 * dstrdV), cstrdV, rV, tx); 132 133 /* second call (i_DIM = 1, i_DIM_U = 0, i_DIM_V = 0) -- 134 output from rV[0][][] into dV (idim = 1) */ 135 magma_grad_2d_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, 136 sTmp); 137 /* there is a sync at the end of magma_grad_2d_device */ 138 write_V_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dV + (1 * dstrdV), cstrdV, rV, tx); 139 } 140 141 //////////////////////////////////////////////////////////////////////////////// 142 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_2D)) __global__ 143 void magma_gradt_2d_kernel(const CeedScalar *dinterp1d, const CeedScalar *dgrad1d, const CeedScalar *dU, const int estrdU, const int cstrdU, 144 const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) { 145 MAGMA_DEVICE_SHARED(CeedScalar, shared_data) 146 147 const int tx = threadIdx.x; 148 const int ty = threadIdx.y; 149 const int elem_id = (blockIdx.x * blockDim.y) + ty; 150 151 if (elem_id >= nelem) return; 152 153 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // here DIM_U = 1, but might be different for a fused operator 154 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // here DIM_V = 1, but might be different for a fused operator 155 CeedScalar rTmp = 0.0; 156 157 // shift global memory pointers by elem stride 158 dU += elem_id * estrdU; 159 dV += elem_id * estrdV; 160 161 // assign shared memory pointers 162 CeedScalar *sTinterp = (CeedScalar *)shared_data; 163 CeedScalar *sTgrad = sTinterp + BASIS_Q * BASIS_P; 164 CeedScalar *sTmp = sTgrad + BASIS_Q * BASIS_P; 165 sTmp += ty * (BASIS_Q * BASIS_MAX_P_Q); 166 167 // read T 168 if (ty == 0) { 169 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dinterp1d, sTinterp); 170 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dgrad1d, sTgrad); 171 } 172 __syncthreads(); 173 174 /* read U (idim = 0 for dU, i_DIM = 0 for rU) -- 175 there is a sync at the end of this function */ 176 read_U_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (0 * dstrdU), cstrdU, rU, sTmp, tx); 177 /* first call (i_DIM = 0, i_DIM_U = 0, i_DIM_V = 0) */ 178 magma_grad_2d_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); 179 /* there is a sync at the end of magma_grad_2d_device */ 180 181 /* read U (idim = 1 for dU, i_DIM = 0 for rU) -- 182 there is a sync at the end of this function */ 183 read_U_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (1 * dstrdU), cstrdU, rU, sTmp, tx); 184 /* second call (i_DIM = 1, i_DIM_U = 0, i_DIM_V = 0) */ 185 magma_grad_2d_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); 186 /* there is a sync at the end of magma_grad_2d_device */ 187 188 // write V 189 write_V_2d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV + (0 * dstrdV), cstrdV, rV, tx); 190 } 191 192 //////////////////////////////////////////////////////////////////////////////// 193 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_2D)) __global__ 194 void magma_gradta_2d_kernel(const CeedScalar *dinterp1d, const CeedScalar *dgrad1d, const CeedScalar *dU, const int estrdU, const int cstrdU, 195 const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) { 196 MAGMA_DEVICE_SHARED(CeedScalar, shared_data) 197 198 const int tx = threadIdx.x; 199 const int ty = threadIdx.y; 200 const int elem_id = (blockIdx.x * blockDim.y) + ty; 201 202 if (elem_id >= nelem) return; 203 204 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // here DIM_U = 1, but might be different for a fused operator 205 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // here DIM_V = 1, but might be different for a fused operator 206 CeedScalar rTmp = 0.0; 207 208 // shift global memory pointers by elem stride 209 dU += elem_id * estrdU; 210 dV += elem_id * estrdV; 211 212 // assign shared memory pointers 213 CeedScalar *sTinterp = (CeedScalar *)shared_data; 214 CeedScalar *sTgrad = sTinterp + BASIS_Q * BASIS_P; 215 CeedScalar *sTmp = sTgrad + BASIS_Q * BASIS_P; 216 sTmp += ty * (BASIS_Q * BASIS_MAX_P_Q); 217 218 // read T 219 if (ty == 0) { 220 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dinterp1d, sTinterp); 221 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dgrad1d, sTgrad); 222 } 223 __syncthreads(); 224 225 /* read U (idim = 0 for dU, i_DIM = 0 for rU) -- 226 there is a sync at the end of this function */ 227 read_U_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (0 * dstrdU), cstrdU, rU, sTmp, tx); 228 /* first call (i_DIM = 0, i_DIM_U = 0, i_DIM_V = 0) */ 229 magma_grad_2d_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); 230 /* there is a sync at the end of magma_grad_2d_device */ 231 232 /* read U (idim = 1 for dU, i_DIM = 0 for rU) -- 233 there is a sync at the end of this function */ 234 read_U_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (1 * dstrdU), cstrdU, rU, sTmp, tx); 235 /* second call (i_DIM = 1, i_DIM_U = 0, i_DIM_V = 0) */ 236 magma_grad_2d_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); 237 /* there is a sync at the end of magma_grad_2d_device */ 238 239 // sum into V 240 sum_V_2d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV + (0 * dstrdV), cstrdV, rV, tx); 241 } 242