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