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