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