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 /// @file 9 /// Internal header for MAGMA tensor basis gradient in 2D 10 #ifndef CEED_MAGMA_BASIS_GRAD_2D_H 11 #define CEED_MAGMA_BASIS_GRAD_2D_H 12 13 #include "magma-common-tensor.h" 14 15 // macros to abstract access of shared memory and reg. file 16 #define sT(i, j) sT[(j)*P + (i)] 17 #define sTmp(i, j, ldw) sTmp[(j) * (ldw) + (i)] 18 19 //////////////////////////////////////////////////////////////////////////////// 20 // grad basis action (2D) 21 // This function is called two times at a higher level for 2D 22 // DIM_U -- for the size of rU[DIM_U * NUM_COMP * MAX_P_Q] 23 // DIM_V -- for the size of rV[DIM_V * NUM_COMP * MAX_P_Q] 24 // i_DIM -- the index of the outermost loop over dimensions in grad 25 // i_DIM_U -- which dim index of rU is accessed (always 0 for notrans, 0 or 1 for trans) 26 // i_DIM_V -- which dim index of rV is accessed (0 or 1 for notrans, always 0 for trans) 27 // the scalar beta is used to specify whether to accumulate to rV, or overwrite it 28 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> 29 static __device__ __inline__ void magma_grad_2d_device(const T *sTinterp, const T *sTgrad, T rU[DIM_U][NUM_COMP][rU_SIZE], 30 T rV[DIM_V][NUM_COMP][rV_SIZE], T beta, const int tx, T rTmp, T *swork) { 31 // Assumptions 32 // 0. This device routine applies grad for one dim only (i_DIM), so it should be called twice for 2D 33 // 1. 1D threads of size max(P,Q) 34 // 2. input: rU[DIM_U x NUM_COMP x P] in registers (per thread) 35 // 3. output: rV[DIM_V x NUM_COMP x Q] in registers (per thread) 36 // 4. Two products per each (dim,component) pair 37 // 4.1 Batch P of (1xP) matrices times (PxQ) matrix => Batch P of (1xQ) matrices 38 // 4.2 Batch 1 of (QxP) matrix times (PxQ) matrix => (QxQ) matrix 39 // 6. Each thread computes one row of the output of each product 40 // 7. Sync is recommended before and after the call 41 42 for (int comp = 0; comp < NUM_COMP; comp++) { 43 // 1st product -- Batch P of (1xP) matrices [reg] x (PxQ) [shmem] => Batch P of (1xQ) matrices 44 // the batch output P x (1xQ) is written on the fly to shmem 45 if (tx < P) { 46 const int batchid = tx; 47 const int sld = 1; 48 const T *sT = (i_DIM == 0) ? sTgrad : sTinterp; 49 T *sTmp = swork + batchid * (1 * Q); 50 for (int j = 0; j < Q; j++) { 51 rTmp = 0.0; 52 for (int i = 0; i < P; i++) { 53 rTmp += rU[i_DIM_U][comp][i] * sT(i, j); 54 } 55 sTmp(0, j, sld) = rTmp; 56 } 57 } // end of: if (tx < P) 58 __syncthreads(); 59 60 // 2nd product -- Batch 1 of a (QxP) matrix [shmem] x (PxQ) [shmem] => (QxQ) matrix [reg] 61 if (tx < Q) { 62 const int batchid = 0; 63 const int sld = Q; 64 const T *sT = (i_DIM == 1) ? sTgrad : sTinterp; 65 T *sTmp = swork + batchid * (Q * P); 66 for (int j = 0; j < Q; j++) { 67 rTmp = 0.0; 68 for (int i = 0; i < P; i++) { 69 rTmp += sTmp(tx, i, sld) * sT(i, j); 70 } 71 rV[i_DIM_V][comp][j] *= beta; 72 rV[i_DIM_V][comp][j] += rTmp; 73 } 74 } 75 __syncthreads(); 76 } // loop over NUM_COMP 77 } 78 79 //////////////////////////////////////////////////////////////////////////////// 80 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_2D)) __global__ 81 void magma_gradn_2d_kernel(const CeedScalar *dinterp1d, const CeedScalar *dgrad1d, const CeedScalar *dU, const int estrdU, const int cstrdU, 82 const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) { 83 MAGMA_DEVICE_SHARED(CeedScalar, shared_data) 84 85 const int tx = threadIdx.x; 86 const int ty = threadIdx.y; 87 const int elem_id = (blockIdx.x * blockDim.y) + ty; 88 magma_trans_t transT = MagmaNoTrans; 89 90 if (elem_id >= nelem) return; 91 92 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // here DIM_U = 1, but might be different for a fused operator 93 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // here DIM_V = 1, but might be different for a fused operator 94 CeedScalar rTmp = 0.0; 95 96 // shift global memory pointers by elem stride 97 dU += elem_id * estrdU; 98 dV += elem_id * estrdV; 99 100 // assign shared memory pointers 101 CeedScalar *sTinterp = (CeedScalar *)shared_data; 102 CeedScalar *sTgrad = sTinterp + BASIS_P * BASIS_Q; 103 CeedScalar *sTmp = sTgrad + BASIS_P * BASIS_Q; 104 sTmp += ty * (BASIS_P * BASIS_MAX_P_Q); 105 106 // read T 107 if (ty == 0) { 108 read_T_notrans_gm2sm<BASIS_P, BASIS_Q>(tx, dinterp1d, sTinterp); 109 read_T_notrans_gm2sm<BASIS_P, BASIS_Q>(tx, dgrad1d, sTgrad); 110 } 111 112 // No need to read V ( required only in transposed grad ) 113 const CeedScalar beta = 0.0; 114 115 /* read U (idim = 0 for dU, i_DIM = 0 for rU) -- 116 there is a sync at the end of this function */ 117 read_U_2d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dU + (0 * dstrdU), cstrdU, rU, sTmp, tx); 118 119 /* first call (i_DIM = 0, i_DIM_U = 0, i_DIM_V = 0) -- 120 output from rV[0][][] into dV (idim = 0) */ 121 magma_grad_2d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_P, BASIS_Q, BASIS_P, BASIS_Q, 0, 0, 0>(sTinterp, sTgrad, rU, rV, beta, tx, rTmp, sTmp); 122 /* there is a sync at the end of magma_grad_2d_device */ 123 write_V_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dV + (0 * dstrdV), cstrdV, rV, tx); 124 125 /* second call (i_DIM = 1, i_DIM_U = 0, i_DIM_V = 0) -- 126 output from rV[0][][] into dV (idim = 1) */ 127 magma_grad_2d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_P, BASIS_Q, BASIS_P, BASIS_Q, 1, 0, 0>(sTinterp, sTgrad, rU, rV, beta, tx, rTmp, sTmp); 128 /* there is a sync at the end of magma_grad_2d_device */ 129 write_V_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dV + (1 * dstrdV), cstrdV, rV, tx); 130 } 131 132 //////////////////////////////////////////////////////////////////////////////// 133 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_2D)) __global__ 134 void magma_gradt_2d_kernel(const CeedScalar *dinterp1d, const CeedScalar *dgrad1d, const CeedScalar *dU, const int estrdU, const int cstrdU, 135 const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) { 136 MAGMA_DEVICE_SHARED(CeedScalar, shared_data) 137 138 const int tx = threadIdx.x; 139 const int ty = threadIdx.y; 140 const int elem_id = (blockIdx.x * blockDim.y) + ty; 141 magma_trans_t transT = MagmaTrans; 142 143 if (elem_id >= nelem) return; 144 145 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // here DIM_U = 1, but might be different for a fused operator 146 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // here DIM_V = 1, but might be different for a fused operator 147 CeedScalar rTmp = 0.0; 148 149 // shift global memory pointers by elem stride 150 dU += elem_id * estrdU; 151 dV += elem_id * estrdV; 152 153 // assign shared memory pointers 154 CeedScalar *sTinterp = (CeedScalar *)shared_data; 155 CeedScalar *sTgrad = sTinterp + BASIS_Q * BASIS_P; 156 CeedScalar *sTmp = sTgrad + BASIS_Q * BASIS_P; 157 sTmp += ty * (BASIS_Q * BASIS_MAX_P_Q); 158 159 // read T 160 if (ty == 0) { 161 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dinterp1d, sTinterp); 162 read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dgrad1d, sTgrad); 163 } 164 __syncthreads(); 165 166 /* read V (since this is transposed mode -- 167 idim = 0 for dV, i_DIM = 0 for rV) */ 168 const CeedScalar beta = 1.0; 169 read_V_2d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV + (0 * dstrdV), cstrdV, rV, tx); 170 171 /* read U (idim = 0 for dU, i_DIM = 0 for rU) -- 172 there is a sync at the end of this function */ 173 read_U_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (0 * dstrdU), cstrdU, rU, sTmp, tx); 174 /* first call (i_DIM = 0, i_DIM_U = 0, i_DIM_V = 0) */ 175 magma_grad_2d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P, 0, 0, 0>(sTinterp, sTgrad, rU, rV, beta, tx, rTmp, sTmp); 176 /* there is a sync at the end of magma_grad_2d_device */ 177 178 /* read U (idim = 1 for dU, i_DIM = 0 for rU) -- 179 there is a sync at the end of this function */ 180 read_U_2d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU + (1 * dstrdU), cstrdU, rU, sTmp, tx); 181 /* second call (i_DIM = 1, i_DIM_U = 0, i_DIM_V = 0) */ 182 magma_grad_2d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P, 1, 0, 0>(sTinterp, sTgrad, rU, rV, beta, tx, rTmp, sTmp); 183 /* there is a sync at the end of magma_grad_2d_device */ 184 185 // write V 186 write_V_2d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV + (0 * dstrdV), cstrdV, rV, tx); 187 } 188 189 #endif // CEED_MAGMA_BASIS_GRAD_2D_H 190