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 #include <ceed.h>
9 #include <ceed/backend.h>
10 #include <ceed/jit-tools.h>
11 #include <stdbool.h>
12 #include <stddef.h>
13 #include <string.h>
14 #include <hip/hip_runtime.h>
15
16 #include "../hip/ceed-hip-common.h"
17 #include "../hip/ceed-hip-compile.h"
18 #include "ceed-hip-ref.h"
19
20 //------------------------------------------------------------------------------
21 // Compile restriction kernels
22 //------------------------------------------------------------------------------
CeedElemRestrictionSetupCompile_Hip(CeedElemRestriction rstr)23 static inline int CeedElemRestrictionSetupCompile_Hip(CeedElemRestriction rstr) {
24 Ceed ceed;
25 bool is_deterministic;
26 CeedInt num_elem, num_comp, elem_size, comp_stride;
27 CeedRestrictionType rstr_type;
28 CeedElemRestriction_Hip *impl;
29
30 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
31 CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed));
32 CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type));
33 CeedCallBackend(CeedElemRestrictionGetNumElements(rstr, &num_elem));
34 CeedCallBackend(CeedElemRestrictionGetNumComponents(rstr, &num_comp));
35 CeedCallBackend(CeedElemRestrictionGetCompStride(rstr, &comp_stride));
36 if (rstr_type == CEED_RESTRICTION_POINTS) {
37 CeedCallBackend(CeedElemRestrictionGetMaxPointsInElement(rstr, &elem_size));
38 } else {
39 CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elem_size));
40 }
41 is_deterministic = impl->d_l_vec_indices != NULL;
42
43 // Compile HIP kernels
44 switch (rstr_type) {
45 case CEED_RESTRICTION_STRIDED: {
46 const char restriction_kernel_source[] = "// Strided restriction source\n#include <ceed/jit-source/hip/hip-ref-restriction-strided.h>\n";
47 bool has_backend_strides;
48 CeedInt strides[3] = {1, num_elem * elem_size, elem_size};
49
50 CeedCallBackend(CeedElemRestrictionHasBackendStrides(rstr, &has_backend_strides));
51 if (!has_backend_strides) {
52 CeedCallBackend(CeedElemRestrictionGetStrides(rstr, strides));
53 }
54 CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem,
55 "RSTR_NUM_COMP", num_comp, "RSTR_STRIDE_NODES", strides[0], "RSTR_STRIDE_COMP", strides[1], "RSTR_STRIDE_ELEM",
56 strides[2]));
57 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "StridedNoTranspose", &impl->ApplyNoTranspose));
58 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "StridedTranspose", &impl->ApplyTranspose));
59 } break;
60 case CEED_RESTRICTION_STANDARD: {
61 const char restriction_kernel_source[] = "// Standard restriction source\n#include <ceed/jit-source/hip/hip-ref-restriction-offset.h>\n";
62
63 CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem,
64 "RSTR_NUM_COMP", num_comp, "RSTR_NUM_NODES", impl->num_nodes, "RSTR_COMP_STRIDE", comp_stride,
65 "USE_DETERMINISTIC", is_deterministic ? 1 : 0));
66 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetNoTranspose", &impl->ApplyNoTranspose));
67 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetTranspose", &impl->ApplyTranspose));
68 } break;
69 case CEED_RESTRICTION_POINTS: {
70 const char restriction_kernel_source[] =
71 "// AtPoints restriction source\n#include <ceed/jit-source/hip/hip-ref-restriction-at-points.h>\n\n"
72 "// Standard restriction source\n#include <ceed/jit-source/hip/hip-ref-restriction-offset.h>\n";
73
74 CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem,
75 "RSTR_NUM_COMP", num_comp, "RSTR_NUM_NODES", impl->num_nodes, "RSTR_COMP_STRIDE", comp_stride,
76 "USE_DETERMINISTIC", is_deterministic ? 1 : 0));
77 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetNoTranspose", &impl->ApplyNoTranspose));
78 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "AtPointsTranspose", &impl->ApplyTranspose));
79 } break;
80 case CEED_RESTRICTION_ORIENTED: {
81 const char restriction_kernel_source[] =
82 "// Oriented restriction source\n#include <ceed/jit-source/hip/hip-ref-restriction-oriented.h>\n\n"
83 "// Standard restriction source\n#include <ceed/jit-source/hip/hip-ref-restriction-offset.h>\n";
84
85 CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem,
86 "RSTR_NUM_COMP", num_comp, "RSTR_NUM_NODES", impl->num_nodes, "RSTR_COMP_STRIDE", comp_stride,
87 "USE_DETERMINISTIC", is_deterministic ? 1 : 0));
88 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OrientedNoTranspose", &impl->ApplyNoTranspose));
89 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetNoTranspose", &impl->ApplyUnsignedNoTranspose));
90 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OrientedTranspose", &impl->ApplyTranspose));
91 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetTranspose", &impl->ApplyUnsignedTranspose));
92 } break;
93 case CEED_RESTRICTION_CURL_ORIENTED: {
94 const char restriction_kernel_source[] =
95 "// Curl oriented restriction source\n#include <ceed/jit-source/hip/hip-ref-restriction-curl-oriented.h>\n\n"
96 "// Standard restriction source\n#include <ceed/jit-source/hip/hip-ref-restriction-offset.h>\n";
97
98 CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem,
99 "RSTR_NUM_COMP", num_comp, "RSTR_NUM_NODES", impl->num_nodes, "RSTR_COMP_STRIDE", comp_stride,
100 "USE_DETERMINISTIC", is_deterministic ? 1 : 0));
101 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedNoTranspose", &impl->ApplyNoTranspose));
102 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedUnsignedNoTranspose", &impl->ApplyUnsignedNoTranspose));
103 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetNoTranspose", &impl->ApplyUnorientedNoTranspose));
104 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedTranspose", &impl->ApplyTranspose));
105 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedUnsignedTranspose", &impl->ApplyUnsignedTranspose));
106 CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetTranspose", &impl->ApplyUnorientedTranspose));
107
108 } break;
109 }
110 CeedCallBackend(CeedDestroy(&ceed));
111 return CEED_ERROR_SUCCESS;
112 }
113
114 //------------------------------------------------------------------------------
115 // Core apply restriction code
116 //------------------------------------------------------------------------------
CeedElemRestrictionApply_Hip_Core(CeedElemRestriction rstr,CeedTransposeMode t_mode,bool use_signs,bool use_orients,CeedVector u,CeedVector v,CeedRequest * request)117 static inline int CeedElemRestrictionApply_Hip_Core(CeedElemRestriction rstr, CeedTransposeMode t_mode, bool use_signs, bool use_orients,
118 CeedVector u, CeedVector v, CeedRequest *request) {
119 Ceed ceed;
120 CeedRestrictionType rstr_type;
121 const CeedScalar *d_u;
122 CeedScalar *d_v;
123 CeedElemRestriction_Hip *impl;
124
125 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
126 CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed));
127 CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type));
128
129 // Assemble kernel if needed
130 if (!impl->module) {
131 CeedCallBackend(CeedElemRestrictionSetupCompile_Hip(rstr));
132 }
133
134 // Get vectors
135 CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
136 if (t_mode == CEED_TRANSPOSE) {
137 // Sum into for transpose mode, e-vec to l-vec
138 CeedCallBackend(CeedVectorGetArray(v, CEED_MEM_DEVICE, &d_v));
139 } else {
140 // Overwrite for notranspose mode, l-vec to e-vec
141 CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
142 }
143
144 // Restrict
145 if (t_mode == CEED_NOTRANSPOSE) {
146 // L-vector -> E-vector
147 CeedInt elem_size;
148
149 CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elem_size));
150 const CeedInt block_size = elem_size < 256 ? (elem_size > 64 ? elem_size : 64) : 256;
151 const CeedInt grid = CeedDivUpInt(impl->num_nodes, block_size);
152
153 switch (rstr_type) {
154 case CEED_RESTRICTION_STRIDED: {
155 void *args[] = {&d_u, &d_v};
156
157 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args));
158 } break;
159 case CEED_RESTRICTION_POINTS:
160 case CEED_RESTRICTION_STANDARD: {
161 void *args[] = {&impl->d_offsets, &d_u, &d_v};
162
163 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args));
164 } break;
165 case CEED_RESTRICTION_ORIENTED: {
166 if (use_signs) {
167 void *args[] = {&impl->d_offsets, &impl->d_orients, &d_u, &d_v};
168
169 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args));
170 } else {
171 void *args[] = {&impl->d_offsets, &d_u, &d_v};
172
173 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedNoTranspose, grid, block_size, args));
174 }
175 } break;
176 case CEED_RESTRICTION_CURL_ORIENTED: {
177 if (use_signs && use_orients) {
178 void *args[] = {&impl->d_offsets, &impl->d_curl_orients, &d_u, &d_v};
179
180 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args));
181 } else if (use_orients) {
182 void *args[] = {&impl->d_offsets, &impl->d_curl_orients, &d_u, &d_v};
183
184 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedNoTranspose, grid, block_size, args));
185 } else {
186 void *args[] = {&impl->d_offsets, &d_u, &d_v};
187
188 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnorientedNoTranspose, grid, block_size, args));
189 }
190 } break;
191 }
192 } else {
193 // E-vector -> L-vector
194 const bool is_deterministic = impl->d_l_vec_indices != NULL;
195 const CeedInt block_size = 64;
196 const CeedInt grid = CeedDivUpInt(impl->num_nodes, block_size);
197
198 switch (rstr_type) {
199 case CEED_RESTRICTION_STRIDED: {
200 void *args[] = {&d_u, &d_v};
201
202 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args));
203 } break;
204 case CEED_RESTRICTION_POINTS: {
205 if (!is_deterministic) {
206 void *args[] = {&impl->d_offsets, &impl->d_points_per_elem, &d_u, &d_v};
207
208 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args));
209 } else {
210 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_points_per_elem, &impl->d_t_offsets, &d_u, &d_v};
211
212 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args));
213 }
214 } break;
215 case CEED_RESTRICTION_STANDARD: {
216 if (!is_deterministic) {
217 void *args[] = {&impl->d_offsets, &d_u, &d_v};
218
219 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args));
220 } else {
221 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &d_u, &d_v};
222
223 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args));
224 }
225 } break;
226 case CEED_RESTRICTION_ORIENTED: {
227 if (use_signs) {
228 if (!is_deterministic) {
229 void *args[] = {&impl->d_offsets, &impl->d_orients, &d_u, &d_v};
230
231 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args));
232 } else {
233 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &impl->d_orients, &d_u, &d_v};
234
235 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args));
236 }
237 } else {
238 if (!is_deterministic) {
239 void *args[] = {&impl->d_offsets, &d_u, &d_v};
240
241 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedTranspose, grid, block_size, args));
242 } else {
243 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &d_u, &d_v};
244
245 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedTranspose, grid, block_size, args));
246 }
247 }
248 } break;
249 case CEED_RESTRICTION_CURL_ORIENTED: {
250 if (use_signs && use_orients) {
251 if (!is_deterministic) {
252 void *args[] = {&impl->d_offsets, &impl->d_curl_orients, &d_u, &d_v};
253
254 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args));
255 } else {
256 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &impl->d_curl_orients, &d_u, &d_v};
257
258 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args));
259 }
260 } else if (use_orients) {
261 if (!is_deterministic) {
262 void *args[] = {&impl->d_offsets, &impl->d_curl_orients, &d_u, &d_v};
263
264 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedTranspose, grid, block_size, args));
265 } else {
266 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &impl->d_curl_orients, &d_u, &d_v};
267
268 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedTranspose, grid, block_size, args));
269 }
270 } else {
271 if (!is_deterministic) {
272 void *args[] = {&impl->d_offsets, &d_u, &d_v};
273
274 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnorientedTranspose, grid, block_size, args));
275 } else {
276 void *args[] = {&impl->d_l_vec_indices, &impl->d_t_indices, &impl->d_t_offsets, &d_u, &d_v};
277
278 CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnorientedTranspose, grid, block_size, args));
279 }
280 }
281 } break;
282 }
283 }
284
285 if (request != CEED_REQUEST_IMMEDIATE && request != CEED_REQUEST_ORDERED) *request = NULL;
286
287 // Restore arrays
288 CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
289 CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
290 CeedCallBackend(CeedDestroy(&ceed));
291 return CEED_ERROR_SUCCESS;
292 }
293
294 //------------------------------------------------------------------------------
295 // Apply restriction
296 //------------------------------------------------------------------------------
CeedElemRestrictionApply_Hip(CeedElemRestriction rstr,CeedTransposeMode t_mode,CeedVector u,CeedVector v,CeedRequest * request)297 static int CeedElemRestrictionApply_Hip(CeedElemRestriction rstr, CeedTransposeMode t_mode, CeedVector u, CeedVector v, CeedRequest *request) {
298 return CeedElemRestrictionApply_Hip_Core(rstr, t_mode, true, true, u, v, request);
299 }
300
301 //------------------------------------------------------------------------------
302 // Apply unsigned restriction
303 //------------------------------------------------------------------------------
CeedElemRestrictionApplyUnsigned_Hip(CeedElemRestriction rstr,CeedTransposeMode t_mode,CeedVector u,CeedVector v,CeedRequest * request)304 static int CeedElemRestrictionApplyUnsigned_Hip(CeedElemRestriction rstr, CeedTransposeMode t_mode, CeedVector u, CeedVector v,
305 CeedRequest *request) {
306 return CeedElemRestrictionApply_Hip_Core(rstr, t_mode, false, true, u, v, request);
307 }
308
309 //------------------------------------------------------------------------------
310 // Apply unoriented restriction
311 //------------------------------------------------------------------------------
CeedElemRestrictionApplyUnoriented_Hip(CeedElemRestriction rstr,CeedTransposeMode t_mode,CeedVector u,CeedVector v,CeedRequest * request)312 static int CeedElemRestrictionApplyUnoriented_Hip(CeedElemRestriction rstr, CeedTransposeMode t_mode, CeedVector u, CeedVector v,
313 CeedRequest *request) {
314 return CeedElemRestrictionApply_Hip_Core(rstr, t_mode, false, false, u, v, request);
315 }
316
317 //------------------------------------------------------------------------------
318 // Get offsets
319 //------------------------------------------------------------------------------
CeedElemRestrictionGetOffsets_Hip(CeedElemRestriction rstr,CeedMemType mem_type,const CeedInt ** offsets)320 static int CeedElemRestrictionGetOffsets_Hip(CeedElemRestriction rstr, CeedMemType mem_type, const CeedInt **offsets) {
321 CeedElemRestriction_Hip *impl;
322 CeedRestrictionType rstr_type;
323
324 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
325 CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type));
326 switch (mem_type) {
327 case CEED_MEM_HOST:
328 *offsets = rstr_type == CEED_RESTRICTION_POINTS ? impl->h_offsets_at_points : impl->h_offsets;
329 break;
330 case CEED_MEM_DEVICE:
331 *offsets = rstr_type == CEED_RESTRICTION_POINTS ? impl->d_offsets_at_points : impl->d_offsets;
332 break;
333 }
334 return CEED_ERROR_SUCCESS;
335 }
336
337 //------------------------------------------------------------------------------
338 // Get orientations
339 //------------------------------------------------------------------------------
CeedElemRestrictionGetOrientations_Hip(CeedElemRestriction rstr,CeedMemType mem_type,const bool ** orients)340 static int CeedElemRestrictionGetOrientations_Hip(CeedElemRestriction rstr, CeedMemType mem_type, const bool **orients) {
341 CeedElemRestriction_Hip *impl;
342 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
343
344 switch (mem_type) {
345 case CEED_MEM_HOST:
346 *orients = impl->h_orients;
347 break;
348 case CEED_MEM_DEVICE:
349 *orients = impl->d_orients;
350 break;
351 }
352 return CEED_ERROR_SUCCESS;
353 }
354
355 //------------------------------------------------------------------------------
356 // Get curl-conforming orientations
357 //------------------------------------------------------------------------------
CeedElemRestrictionGetCurlOrientations_Hip(CeedElemRestriction rstr,CeedMemType mem_type,const CeedInt8 ** curl_orients)358 static int CeedElemRestrictionGetCurlOrientations_Hip(CeedElemRestriction rstr, CeedMemType mem_type, const CeedInt8 **curl_orients) {
359 CeedElemRestriction_Hip *impl;
360 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
361
362 switch (mem_type) {
363 case CEED_MEM_HOST:
364 *curl_orients = impl->h_curl_orients;
365 break;
366 case CEED_MEM_DEVICE:
367 *curl_orients = impl->d_curl_orients;
368 break;
369 }
370 return CEED_ERROR_SUCCESS;
371 }
372
373 //------------------------------------------------------------------------------
374 // Get offset for padded AtPoints E-layout
375 //------------------------------------------------------------------------------
CeedElemRestrictionGetAtPointsElementOffset_Hip(CeedElemRestriction rstr,CeedInt elem,CeedSize * elem_offset)376 static int CeedElemRestrictionGetAtPointsElementOffset_Hip(CeedElemRestriction rstr, CeedInt elem, CeedSize *elem_offset) {
377 CeedInt layout[3];
378
379 CeedCallBackend(CeedElemRestrictionGetELayout(rstr, layout));
380 *elem_offset = 0 * layout[0] + 0 * layout[1] + elem * layout[2];
381 return CEED_ERROR_SUCCESS;
382 }
383
384 //------------------------------------------------------------------------------
385 // Destroy restriction
386 //------------------------------------------------------------------------------
CeedElemRestrictionDestroy_Hip(CeedElemRestriction rstr)387 static int CeedElemRestrictionDestroy_Hip(CeedElemRestriction rstr) {
388 Ceed ceed;
389 CeedElemRestriction_Hip *impl;
390
391 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
392 CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed));
393 if (impl->module) {
394 CeedCallHip(ceed, hipModuleUnload(impl->module));
395 }
396 CeedCallBackend(CeedFree(&impl->h_offsets_owned));
397 CeedCallHip(ceed, hipFree((CeedInt *)impl->d_offsets_owned));
398 CeedCallHip(ceed, hipFree((CeedInt *)impl->d_t_offsets));
399 CeedCallHip(ceed, hipFree((CeedInt *)impl->d_t_indices));
400 CeedCallHip(ceed, hipFree((CeedInt *)impl->d_l_vec_indices));
401 CeedCallBackend(CeedFree(&impl->h_orients_owned));
402 CeedCallHip(ceed, hipFree((bool *)impl->d_orients_owned));
403 CeedCallBackend(CeedFree(&impl->h_curl_orients_owned));
404 CeedCallHip(ceed, hipFree((CeedInt8 *)impl->d_curl_orients_owned));
405 CeedCallBackend(CeedFree(&impl->h_offsets_at_points_owned));
406 CeedCallHip(ceed, hipFree((CeedInt8 *)impl->d_offsets_at_points_owned));
407 CeedCallBackend(CeedFree(&impl->h_points_per_elem_owned));
408 CeedCallHip(ceed, hipFree((CeedInt *)impl->d_points_per_elem_owned));
409 CeedCallBackend(CeedFree(&impl));
410 CeedCallBackend(CeedDestroy(&ceed));
411 return CEED_ERROR_SUCCESS;
412 }
413
414 //------------------------------------------------------------------------------
415 // Create transpose offsets and indices
416 //------------------------------------------------------------------------------
CeedElemRestrictionOffset_Hip(const CeedElemRestriction rstr,const CeedInt elem_size,const CeedInt * indices)417 static int CeedElemRestrictionOffset_Hip(const CeedElemRestriction rstr, const CeedInt elem_size, const CeedInt *indices) {
418 Ceed ceed;
419 bool *is_node;
420 CeedSize l_size;
421 CeedInt num_elem, num_comp, num_nodes = 0;
422 CeedInt *ind_to_offset, *l_vec_indices, *t_offsets, *t_indices;
423 CeedRestrictionType rstr_type;
424 CeedElemRestriction_Hip *impl;
425
426 CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed));
427 CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
428 CeedCallBackend(CeedElemRestrictionGetNumElements(rstr, &num_elem));
429 CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type));
430 CeedCallBackend(CeedElemRestrictionGetLVectorSize(rstr, &l_size));
431 CeedCallBackend(CeedElemRestrictionGetNumComponents(rstr, &num_comp));
432 const CeedInt size_indices = num_elem * elem_size;
433
434 // Count num_nodes
435 CeedCallBackend(CeedCalloc(l_size, &is_node));
436
437 for (CeedInt i = 0; i < size_indices; i++) is_node[indices[i]] = 1;
438 for (CeedInt i = 0; i < l_size; i++) num_nodes += is_node[i];
439 impl->num_nodes = num_nodes;
440
441 // L-vector offsets array
442 CeedCallBackend(CeedCalloc(l_size, &ind_to_offset));
443 CeedCallBackend(CeedCalloc(num_nodes, &l_vec_indices));
444 for (CeedInt i = 0, j = 0; i < l_size; i++) {
445 if (is_node[i]) {
446 l_vec_indices[j] = i;
447 ind_to_offset[i] = j++;
448 }
449 }
450 CeedCallBackend(CeedFree(&is_node));
451
452 // Compute transpose offsets and indices
453 const CeedInt size_offsets = num_nodes + 1;
454
455 CeedCallBackend(CeedCalloc(size_offsets, &t_offsets));
456 CeedCallBackend(CeedMalloc(size_indices, &t_indices));
457 // Count node multiplicity
458 for (CeedInt e = 0; e < num_elem; ++e) {
459 for (CeedInt i = 0; i < elem_size; ++i) ++t_offsets[ind_to_offset[indices[elem_size * e + i]] + 1];
460 }
461 // Convert to running sum
462 for (CeedInt i = 1; i < size_offsets; ++i) t_offsets[i] += t_offsets[i - 1];
463 // List all E-vec indices associated with L-vec node
464 for (CeedInt e = 0; e < num_elem; ++e) {
465 for (CeedInt i = 0; i < elem_size; ++i) {
466 const CeedInt lid = elem_size * e + i;
467 const CeedInt gid = indices[lid];
468
469 t_indices[t_offsets[ind_to_offset[gid]]++] = lid;
470 }
471 }
472 // Reset running sum
473 for (int i = size_offsets - 1; i > 0; --i) t_offsets[i] = t_offsets[i - 1];
474 t_offsets[0] = 0;
475
476 // Copy data to device
477 // -- L-vector indices
478 CeedCallHip(ceed, hipMalloc((void **)&impl->d_l_vec_indices, num_nodes * sizeof(CeedInt)));
479 CeedCallHip(ceed, hipMemcpy((CeedInt *)impl->d_l_vec_indices, l_vec_indices, num_nodes * sizeof(CeedInt), hipMemcpyHostToDevice));
480 // -- Transpose offsets
481 CeedCallHip(ceed, hipMalloc((void **)&impl->d_t_offsets, size_offsets * sizeof(CeedInt)));
482 CeedCallHip(ceed, hipMemcpy((CeedInt *)impl->d_t_offsets, t_offsets, size_offsets * sizeof(CeedInt), hipMemcpyHostToDevice));
483 // -- Transpose indices
484 CeedCallHip(ceed, hipMalloc((void **)&impl->d_t_indices, size_indices * sizeof(CeedInt)));
485 CeedCallHip(ceed, hipMemcpy((CeedInt *)impl->d_t_indices, t_indices, size_indices * sizeof(CeedInt), hipMemcpyHostToDevice));
486
487 // Cleanup
488 CeedCallBackend(CeedFree(&ind_to_offset));
489 CeedCallBackend(CeedFree(&l_vec_indices));
490 CeedCallBackend(CeedFree(&t_offsets));
491 CeedCallBackend(CeedFree(&t_indices));
492 CeedCallBackend(CeedDestroy(&ceed));
493 return CEED_ERROR_SUCCESS;
494 }
495
496 //------------------------------------------------------------------------------
497 // Create restriction
498 //------------------------------------------------------------------------------
CeedElemRestrictionCreate_Hip(CeedMemType mem_type,CeedCopyMode copy_mode,const CeedInt * offsets,const bool * orients,const CeedInt8 * curl_orients,CeedElemRestriction rstr)499 int CeedElemRestrictionCreate_Hip(CeedMemType mem_type, CeedCopyMode copy_mode, const CeedInt *offsets, const bool *orients,
500 const CeedInt8 *curl_orients, CeedElemRestriction rstr) {
501 Ceed ceed, ceed_parent;
502 bool is_deterministic;
503 CeedInt num_elem, num_comp, elem_size;
504 CeedRestrictionType rstr_type;
505 CeedElemRestriction_Hip *impl;
506
507 CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed));
508 CeedCallBackend(CeedGetParent(ceed, &ceed_parent));
509 CeedCallBackend(CeedIsDeterministic(ceed_parent, &is_deterministic));
510 CeedCallBackend(CeedDestroy(&ceed_parent));
511 CeedCallBackend(CeedElemRestrictionGetNumElements(rstr, &num_elem));
512 CeedCallBackend(CeedElemRestrictionGetNumComponents(rstr, &num_comp));
513 CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elem_size));
514 CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type));
515 // Use max number of points as elem size for AtPoints restrictions
516 if (rstr_type == CEED_RESTRICTION_POINTS) {
517 CeedInt max_points = 0;
518
519 for (CeedInt i = 0; i < num_elem; i++) {
520 max_points = CeedIntMax(max_points, offsets[i + 1] - offsets[i]);
521 }
522 elem_size = max_points;
523 }
524 const CeedInt size = num_elem * elem_size;
525
526 CeedCallBackend(CeedCalloc(1, &impl));
527 impl->num_nodes = size;
528 CeedCallBackend(CeedElemRestrictionSetData(rstr, impl));
529
530 // Set layouts
531 {
532 bool has_backend_strides;
533 CeedInt layout[3] = {1, size, elem_size};
534
535 CeedCallBackend(CeedElemRestrictionSetELayout(rstr, layout));
536 if (rstr_type == CEED_RESTRICTION_STRIDED) {
537 CeedCallBackend(CeedElemRestrictionHasBackendStrides(rstr, &has_backend_strides));
538 if (has_backend_strides) {
539 CeedCallBackend(CeedElemRestrictionSetLLayout(rstr, layout));
540 }
541 }
542 }
543
544 // Pad AtPoints indices
545 if (rstr_type == CEED_RESTRICTION_POINTS) {
546 CeedSize offsets_len = elem_size * num_elem, at_points_size = num_elem + 1;
547 CeedInt max_points = elem_size, *offsets_padded, *points_per_elem;
548
549 CeedCheck(mem_type == CEED_MEM_HOST, ceed, CEED_ERROR_BACKEND, "only MemType Host supported when creating AtPoints restriction");
550 CeedCallBackend(CeedMalloc(offsets_len, &offsets_padded));
551 CeedCallBackend(CeedMalloc(num_elem, &points_per_elem));
552 for (CeedInt i = 0; i < num_elem; i++) {
553 CeedInt num_points = offsets[i + 1] - offsets[i];
554 CeedInt last_point = 0;
555
556 points_per_elem[i] = num_points;
557 at_points_size += num_points;
558 // -- Copy all points in element
559 for (CeedInt j = 0; j < num_points; j++) {
560 offsets_padded[i * max_points + j] = offsets[offsets[i] + j] * num_comp;
561 last_point = offsets_padded[i * max_points + j];
562 }
563 // -- Replicate out last point in element
564 for (CeedInt j = num_points; j < max_points; j++) {
565 offsets_padded[i * max_points + j] = last_point;
566 }
567 }
568 CeedCallBackend(CeedSetHostCeedIntArray(offsets, copy_mode, at_points_size, &impl->h_offsets_at_points_owned, &impl->h_offsets_at_points_borrowed,
569 &impl->h_offsets_at_points));
570 CeedCallHip(ceed, hipMalloc((void **)&impl->d_offsets_at_points_owned, at_points_size * sizeof(CeedInt)));
571 CeedCallHip(ceed, hipMemcpy((CeedInt **)impl->d_offsets_at_points_owned, impl->h_offsets_at_points, at_points_size * sizeof(CeedInt),
572 hipMemcpyHostToDevice));
573 impl->d_offsets_at_points = (CeedInt *)impl->d_offsets_at_points_owned;
574
575 // -- Use padded offsets for the rest of the setup
576 offsets = (const CeedInt *)offsets_padded;
577 copy_mode = CEED_OWN_POINTER;
578 CeedCallBackend(CeedElemRestrictionSetAtPointsEVectorSize(rstr, elem_size * num_elem * num_comp));
579
580 // -- Points per element
581 CeedCallBackend(CeedSetHostCeedIntArray(points_per_elem, CEED_OWN_POINTER, num_elem, &impl->h_points_per_elem_owned,
582 &impl->h_points_per_elem_borrowed, &impl->h_points_per_elem));
583 CeedCallHip(ceed, hipMalloc((void **)&impl->d_points_per_elem_owned, num_elem * sizeof(CeedInt)));
584 CeedCallHip(ceed,
585 hipMemcpy((CeedInt **)impl->d_points_per_elem_owned, impl->h_points_per_elem, num_elem * sizeof(CeedInt), hipMemcpyHostToDevice));
586 impl->d_points_per_elem = (CeedInt *)impl->d_points_per_elem_owned;
587 }
588
589 // Set up device offset/orientation arrays
590 if (rstr_type != CEED_RESTRICTION_STRIDED) {
591 switch (mem_type) {
592 case CEED_MEM_HOST: {
593 CeedCallBackend(CeedSetHostCeedIntArray(offsets, copy_mode, size, &impl->h_offsets_owned, &impl->h_offsets_borrowed, &impl->h_offsets));
594 CeedCallHip(ceed, hipMalloc((void **)&impl->d_offsets_owned, size * sizeof(CeedInt)));
595 CeedCallHip(ceed, hipMemcpy((CeedInt **)impl->d_offsets_owned, impl->h_offsets, size * sizeof(CeedInt), hipMemcpyHostToDevice));
596 impl->d_offsets = (CeedInt *)impl->d_offsets_owned;
597 if (is_deterministic) CeedCallBackend(CeedElemRestrictionOffset_Hip(rstr, elem_size, offsets));
598 } break;
599 case CEED_MEM_DEVICE: {
600 CeedCallBackend(CeedSetDeviceCeedIntArray_Hip(ceed, offsets, copy_mode, size, &impl->d_offsets_owned, &impl->d_offsets_borrowed,
601 (const CeedInt **)&impl->d_offsets));
602 CeedCallBackend(CeedMalloc(size, &impl->h_offsets_owned));
603 CeedCallHip(ceed, hipMemcpy((CeedInt **)impl->h_offsets_owned, impl->d_offsets, size * sizeof(CeedInt), hipMemcpyDeviceToHost));
604 impl->h_offsets = impl->h_offsets_owned;
605 if (is_deterministic) CeedCallBackend(CeedElemRestrictionOffset_Hip(rstr, elem_size, offsets));
606 } break;
607 }
608
609 // Orientation data
610 if (rstr_type == CEED_RESTRICTION_ORIENTED) {
611 switch (mem_type) {
612 case CEED_MEM_HOST: {
613 CeedCallBackend(CeedSetHostBoolArray(orients, copy_mode, size, &impl->h_orients_owned, &impl->h_orients_borrowed, &impl->h_orients));
614 CeedCallHip(ceed, hipMalloc((void **)&impl->d_orients_owned, size * sizeof(bool)));
615 CeedCallHip(ceed, hipMemcpy((bool *)impl->d_orients_owned, impl->h_orients, size * sizeof(bool), hipMemcpyHostToDevice));
616 impl->d_orients = impl->d_orients_owned;
617 } break;
618 case CEED_MEM_DEVICE: {
619 CeedCallBackend(CeedSetDeviceBoolArray_Hip(ceed, orients, copy_mode, size, &impl->d_orients_owned, &impl->d_orients_borrowed,
620 (const bool **)&impl->d_orients));
621 CeedCallBackend(CeedMalloc(size, &impl->h_orients_owned));
622 CeedCallHip(ceed, hipMemcpy((bool *)impl->h_orients_owned, impl->d_orients, size * sizeof(bool), hipMemcpyDeviceToHost));
623 impl->h_orients = impl->h_orients_owned;
624 } break;
625 }
626 } else if (rstr_type == CEED_RESTRICTION_CURL_ORIENTED) {
627 switch (mem_type) {
628 case CEED_MEM_HOST: {
629 CeedCallBackend(CeedSetHostCeedInt8Array(curl_orients, copy_mode, 3 * size, &impl->h_curl_orients_owned, &impl->h_curl_orients_borrowed,
630 &impl->h_curl_orients));
631 CeedCallHip(ceed, hipMalloc((void **)&impl->d_curl_orients_owned, 3 * size * sizeof(CeedInt8)));
632 CeedCallHip(ceed,
633 hipMemcpy((CeedInt8 *)impl->d_curl_orients_owned, impl->h_curl_orients, 3 * size * sizeof(CeedInt8), hipMemcpyHostToDevice));
634 impl->d_curl_orients = impl->d_curl_orients_owned;
635 } break;
636 case CEED_MEM_DEVICE: {
637 CeedCallBackend(CeedSetDeviceCeedInt8Array_Hip(ceed, curl_orients, copy_mode, 3 * size, &impl->d_curl_orients_owned,
638 &impl->d_curl_orients_borrowed, (const CeedInt8 **)&impl->d_curl_orients));
639 CeedCallBackend(CeedMalloc(3 * size, &impl->h_curl_orients_owned));
640 CeedCallHip(ceed,
641 hipMemcpy((CeedInt8 *)impl->h_curl_orients_owned, impl->d_curl_orients, 3 * size * sizeof(CeedInt8), hipMemcpyDeviceToHost));
642 impl->h_curl_orients = impl->h_curl_orients_owned;
643 } break;
644 }
645 }
646 }
647
648 // Register backend functions
649 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "Apply", CeedElemRestrictionApply_Hip));
650 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "ApplyUnsigned", CeedElemRestrictionApplyUnsigned_Hip));
651 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "ApplyUnoriented", CeedElemRestrictionApplyUnoriented_Hip));
652 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetOffsets", CeedElemRestrictionGetOffsets_Hip));
653 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetOrientations", CeedElemRestrictionGetOrientations_Hip));
654 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetCurlOrientations", CeedElemRestrictionGetCurlOrientations_Hip));
655 if (rstr_type == CEED_RESTRICTION_POINTS) {
656 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetAtPointsElementOffset",
657 CeedElemRestrictionGetAtPointsElementOffset_Hip));
658 }
659 CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "Destroy", CeedElemRestrictionDestroy_Hip));
660 CeedCallBackend(CeedDestroy(&ceed));
661 return CEED_ERROR_SUCCESS;
662 }
663
664 //------------------------------------------------------------------------------
665