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 interpolation in 3D 10 #ifndef CEED_MAGMA_BASIS_INTERP_3D_H 11 #define CEED_MAGMA_BASIS_INTERP_3D_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 // interp basis action (3D) 21 template <typename T, int DIM_U, int DIM_V, int NUM_COMP, int P, int Q, int rU_SIZE, int rV_SIZE> 22 static __device__ __inline__ void magma_interp_3d_device(const T *sT, magma_trans_t transT, T rU[DIM_U][NUM_COMP][rU_SIZE], 23 T rV[DIM_V][NUM_COMP][rV_SIZE], const int tx, T rTmp[Q], T *swork) { 24 // Assumptions 25 // 1. 1D threads of size max(P,Q)^2 26 // 2. input: rU[DIM_U x NUM_COMP x rU_SIZE] in registers (per thread) 27 // 3. output: rV[DIM_V x NUM_COMP x rV_SIZE] in registers (per thread) 28 // 4. Three products per component 29 // 4.1 Batch P^2 of (1xP) matrices times (PxQ) matrix => Batch P^2 of (1xQ) matrices 30 // 4.2 Batch P of (QxP) matrices times (PxQ) matrix => Batch P of (QxQ) matrices 31 // 4.3 Batch 1 of (Q^2xP_) matrix times (PxQ) matrix => (Q^2xQ_) matrix 32 // 5. Each thread computes one row of the output of each product 33 // 6. Sync is recommended before and after the call 34 35 for (int comp = 0; comp < NUM_COMP; comp++) { 36 // Batch P^2 of (1xP) matrices [reg] times (PxQ) matrix [shmem] => Batch P^2 of (1xQ) matrices [shmem] 37 if (tx < (P * P)) { 38 const int batchid = tx; 39 const int sld = 1; 40 T *sTmp = swork + batchid * (1 * Q); 41 for (int j = 0; j < Q; j++) { 42 rTmp[0] = 0.0; 43 for (int i = 0; i < P; i++) { 44 rTmp[0] += rU[0][comp][i] * sT(i, j); 45 } 46 sTmp(0, j, sld) = rTmp[0]; 47 } 48 } // end of: if (tx < P*P) 49 __syncthreads(); 50 51 // Batch P of (QxP) matrices [shmem] times (PxQ) matrix [shmem] => Batch P of (QxQ) matrices [reg] 52 if (tx < (P * Q)) { 53 const int batchid = tx / Q; 54 const int tx_ = tx % Q; 55 const int sld = Q; 56 T *sTmp = swork + batchid * (Q * P); // sTmp is input 57 for (int j = 0; j < Q; j++) { 58 rTmp[j] = 0.0; 59 for (int i = 0; i < P; i++) { 60 rTmp[j] += sTmp(tx_, i, sld) * sT(i, j); 61 } 62 } 63 } 64 __syncthreads(); 65 66 // write rTmp[] into shmem as batch P of QxQ matrices 67 if (tx < (P * Q)) { 68 const int batchid = tx / Q; 69 const int tx_ = tx % Q; 70 const int sld = Q; 71 T *sTmp = swork + batchid * (Q * Q); 72 for (int j = 0; j < Q; j++) { 73 sTmp(tx_, j, sld) = rTmp[j]; 74 } 75 } 76 __syncthreads(); 77 78 // Batch 1 of (Q^2xP_) matrices [shmem] times (PxQ) matrix [shmem] => Batch 1 of (Q^2xQ_) matrices [reg] 79 if (tx < (Q * Q)) { 80 // No need to declare batchid = (tx / Q^2) = always zero 81 // No need to declare tx_ = (tx_ % Q^2) = always tx 82 const int sld = Q * Q; 83 T *sTmp = swork; 84 for (int j = 0; j < Q; j++) { 85 rTmp[0] = 0.0; 86 for (int i = 0; i < P; i++) { 87 rTmp[0] += sTmp(tx, i, sld) * sT(i, j); 88 } 89 rV[0][comp][j] += rTmp[0]; 90 } 91 } 92 __syncthreads(); 93 } 94 } 95 96 ////////////////////////////////////////////////////////////////////////////////////////// 97 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q *BASIS_MAX_P_Q, MAGMA_MAXTHREADS_3D)) __global__ 98 void magma_interpn_3d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV, 99 const int cstrdV, const int nelem) { 100 MAGMA_DEVICE_SHARED(CeedScalar, shared_data) 101 102 const int tx = threadIdx.x; 103 const int ty = threadIdx.y; 104 const int elem_id = (blockIdx.x * blockDim.y) + ty; 105 magma_trans_t transT = MagmaNoTrans; 106 107 if (elem_id >= nelem) return; 108 109 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // for a non-fused operator BASIS_DIM is always 1 110 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // for a non-fused operator BASIS_DIM is always 1 111 CeedScalar rTmp[BASIS_Q] = {0.0}; 112 113 // shift global memory pointers by elem stride 114 dU += elem_id * estrdU; 115 dV += elem_id * estrdV; 116 117 // assign shared memory pointers 118 CeedScalar *sT = (CeedScalar *)shared_data; 119 CeedScalar *sTmp = sT + BASIS_P * BASIS_Q; 120 sTmp += ty * (max(BASIS_P * BASIS_P * BASIS_MAX_P_Q, BASIS_P * BASIS_Q * BASIS_Q)); 121 122 // read T 123 if (ty == 0) { 124 dread_T_gm2sm<BASIS_P, BASIS_Q>(tx, transT, dT, sT); 125 } 126 127 // read U (idim = 0 for dU, i_DIM = 0 for rU, u_dimstride is always 0) 128 readU_3d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dU, cstrdU, rU, sTmp, tx); 129 // there is a sync at the end of this function 130 131 magma_interp_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_P, BASIS_Q, BASIS_P, BASIS_Q>(sT, transT, rU, rV, tx, rTmp, sTmp); 132 __syncthreads(); 133 134 // write V 135 writeV_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dV, cstrdV, rV, tx); 136 } 137 138 ////////////////////////////////////////////////////////////////////////////////////////// 139 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q *BASIS_MAX_P_Q, MAGMA_MAXTHREADS_3D)) __global__ 140 void magma_interpt_3d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV, 141 const int cstrdV, const int nelem) { 142 MAGMA_DEVICE_SHARED(CeedScalar, shared_data) 143 144 const int tx = threadIdx.x; 145 const int ty = threadIdx.y; 146 const int elem_id = (blockIdx.x * blockDim.y) + ty; 147 magma_trans_t transT = MagmaTrans; 148 149 if (elem_id >= nelem) return; 150 151 CeedScalar rU[1][BASIS_NUM_COMP][BASIS_Q] = {0.0}; // for a non-fused operator BASIS_DIM is always 1 152 CeedScalar rV[1][BASIS_NUM_COMP][BASIS_P] = {0.0}; // for a non-fused operator BASIS_DIM is always 1 153 CeedScalar rTmp[BASIS_P] = {0.0}; 154 155 // shift global memory pointers by elem stride 156 dU += elem_id * estrdU; 157 dV += elem_id * estrdV; 158 159 // assign shared memory pointers 160 CeedScalar *sT = (CeedScalar *)shared_data; 161 CeedScalar *sTmp = sT + BASIS_Q * BASIS_P; 162 sTmp += ty * (max(BASIS_Q * BASIS_Q * BASIS_MAX_P_Q, BASIS_Q * BASIS_P * BASIS_P)); 163 164 // read T 165 if (ty == 0) { 166 dread_T_gm2sm<BASIS_Q, BASIS_P>(tx, transT, dT, sT); 167 } 168 169 // read V 170 readV_3d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV, cstrdV, rV, tx); 171 172 // read U (idim = 0 for dU, i_DIM = 0 for rU, u_dimstride is always 0) 173 readU_3d<CeedScalar, BASIS_Q, 1, BASIS_NUM_COMP, BASIS_Q, 0>(dU, cstrdU, rU, sTmp, tx); 174 // there is a sync at the end of this function 175 176 magma_interp_3d_device<CeedScalar, 1, 1, BASIS_NUM_COMP, BASIS_Q, BASIS_P, BASIS_Q, BASIS_P>(sT, transT, rU, rV, tx, rTmp, sTmp); 177 __syncthreads(); 178 179 // write V 180 writeV_3d<CeedScalar, BASIS_P, 1, BASIS_NUM_COMP, BASIS_P, 0>(dV, cstrdV, rV, tx); 181 } 182 183 #endif // CEED_MAGMA_BASIS_INTERP_3D_H 184