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 SYCL shared memory tensor product basis 10 #ifndef _ceed_sycl_shared_basis_tensor_h 11 #define _ceed_sycl_shared_basis_tensor_h 12 13 #include <ceed.h> 14 15 #include "sycl-shared-basis-read-write-templates.h" 16 #include "sycl-shared-basis-tensor-templates.h" 17 18 // 19 // BASIS_NUM_NODES = CeedIntPow(BASIS_P_1D,DIM) 20 // BASIS_NUM_QPTS = CeedIntPow(BASIS_Q_1D,DIM) 21 22 //------------------------------------------------------------------------------ 23 // Interp kernel by dim 24 //------------------------------------------------------------------------------ 25 kernel void Interp(const CeedInt num_elem, global const CeedScalar* restrict d_interp_1d, global const CeedScalar* restrict d_U, 26 global CeedScalar* restrict d_V) { 27 local CeedScalar s_B[BASIS_P_1D * BASIS_Q_1D]; 28 private 29 CeedScalar r_U[BASIS_NUM_COMP * (BASIS_DIM > 2 ? BASIS_P_1D : 1)]; 30 private 31 CeedScalar r_V[BASIS_NUM_COMP * (BASIS_DIM > 2 ? BASIS_Q_1D : 1)]; 32 33 local CeedScalar scratch[BASIS_INTERP_SCRATCH_SIZE]; 34 local CeedScalar* elem_scratch = scratch + get_local_id(2) * T_1D * (BASIS_DIM > 1 ? T_1D : 1); 35 36 loadMatrix(BASIS_P_1D * BASIS_Q_1D, d_interp_1d, s_B); 37 work_group_barrier(CLK_LOCAL_MEM_FENCE); 38 39 if (BASIS_DIM == 1) { 40 ReadElementStrided1d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, d_U, r_U); 41 Interp1d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, r_V, elem_scratch); 42 WriteElementStrided1d(BASIS_NUM_COMP, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, r_V, d_V); 43 44 } else if (BASIS_DIM == 2) { 45 ReadElementStrided2d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, d_U, r_U); 46 InterpTensor2d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, r_V, elem_scratch); 47 WriteElementStrided2d(BASIS_NUM_COMP, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, r_V, d_V); 48 49 } else if (BASIS_DIM == 3) { 50 ReadElementStrided3d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, d_U, r_U); 51 InterpTensor3d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, r_V, elem_scratch); 52 WriteElementStrided3d(BASIS_NUM_COMP, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, r_V, d_V); 53 } 54 } 55 56 kernel void InterpTranspose(const CeedInt num_elem, global const CeedScalar* restrict d_interp_1d, global const CeedScalar* restrict d_U, 57 global CeedScalar* restrict d_V) { 58 // local size: 59 // 1d: elems_per_block * T_1d 60 // 2d,3d: elems_per_block * T_1d * T_1d 61 local CeedScalar s_B[BASIS_P_1D * BASIS_Q_1D]; 62 private 63 CeedScalar r_U[BASIS_NUM_COMP * (BASIS_DIM > 2 ? BASIS_Q_1D : 1)]; 64 private 65 CeedScalar r_V[BASIS_NUM_COMP * (BASIS_DIM > 2 ? BASIS_P_1D : 1)]; 66 67 local CeedScalar scratch[BASIS_INTERP_SCRATCH_SIZE]; 68 local CeedScalar* elem_scratch = scratch + get_local_id(2) * T_1D * (BASIS_DIM > 1 ? T_1D : 1); 69 70 loadMatrix(BASIS_P_1D * BASIS_Q_1D, d_interp_1d, s_B); 71 work_group_barrier(CLK_LOCAL_MEM_FENCE); 72 73 if (BASIS_DIM == 1) { 74 ReadElementStrided1d(BASIS_NUM_COMP, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, d_U, r_U); 75 InterpTranspose1d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, r_V, elem_scratch); 76 WriteElementStrided1d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, r_V, d_V); 77 78 } else if (BASIS_DIM == 2) { 79 ReadElementStrided2d(BASIS_NUM_COMP, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, d_U, r_U); 80 InterpTransposeTensor2d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, r_V, elem_scratch); 81 WriteElementStrided2d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, r_V, d_V); 82 83 } else if (BASIS_DIM == 3) { 84 ReadElementStrided3d(BASIS_NUM_COMP, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, d_U, r_U); 85 InterpTransposeTensor3d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, r_V, elem_scratch); 86 WriteElementStrided3d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, r_V, d_V); 87 } 88 } 89 90 //------------------------------------------------------------------------------ 91 // Grad kernel by dim 92 //------------------------------------------------------------------------------ 93 kernel void Grad(const CeedInt num_elem, global const CeedScalar* restrict d_interp_1d, global const CeedScalar* restrict d_grad_1d, 94 global const CeedScalar* restrict d_U, global CeedScalar* restrict d_V) { 95 local CeedScalar s_B[BASIS_P_1D * BASIS_Q_1D]; // Todo, don't allocate s_B for dimension 1 96 local CeedScalar s_G[BASIS_Q_1D * (BASIS_HAS_COLLOCATED_GRAD ? BASIS_Q_1D : BASIS_P_1D)]; 97 98 private 99 CeedScalar r_U[BASIS_NUM_COMP * (BASIS_DIM > 2 ? BASIS_P_1D : 1)]; 100 private 101 CeedScalar r_V[BASIS_NUM_COMP * BASIS_DIM * (BASIS_DIM > 2 ? BASIS_Q_1D : 1)]; 102 103 local CeedScalar scratch[BASIS_GRAD_SCRATCH_SIZE]; 104 local CeedScalar* elem_scratch = scratch + get_local_id(2) * T_1D * (BASIS_DIM > 1 ? T_1D : 1); 105 106 loadMatrix(BASIS_P_1D * BASIS_Q_1D, d_interp_1d, s_B); 107 loadMatrix(BASIS_Q_1D * (BASIS_HAS_COLLOCATED_GRAD ? BASIS_Q_1D : BASIS_P_1D), d_grad_1d, s_G); 108 work_group_barrier(CLK_LOCAL_MEM_FENCE); 109 110 if (BASIS_DIM == 1) { 111 ReadElementStrided1d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, d_U, r_U); 112 Grad1d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_G, r_V, elem_scratch); 113 WriteElementStrided1d(BASIS_NUM_COMP, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, r_V, d_V); 114 115 } else if (BASIS_DIM == 2) { 116 ReadElementStrided2d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, d_U, r_U); 117 GradTensor2d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, s_G, r_V, elem_scratch); 118 WriteElementStrided2d(BASIS_NUM_COMP * BASIS_DIM, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, r_V, d_V); 119 120 } else if (BASIS_DIM == 3) { 121 ReadElementStrided3d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, d_U, r_U); 122 if (BASIS_HAS_COLLOCATED_GRAD) GradTensorCollocated3d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, s_G, r_V, elem_scratch); 123 else GradTensor3d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, s_G, r_V, elem_scratch); 124 WriteElementStrided3d(BASIS_NUM_COMP * BASIS_DIM, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, r_V, d_V); 125 } 126 } 127 128 kernel void GradTranspose(const CeedInt num_elem, global const CeedScalar* restrict d_interp_1d, global const CeedScalar* restrict d_grad_1d, 129 global const CeedScalar* restrict d_U, global CeedScalar* restrict d_V) { 130 local CeedScalar s_B[BASIS_P_1D * BASIS_Q_1D]; // Todo, don't allocate s_B for dimension 1 131 local CeedScalar s_G[BASIS_Q_1D * (BASIS_HAS_COLLOCATED_GRAD ? BASIS_Q_1D : BASIS_P_1D)]; 132 133 private 134 CeedScalar r_U[BASIS_NUM_COMP * BASIS_DIM * (BASIS_DIM > 2 ? BASIS_Q_1D : 1)]; 135 private 136 CeedScalar r_V[BASIS_NUM_COMP * (BASIS_DIM > 2 ? BASIS_P_1D : 1)]; 137 138 local CeedScalar scratch[BASIS_GRAD_SCRATCH_SIZE]; 139 local CeedScalar* elem_scratch = scratch + get_local_id(2) * T_1D * (BASIS_DIM > 1 ? T_1D : 1); 140 141 loadMatrix(BASIS_P_1D * BASIS_Q_1D, d_interp_1d, s_B); 142 loadMatrix(BASIS_Q_1D * (BASIS_HAS_COLLOCATED_GRAD ? BASIS_Q_1D : BASIS_P_1D), d_grad_1d, s_G); 143 work_group_barrier(CLK_LOCAL_MEM_FENCE); 144 145 if (BASIS_DIM == 1) { 146 ReadElementStrided1d(BASIS_NUM_COMP, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, d_U, r_U); 147 GradTranspose1d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_G, r_V, elem_scratch); 148 WriteElementStrided1d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, r_V, d_V); 149 150 } else if (BASIS_DIM == 2) { 151 ReadElementStrided2d(BASIS_NUM_COMP * BASIS_DIM, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, d_U, r_U); 152 GradTransposeTensor2d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, s_G, r_V, elem_scratch); 153 WriteElementStrided2d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, r_V, d_V); 154 155 } else if (BASIS_DIM == 3) { 156 ReadElementStrided3d(BASIS_NUM_COMP * BASIS_DIM, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, d_U, r_U); 157 if (BASIS_HAS_COLLOCATED_GRAD) GradTransposeTensorCollocated3d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, s_G, r_V, elem_scratch); 158 else GradTransposeTensor3d(BASIS_NUM_COMP, BASIS_P_1D, BASIS_Q_1D, r_U, s_B, s_G, r_V, elem_scratch); 159 WriteElementStrided3d(BASIS_NUM_COMP, BASIS_P_1D, num_elem, 1, BASIS_NUM_NODES * num_elem, BASIS_NUM_NODES, r_V, d_V); 160 } 161 } 162 163 //------------------------------------------------------------------------------ 164 // Weight kernels by dim 165 //------------------------------------------------------------------------------ 166 kernel void Weight(const CeedInt num_elem, global const CeedScalar* restrict q_weight_1d, global CeedScalar* restrict d_W) { 167 private 168 CeedScalar r_W[BASIS_DIM > 2 ? BASIS_Q_1D : 1]; 169 170 // void prefetch(q_weight_1d,BASIS_Q_1D); 171 172 if (BASIS_DIM == 1) { 173 Weight1d(BASIS_Q_1D, q_weight_1d, r_W); 174 WriteElementStrided1d(1, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, r_W, d_W); 175 176 } else if (BASIS_DIM == 2) { 177 WeightTensor2d(BASIS_Q_1D, q_weight_1d, r_W); 178 WriteElementStrided2d(1, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, r_W, d_W); 179 180 } else if (BASIS_DIM == 3) { 181 WeightTensor3d(BASIS_Q_1D, q_weight_1d, r_W); 182 WriteElementStrided3d(1, BASIS_Q_1D, num_elem, 1, BASIS_NUM_QPTS * num_elem, BASIS_NUM_QPTS, r_W, d_W); 183 } 184 } 185 186 //------------------------------------------------------------------------------ 187 188 #endif 189