xref: /libCEED/include/ceed/jit-source/hip/hip-ref-basis-tensor.h (revision 9e201c85545dd39529c090846df629a32c15659b)
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 #include <ceed.h>
9 
10 //------------------------------------------------------------------------------
11 // Tensor Basis Kernels
12 //------------------------------------------------------------------------------
13 
14 //------------------------------------------------------------------------------
15 // Interp
16 //------------------------------------------------------------------------------
17 extern "C" __global__ void Interp(const CeedInt num_elem, const CeedInt transpose,
18                                   const CeedScalar *__restrict__ interp_1d,
19                                   const CeedScalar *__restrict__ u,
20                                   CeedScalar *__restrict__ v) {
21   const CeedInt i = threadIdx.x;
22 
23   __shared__ CeedScalar s_mem[BASIS_Q_1D * BASIS_P_1D + 2 * BASIS_BUF_LEN];
24   CeedScalar *s_interp_1d = s_mem;
25   CeedScalar *s_buffer_1 = s_mem + BASIS_Q_1D * BASIS_P_1D;
26   CeedScalar *s_buffer_2 = s_buffer_1 + BASIS_BUF_LEN;
27   for (CeedInt k = i; k < BASIS_Q_1D * BASIS_P_1D; k += blockDim.x) {
28     s_interp_1d[k] = interp_1d[k];
29   }
30 
31   const CeedInt P = transpose ? BASIS_Q_1D : BASIS_P_1D;
32   const CeedInt Q = transpose ? BASIS_P_1D : BASIS_Q_1D;
33   const CeedInt stride0 = transpose ? 1 : BASIS_P_1D;
34   const CeedInt stride1 = transpose ? BASIS_P_1D : 1;
35   const CeedInt u_stride = transpose ? BASIS_NUM_QPTS : BASIS_NUM_NODES;
36   const CeedInt v_stride = transpose ? BASIS_NUM_NODES : BASIS_NUM_QPTS;
37   const CeedInt u_comp_stride = num_elem * (transpose ? BASIS_NUM_QPTS : BASIS_NUM_NODES);
38   const CeedInt v_comp_stride = num_elem * (transpose ? BASIS_NUM_NODES : BASIS_NUM_QPTS);
39   const CeedInt u_size = transpose ? BASIS_NUM_QPTS : BASIS_NUM_NODES;
40 
41   // Apply basis element by element
42   for (CeedInt elem = blockIdx.x; elem < num_elem; elem += gridDim.x) {
43     for (CeedInt comp = 0; comp < BASIS_NUM_COMP; comp++) {
44       const CeedScalar *cur_u = u + elem * u_stride + comp * u_comp_stride;
45       CeedScalar *cur_v = v + elem * v_stride + comp * v_comp_stride;
46       for (CeedInt k = i; k < u_size; k += blockDim.x) {
47         s_buffer_1[k] = cur_u[k];
48       }
49       CeedInt pre = u_size;
50       CeedInt post = 1;
51       for (CeedInt d = 0; d < BASIS_DIM; d++) {
52         __syncthreads();
53         // Update buffers used
54         pre /= P;
55         const CeedScalar *in = d % 2 ? s_buffer_2 : s_buffer_1;
56         CeedScalar *out = d == BASIS_DIM - 1 ? cur_v : (d % 2 ? s_buffer_1 : s_buffer_2);
57 
58         // Contract along middle index
59         const CeedInt writeLen = pre * post * Q;
60         for (CeedInt k = i; k < writeLen; k += blockDim.x) {
61           const CeedInt c = k % post;
62           const CeedInt j = (k / post) % Q;
63           const CeedInt a = k / (post * Q);
64 
65           CeedScalar vk = 0;
66           for (CeedInt b = 0; b < P; b++)
67             vk += s_interp_1d[j*stride0 + b*stride1] * in[(a*P + b)*post + c];
68 
69           out[k] = vk;
70         }
71 
72         post *= Q;
73       }
74     }
75   }
76 }
77 
78 //------------------------------------------------------------------------------
79 // Grad
80 //------------------------------------------------------------------------------
81 extern "C" __global__ void Grad(const CeedInt num_elem, const CeedInt transpose,
82                                 const CeedScalar *__restrict__ interp_1d,
83                                 const CeedScalar *__restrict__ grad_1d,
84                                 const CeedScalar *__restrict__ u,
85                                 CeedScalar *__restrict__ v) {
86   const CeedInt i = threadIdx.x;
87 
88   __shared__ CeedScalar s_mem[2 * (BASIS_Q_1D * BASIS_P_1D + BASIS_BUF_LEN)];
89   CeedScalar *s_interp_1d = s_mem;
90   CeedScalar *s_grad_1d = s_interp_1d + BASIS_Q_1D * BASIS_P_1D;
91   CeedScalar *s_buffer_1 = s_grad_1d + BASIS_Q_1D * BASIS_P_1D;
92   CeedScalar *s_buffer_2 = s_buffer_1 + BASIS_BUF_LEN;
93   for (CeedInt k = i; k < BASIS_Q_1D * BASIS_P_1D; k += blockDim.x) {
94     s_interp_1d[k] = interp_1d[k];
95     s_grad_1d[k] = grad_1d[k];
96   }
97 
98   const CeedInt P = transpose ? BASIS_Q_1D : BASIS_P_1D;
99   const CeedInt Q = transpose ? BASIS_P_1D : BASIS_Q_1D;
100   const CeedInt stride0 = transpose ? 1 : BASIS_P_1D;
101   const CeedInt stride1 = transpose ? BASIS_P_1D : 1;
102   const CeedInt u_stride = transpose ? BASIS_NUM_QPTS : BASIS_NUM_NODES;
103   const CeedInt v_stride = transpose ? BASIS_NUM_NODES : BASIS_NUM_QPTS;
104   const CeedInt u_comp_stride = num_elem * (transpose ? BASIS_NUM_QPTS : BASIS_NUM_NODES);
105   const CeedInt v_comp_stride = num_elem * (transpose ? BASIS_NUM_NODES : BASIS_NUM_QPTS);
106   const CeedInt u_dim_stride = transpose ? num_elem * BASIS_NUM_QPTS * BASIS_NUM_COMP : 0;
107   const CeedInt v_dim_stride = transpose ? 0 : num_elem * BASIS_NUM_QPTS * BASIS_NUM_COMP;
108 
109   // Apply basis element by element
110   for (CeedInt elem = blockIdx.x; elem < num_elem; elem += gridDim.x) {
111     for (CeedInt comp = 0; comp < BASIS_NUM_COMP; comp++) {
112 
113       // dim*dim contractions for grad
114       for (CeedInt dim_1 = 0; dim_1 < BASIS_DIM; dim_1++) {
115         CeedInt pre = transpose ? BASIS_NUM_QPTS : BASIS_NUM_NODES;
116         CeedInt post = 1;
117         const CeedScalar *cur_u = u + elem * u_stride + dim_1 * u_dim_stride +
118                                   comp * u_comp_stride;
119         CeedScalar *cur_v = v + elem * v_stride + dim_1 * v_dim_stride + comp *
120                             v_comp_stride;
121         for (CeedInt dim_2 = 0; dim_2 < BASIS_DIM; dim_2++) {
122           __syncthreads();
123           // Update buffers used
124           pre /= P;
125           const CeedScalar *op = dim_1 == dim_2 ? s_grad_1d : s_interp_1d;
126           const CeedScalar *in = dim_2 == 0
127                                  ? cur_u
128                                  : (dim_2 % 2 ? s_buffer_2 : s_buffer_1);
129           CeedScalar *out = dim_2 == BASIS_DIM - 1
130                             ? cur_v
131                             : (dim_2 % 2 ? s_buffer_1 : s_buffer_2);
132 
133           // Contract along middle index
134           const CeedInt writeLen = pre * post * Q;
135           for (CeedInt k = i; k < writeLen; k += blockDim.x) {
136             const CeedInt c = k % post;
137             const CeedInt j = (k / post) % Q;
138             const CeedInt a = k / (post * Q);
139             CeedScalar v_k = 0;
140             for (CeedInt b = 0; b < P; b++)
141               v_k += op[j * stride0 + b * stride1] * in[(a * P + b) * post + c];
142 
143             if (transpose && dim_2 == BASIS_DIM - 1)
144               out[k] += v_k;
145             else
146               out[k] = v_k;
147           }
148 
149           post *= Q;
150         }
151       }
152     }
153   }
154 }
155 
156 //------------------------------------------------------------------------------
157 // 1D quadrature weights
158 //------------------------------------------------------------------------------
159 __device__ void Weight1d(const CeedInt num_elem, const CeedScalar *q_weight_1d,
160                          CeedScalar *w) {
161   CeedScalar w1d[BASIS_Q_1D];
162   for (CeedInt i = 0; i < BASIS_Q_1D; i++)
163     w1d[i] = q_weight_1d[i];
164 
165   for (CeedInt e = blockIdx.x * blockDim.x + threadIdx.x;
166        e < num_elem;
167        e += blockDim.x * gridDim.x)
168     for (CeedInt i = 0; i < BASIS_Q_1D; i++) {
169       const CeedInt ind = e*BASIS_Q_1D + i; // sequential
170       w[ind] = w1d[i];
171     }
172 }
173 
174 //------------------------------------------------------------------------------
175 // 2D quadrature weights
176 //------------------------------------------------------------------------------
177 __device__ void Weight2d(const CeedInt num_elem, const CeedScalar *q_weight_1d,
178                          CeedScalar *w) {
179   CeedScalar w1d[BASIS_Q_1D];
180   for (CeedInt i = 0; i < BASIS_Q_1D; i++)
181     w1d[i] = q_weight_1d[i];
182 
183   for (CeedInt e = blockIdx.x * blockDim.x + threadIdx.x;
184        e < num_elem;
185        e += blockDim.x * gridDim.x)
186     for (CeedInt i = 0; i < BASIS_Q_1D; i++)
187       for (CeedInt j = 0; j < BASIS_Q_1D; j++) {
188         const CeedInt ind = e*BASIS_Q_1D*BASIS_Q_1D + i + j*BASIS_Q_1D; // sequential
189         w[ind] = w1d[i]*w1d[j];
190       }
191 }
192 
193 //------------------------------------------------------------------------------
194 // 3D quadrature weights
195 //------------------------------------------------------------------------------
196 __device__ void Weight3d(const CeedInt num_elem, const CeedScalar *q_weight_1d,
197                          CeedScalar *w) {
198   CeedScalar w1d[BASIS_Q_1D];
199   for (CeedInt i = 0; i < BASIS_Q_1D; i++)
200     w1d[i] = q_weight_1d[i];
201 
202   for (CeedInt e = blockIdx.x * blockDim.x + threadIdx.x;
203        e < num_elem;
204        e += blockDim.x * gridDim.x)
205     for (CeedInt i = 0; i < BASIS_Q_1D; i++)
206       for (CeedInt j = 0; j < BASIS_Q_1D; j++)
207         for (CeedInt k = 0; k < BASIS_Q_1D; k++) {
208           const CeedInt ind = e*BASIS_Q_1D*BASIS_Q_1D*BASIS_Q_1D + i +
209                               j*BASIS_Q_1D + k*BASIS_Q_1D*BASIS_Q_1D; // sequential
210           w[ind] = w1d[i]*w1d[j]*w1d[k];
211         }
212 }
213 
214 //------------------------------------------------------------------------------
215 // Quadrature weights
216 //------------------------------------------------------------------------------
217 extern "C" __global__ void Weight(const CeedInt num_elem,
218                                   const CeedScalar *__restrict__ q_weight_1d,
219                                   CeedScalar *__restrict__ v) {
220   if (BASIS_DIM == 1)
221     Weight1d(num_elem, q_weight_1d, v);
222   else if (BASIS_DIM == 2)
223     Weight2d(num_elem, q_weight_1d, v);
224   else if (BASIS_DIM == 3)
225     Weight3d(num_elem, q_weight_1d, v);
226 }
227 
228 //------------------------------------------------------------------------------
229