xref: /libCEED/rust/libceed-sys/c-src/include/ceed/jit-source/hip/hip-gen-templates.h (revision 688b547317fb91b3911fc39af1ff501a468744ac)
1 // Copyright (c) 2017-2024, 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 HIP backend macro and type definitions for JiT source
10 #include <ceed/types.h>
11 
12 //------------------------------------------------------------------------------
13 // Load matrices for basis actions
14 //------------------------------------------------------------------------------
15 template <int P, int Q>
16 inline __device__ void LoadMatrix(SharedData_Hip &data, const CeedScalar *__restrict__ d_B, CeedScalar *B) {
17   for (CeedInt i = data.t_id; i < P * Q; i += blockDim.x * blockDim.y * blockDim.z) B[i] = d_B[i];
18 }
19 
20 //------------------------------------------------------------------------------
21 // 1D
22 //------------------------------------------------------------------------------
23 
24 //------------------------------------------------------------------------------
25 // L-vector -> E-vector, offsets provided
26 //------------------------------------------------------------------------------
27 template <int NUM_COMP, int COMP_STRIDE, int P_1d>
28 inline __device__ void ReadLVecStandard1d(SharedData_Hip &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
29                                           const CeedScalar *__restrict__ d_u, CeedScalar *__restrict__ r_u) {
30   if (data.t_id_x < P_1d) {
31     const CeedInt node = data.t_id_x;
32     const CeedInt ind  = indices[node + elem * P_1d];
33 
34     for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + COMP_STRIDE * comp];
35   }
36 }
37 
38 //------------------------------------------------------------------------------
39 // L-vector -> E-vector, strided
40 //------------------------------------------------------------------------------
41 template <int NUM_COMP, int P_1d, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
42 inline __device__ void ReadLVecStrided1d(SharedData_Hip &data, const CeedInt elem, const CeedScalar *__restrict__ d_u, CeedScalar *__restrict__ r_u) {
43   if (data.t_id_x < P_1d) {
44     const CeedInt node = data.t_id_x;
45     const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;
46 
47     for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + comp * STRIDES_COMP];
48   }
49 }
50 
51 //------------------------------------------------------------------------------
52 // E-vector -> L-vector, offsets provided
53 //------------------------------------------------------------------------------
54 template <int NUM_COMP, int COMP_STRIDE, int P_1d>
55 inline __device__ void WriteLVecStandard1d(SharedData_Hip &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
56                                            const CeedScalar *__restrict__ r_v, CeedScalar *__restrict__ d_v) {
57   if (data.t_id_x < P_1d) {
58     const CeedInt node = data.t_id_x;
59     const CeedInt ind  = indices[node + elem * P_1d];
60 
61     for (CeedInt comp = 0; comp < NUM_COMP; comp++) atomicAdd(&d_v[ind + COMP_STRIDE * comp], r_v[comp]);
62   }
63 }
64 
65 //------------------------------------------------------------------------------
66 // E-vector -> L-vector, strided
67 //------------------------------------------------------------------------------
68 template <int NUM_COMP, int P_1d, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
69 inline __device__ void WriteLVecStrided1d(SharedData_Hip &data, const CeedInt elem, const CeedScalar *__restrict__ r_v,
70                                           CeedScalar *__restrict__ d_v) {
71   if (data.t_id_x < P_1d) {
72     const CeedInt node = data.t_id_x;
73     const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;
74 
75     for (CeedInt comp = 0; comp < NUM_COMP; comp++) d_v[ind + comp * STRIDES_COMP] += r_v[comp];
76   }
77 }
78 
79 //------------------------------------------------------------------------------
80 // 2D
81 //------------------------------------------------------------------------------
82 
83 //------------------------------------------------------------------------------
84 // L-vector -> E-vector, offsets provided
85 //------------------------------------------------------------------------------
86 template <int NUM_COMP, int COMP_STRIDE, int P_1d>
87 inline __device__ void ReadLVecStandard2d(SharedData_Hip &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
88                                           const CeedScalar *__restrict__ d_u, CeedScalar *__restrict__ r_u) {
89   if (data.t_id_x < P_1d && data.t_id_y < P_1d) {
90     const CeedInt node = data.t_id_x + data.t_id_y * P_1d;
91     const CeedInt ind  = indices[node + elem * P_1d * P_1d];
92 
93     for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + COMP_STRIDE * comp];
94   }
95 }
96 
97 //------------------------------------------------------------------------------
98 // L-vector -> E-vector, strided
99 //------------------------------------------------------------------------------
100 template <int NUM_COMP, int P_1d, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
101 inline __device__ void ReadLVecStrided2d(SharedData_Hip &data, const CeedInt elem, const CeedScalar *__restrict__ d_u, CeedScalar *__restrict__ r_u) {
102   if (data.t_id_x < P_1d && data.t_id_y < P_1d) {
103     const CeedInt node = data.t_id_x + data.t_id_y * P_1d;
104     const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;
105 
106     for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + comp * STRIDES_COMP];
107   }
108 }
109 
110 //------------------------------------------------------------------------------
111 // E-vector -> L-vector, offsets provided
112 //------------------------------------------------------------------------------
113 template <int NUM_COMP, int COMP_STRIDE, int P_1d>
114 inline __device__ void WriteLVecStandard2d(SharedData_Hip &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
115                                            const CeedScalar *__restrict__ r_v, CeedScalar *__restrict__ d_v) {
116   if (data.t_id_x < P_1d && data.t_id_y < P_1d) {
117     const CeedInt node = data.t_id_x + data.t_id_y * P_1d;
118     const CeedInt ind  = indices[node + elem * P_1d * P_1d];
119 
120     for (CeedInt comp = 0; comp < NUM_COMP; comp++) atomicAdd(&d_v[ind + COMP_STRIDE * comp], r_v[comp]);
121   }
122 }
123 
124 //------------------------------------------------------------------------------
125 // E-vector -> L-vector, strided
126 //------------------------------------------------------------------------------
127 template <int NUM_COMP, int P_1d, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
128 inline __device__ void WriteLVecStrided2d(SharedData_Hip &data, const CeedInt elem, const CeedScalar *__restrict__ r_v,
129                                           CeedScalar *__restrict__ d_v) {
130   if (data.t_id_x < P_1d && data.t_id_y < P_1d) {
131     const CeedInt node = data.t_id_x + data.t_id_y * P_1d;
132     const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;
133 
134     for (CeedInt comp = 0; comp < NUM_COMP; comp++) d_v[ind + comp * STRIDES_COMP] += r_v[comp];
135   }
136 }
137 
138 //------------------------------------------------------------------------------
139 // 3D
140 //------------------------------------------------------------------------------
141 
142 //------------------------------------------------------------------------------
143 // L-vector -> E-vector, offsets provided
144 //------------------------------------------------------------------------------
145 template <int NUM_COMP, int COMP_STRIDE, int P_1d>
146 inline __device__ void ReadLVecStandard3d(SharedData_Hip &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
147                                           const CeedScalar *__restrict__ d_u, CeedScalar *__restrict__ r_u) {
148   if (data.t_id_x < P_1d && data.t_id_y < P_1d) {
149     for (CeedInt z = 0; z < P_1d; z++) {
150       const CeedInt node = data.t_id_x + data.t_id_y * P_1d + z * P_1d * P_1d;
151       const CeedInt ind  = indices[node + elem * P_1d * P_1d * P_1d];
152 
153       for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[z + comp * P_1d] = d_u[ind + COMP_STRIDE * comp];
154     }
155   }
156 }
157 
158 //------------------------------------------------------------------------------
159 // L-vector -> E-vector, strided
160 //------------------------------------------------------------------------------
161 template <int NUM_COMP, int P_1d, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
162 inline __device__ void ReadLVecStrided3d(SharedData_Hip &data, const CeedInt elem, const CeedScalar *__restrict__ d_u, CeedScalar *__restrict__ r_u) {
163   if (data.t_id_x < P_1d && data.t_id_y < P_1d) {
164     for (CeedInt z = 0; z < P_1d; z++) {
165       const CeedInt node = data.t_id_x + data.t_id_y * P_1d + z * P_1d * P_1d;
166       const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;
167 
168       for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[z + comp * P_1d] = d_u[ind + comp * STRIDES_COMP];
169     }
170   }
171 }
172 
173 //------------------------------------------------------------------------------
174 // E-vector -> Q-vector, offests provided
175 //------------------------------------------------------------------------------
176 template <int NUM_COMP, int COMP_STRIDE, int Q_1d>
177 inline __device__ void ReadEVecSliceStandard3d(SharedData_Hip &data, const CeedInt nquads, const CeedInt elem, const CeedInt q,
178                                                const CeedInt *__restrict__ indices, const CeedScalar *__restrict__ d_u,
179                                                CeedScalar *__restrict__ r_u) {
180   if (data.t_id_x < Q_1d && data.t_id_y < Q_1d) {
181     const CeedInt node = data.t_id_x + data.t_id_y * Q_1d + q * Q_1d * Q_1d;
182     const CeedInt ind  = indices[node + elem * Q_1d * Q_1d * Q_1d];
183 
184     for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + COMP_STRIDE * comp];
185   }
186 }
187 
188 //------------------------------------------------------------------------------
189 // E-vector -> Q-vector, strided
190 //------------------------------------------------------------------------------
191 template <int NUM_COMP, int Q_1d, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
192 inline __device__ void ReadEVecSliceStrided3d(SharedData_Hip &data, const CeedInt elem, const CeedInt q, const CeedScalar *__restrict__ d_u,
193                                               CeedScalar *__restrict__ r_u) {
194   if (data.t_id_x < Q_1d && data.t_id_y < Q_1d) {
195     const CeedInt node = data.t_id_x + data.t_id_y * Q_1d + q * Q_1d * Q_1d;
196     const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;
197 
198     for (CeedInt comp = 0; comp < NUM_COMP; comp++) r_u[comp] = d_u[ind + comp * STRIDES_COMP];
199   }
200 }
201 
202 //------------------------------------------------------------------------------
203 // E-vector -> L-vector, offsets provided
204 //------------------------------------------------------------------------------
205 template <int NUM_COMP, int COMP_STRIDE, int P_1d>
206 inline __device__ void WriteLVecStandard3d(SharedData_Hip &data, const CeedInt num_nodes, const CeedInt elem, const CeedInt *__restrict__ indices,
207                                            const CeedScalar *__restrict__ r_v, CeedScalar *__restrict__ d_v) {
208   if (data.t_id_x < P_1d && data.t_id_y < P_1d) {
209     for (CeedInt z = 0; z < P_1d; z++) {
210       const CeedInt node = data.t_id_x + data.t_id_y * P_1d + z * P_1d * P_1d;
211       const CeedInt ind  = indices[node + elem * P_1d * P_1d * P_1d];
212 
213       for (CeedInt comp = 0; comp < NUM_COMP; comp++) atomicAdd(&d_v[ind + COMP_STRIDE * comp], r_v[z + comp * P_1d]);
214     }
215   }
216 }
217 
218 //------------------------------------------------------------------------------
219 // E-vector -> L-vector, strided
220 //------------------------------------------------------------------------------
221 template <int NUM_COMP, int P_1d, int STRIDES_NODE, int STRIDES_COMP, int STRIDES_ELEM>
222 inline __device__ void WriteLVecStrided3d(SharedData_Hip &data, const CeedInt elem, const CeedScalar *__restrict__ r_v,
223                                           CeedScalar *__restrict__ d_v) {
224   if (data.t_id_x < P_1d && data.t_id_y < P_1d) {
225     for (CeedInt z = 0; z < P_1d; z++) {
226       const CeedInt node = data.t_id_x + data.t_id_y * P_1d + z * P_1d * P_1d;
227       const CeedInt ind  = node * STRIDES_NODE + elem * STRIDES_ELEM;
228 
229       for (CeedInt comp = 0; comp < NUM_COMP; comp++) d_v[ind + comp * STRIDES_COMP] += r_v[z + comp * P_1d];
230     }
231   }
232 }
233 
234 //------------------------------------------------------------------------------
235 // 3D collocated derivatives computation
236 //------------------------------------------------------------------------------
237 template <int NUM_COMP, int Q_1d>
238 inline __device__ void GradColloSlice3d(SharedData_Hip &data, const CeedInt q, const CeedScalar *__restrict__ r_U, const CeedScalar *c_G,
239                                         CeedScalar *__restrict__ r_V) {
240   if (data.t_id_x < Q_1d && data.t_id_y < Q_1d) {
241     for (CeedInt comp = 0; comp < NUM_COMP; comp++) {
242       data.slice[data.t_id_x + data.t_id_y * T_1D] = r_U[q + comp * Q_1d];
243       __syncthreads();
244       // X derivative
245       r_V[comp + 0 * NUM_COMP] = 0.0;
246       for (CeedInt i = 0; i < Q_1d; i++) {
247         r_V[comp + 0 * NUM_COMP] += c_G[i + data.t_id_x * Q_1d] * data.slice[i + data.t_id_y * T_1D];
248       }
249       // Y derivative
250       r_V[comp + 1 * NUM_COMP] = 0.0;
251       for (CeedInt i = 0; i < Q_1d; i++) {
252         r_V[comp + 1 * NUM_COMP] += c_G[i + data.t_id_y * Q_1d] * data.slice[data.t_id_x + i * T_1D];
253       }
254       // Z derivative
255       r_V[comp + 2 * NUM_COMP] = 0.0;
256       for (CeedInt i = 0; i < Q_1d; i++) {
257         r_V[comp + 2 * NUM_COMP] += c_G[i + q * Q_1d] * r_U[i + comp * Q_1d];
258       }
259       __syncthreads();
260     }
261   }
262 }
263 
264 //------------------------------------------------------------------------------
265 // 3D collocated derivatives transpose
266 //------------------------------------------------------------------------------
267 template <int NUM_COMP, int Q_1d>
268 inline __device__ void GradColloSliceTranspose3d(SharedData_Hip &data, const CeedInt q, const CeedScalar *__restrict__ r_U, const CeedScalar *c_G,
269                                                  CeedScalar *__restrict__ r_V) {
270   if (data.t_id_x < Q_1d && data.t_id_y < Q_1d) {
271     for (CeedInt comp = 0; comp < NUM_COMP; comp++) {
272       // X derivative
273       data.slice[data.t_id_x + data.t_id_y * T_1D] = r_U[comp + 0 * NUM_COMP];
274       __syncthreads();
275       for (CeedInt i = 0; i < Q_1d; i++) {
276         r_V[q + comp * Q_1d] += c_G[data.t_id_x + i * Q_1d] * data.slice[i + data.t_id_y * T_1D];
277       }
278       __syncthreads();
279       // Y derivative
280       data.slice[data.t_id_x + data.t_id_y * T_1D] = r_U[comp + 1 * NUM_COMP];
281       __syncthreads();
282       for (CeedInt i = 0; i < Q_1d; i++) {
283         r_V[q + comp * Q_1d] += c_G[data.t_id_y + i * Q_1d] * data.slice[data.t_id_x + i * T_1D];
284       }
285       __syncthreads();
286       // Z derivative
287       for (CeedInt i = 0; i < Q_1d; i++) {
288         r_V[i + comp * Q_1d] += c_G[i + q * Q_1d] * r_U[comp + 2 * NUM_COMP];
289       }
290     }
291   }
292 }
293