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