xref: /libCEED/backends/hip-ref/ceed-hip-ref-restriction.c (revision 9a25c3513918281076c2babe50808fbe5e3a546e)
1 // Copyright (c) 2017-2022, Lawrence Livermore National Security, LLC and other CEED contributors.
2 // All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
3 //
4 // SPDX-License-Identifier: BSD-2-Clause
5 //
6 // This file is part of CEED:  http://github.com/ceed
7 
8 #include <ceed.h>
9 #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(CeedElemRestrictionGetNumElements(rstr, &num_elem));
35   CeedCallBackend(CeedElemRestrictionGetNumComponents(rstr, &num_comp));
36   CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elem_size));
37   CeedCallBackend(CeedElemRestrictionGetCompStride(rstr, &comp_stride));
38   CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type));
39   is_deterministic = impl->d_l_vec_indices != NULL;
40 
41   // Compile HIP kernels
42   switch (rstr_type) {
43     case CEED_RESTRICTION_STRIDED: {
44       CeedInt strides[3] = {1, num_elem * elem_size, elem_size};
45       bool    has_backend_strides;
46 
47       CeedCallBackend(CeedElemRestrictionHasBackendStrides(rstr, &has_backend_strides));
48       if (!has_backend_strides) {
49         CeedCallBackend(CeedElemRestrictionGetStrides(rstr, strides));
50       }
51 
52       CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-strided.h", &restriction_kernel_path));
53       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source -----\n");
54       CeedCallBackend(CeedLoadSourceToBuffer(ceed, restriction_kernel_path, &restriction_kernel_source));
55       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source Complete! -----\n");
56       CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem,
57                                       "RSTR_NUM_COMP", num_comp, "RSTR_STRIDE_NODES", strides[0], "RSTR_STRIDE_COMP", strides[1], "RSTR_STRIDE_ELEM",
58                                       strides[2]));
59       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "StridedNoTranspose", &impl->ApplyNoTranspose));
60       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "StridedTranspose", &impl->ApplyTranspose));
61     } break;
62     case CEED_RESTRICTION_STANDARD: {
63       CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-offset.h", &restriction_kernel_path));
64       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source -----\n");
65       CeedCallBackend(CeedLoadSourceToBuffer(ceed, restriction_kernel_path, &restriction_kernel_source));
66       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source Complete! -----\n");
67       CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem,
68                                       "RSTR_NUM_COMP", num_comp, "RSTR_NUM_NODES", impl->num_nodes, "RSTR_COMP_STRIDE", comp_stride,
69                                       "USE_DETERMINISTIC", is_deterministic ? 1 : 0));
70       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetNoTranspose", &impl->ApplyNoTranspose));
71       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetTranspose", &impl->ApplyTranspose));
72     } break;
73     case CEED_RESTRICTION_ORIENTED: {
74       const char *offset_kernel_path;
75 
76       CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-oriented.h", &restriction_kernel_path));
77       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source -----\n");
78       CeedCallBackend(CeedLoadSourceToBuffer(ceed, restriction_kernel_path, &restriction_kernel_source));
79       CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-offset.h", &offset_kernel_path));
80       CeedCallBackend(CeedLoadSourceToInitializedBuffer(ceed, offset_kernel_path, &restriction_kernel_source));
81       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source Complete! -----\n");
82       CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem,
83                                       "RSTR_NUM_COMP", num_comp, "RSTR_NUM_NODES", impl->num_nodes, "RSTR_COMP_STRIDE", comp_stride,
84                                       "USE_DETERMINISTIC", is_deterministic ? 1 : 0));
85       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OrientedNoTranspose", &impl->ApplyNoTranspose));
86       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetNoTranspose", &impl->ApplyUnsignedNoTranspose));
87       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OrientedTranspose", &impl->ApplyTranspose));
88       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetTranspose", &impl->ApplyUnsignedTranspose));
89       CeedCallBackend(CeedFree(&offset_kernel_path));
90     } break;
91     case CEED_RESTRICTION_CURL_ORIENTED: {
92       const char *offset_kernel_path;
93 
94       CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-curl-oriented.h", &restriction_kernel_path));
95       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source -----\n");
96       CeedCallBackend(CeedLoadSourceToBuffer(ceed, restriction_kernel_path, &restriction_kernel_source));
97       CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-restriction-offset.h", &offset_kernel_path));
98       CeedCallBackend(CeedLoadSourceToInitializedBuffer(ceed, offset_kernel_path, &restriction_kernel_source));
99       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Restriction Kernel Source Complete! -----\n");
100       CeedCallBackend(CeedCompile_Hip(ceed, restriction_kernel_source, &impl->module, 6, "RSTR_ELEM_SIZE", elem_size, "RSTR_NUM_ELEM", num_elem,
101                                       "RSTR_NUM_COMP", num_comp, "RSTR_NUM_NODES", impl->num_nodes, "RSTR_COMP_STRIDE", comp_stride,
102                                       "USE_DETERMINISTIC", is_deterministic ? 1 : 0));
103       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedNoTranspose", &impl->ApplyNoTranspose));
104       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedUnsignedNoTranspose", &impl->ApplyUnsignedNoTranspose));
105       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetNoTranspose", &impl->ApplyUnorientedNoTranspose));
106       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedTranspose", &impl->ApplyTranspose));
107       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "CurlOrientedUnsignedTranspose", &impl->ApplyUnsignedTranspose));
108       CeedCallBackend(CeedGetKernel_Hip(ceed, impl->module, "OffsetTranspose", &impl->ApplyUnorientedTranspose));
109       CeedCallBackend(CeedFree(&offset_kernel_path));
110     } break;
111     case CEED_RESTRICTION_POINTS: {
112       // LCOV_EXCL_START
113       return CeedError(ceed, CEED_ERROR_UNSUPPORTED, "Backend does not implement restriction CeedElemRestrictionAtPoints");
114       // LCOV_EXCL_STOP
115     } break;
116   }
117   CeedCallBackend(CeedFree(&restriction_kernel_path));
118   CeedCallBackend(CeedFree(&restriction_kernel_source));
119   return CEED_ERROR_SUCCESS;
120 }
121 
122 //------------------------------------------------------------------------------
123 // Core apply restriction code
124 //------------------------------------------------------------------------------
125 static inline int CeedElemRestrictionApply_Hip_Core(CeedElemRestriction rstr, CeedTransposeMode t_mode, bool use_signs, bool use_orients,
126                                                     CeedVector u, CeedVector v, CeedRequest *request) {
127   Ceed                     ceed;
128   CeedRestrictionType      rstr_type;
129   const CeedScalar        *d_u;
130   CeedScalar              *d_v;
131   CeedElemRestriction_Hip *impl;
132 
133   CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
134   CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed));
135   CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type));
136 
137   // Assemble kernel if needed
138   if (!impl->module) {
139     CeedCallBackend(CeedElemRestrictionSetupCompile_Hip(rstr));
140   }
141 
142   // Get vectors
143   CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
144   if (t_mode == CEED_TRANSPOSE) {
145     // Sum into for transpose mode, e-vec to l-vec
146     CeedCallBackend(CeedVectorGetArray(v, CEED_MEM_DEVICE, &d_v));
147   } else {
148     // Overwrite for notranspose mode, l-vec to e-vec
149     CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
150   }
151 
152   // Restrict
153   if (t_mode == CEED_NOTRANSPOSE) {
154     // L-vector -> E-vector
155     CeedInt elem_size;
156 
157     CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elem_size));
158     const CeedInt block_size = elem_size < 256 ? (elem_size > 64 ? elem_size : 64) : 256;
159     const CeedInt grid       = CeedDivUpInt(impl->num_nodes, block_size);
160 
161     switch (rstr_type) {
162       case CEED_RESTRICTION_STRIDED: {
163         void *args[] = {&d_u, &d_v};
164 
165         CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args));
166       } break;
167       case CEED_RESTRICTION_STANDARD: {
168         void *args[] = {&impl->d_ind, &d_u, &d_v};
169 
170         CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args));
171       } break;
172       case CEED_RESTRICTION_ORIENTED: {
173         if (use_signs) {
174           void *args[] = {&impl->d_ind, &impl->d_orients, &d_u, &d_v};
175 
176           CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args));
177         } else {
178           void *args[] = {&impl->d_ind, &d_u, &d_v};
179 
180           CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedNoTranspose, grid, block_size, args));
181         }
182       } break;
183       case CEED_RESTRICTION_CURL_ORIENTED: {
184         if (use_signs && use_orients) {
185           void *args[] = {&impl->d_ind, &impl->d_curl_orients, &d_u, &d_v};
186 
187           CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyNoTranspose, grid, block_size, args));
188         } else if (use_orients) {
189           void *args[] = {&impl->d_ind, &impl->d_curl_orients, &d_u, &d_v};
190 
191           CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnsignedNoTranspose, grid, block_size, args));
192         } else {
193           void *args[] = {&impl->d_ind, &d_u, &d_v};
194 
195           CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyUnorientedNoTranspose, grid, block_size, args));
196         }
197       } break;
198       case CEED_RESTRICTION_POINTS: {
199         // LCOV_EXCL_START
200         return CeedError(ceed, CEED_ERROR_UNSUPPORTED, "Backend does not implement restriction CeedElemRestrictionAtPoints");
201         // LCOV_EXCL_STOP
202       } break;
203     }
204   } else {
205     // E-vector -> L-vector
206     const bool    is_deterministic = impl->d_l_vec_indices != NULL;
207     const CeedInt block_size       = 64;
208     const CeedInt grid             = CeedDivUpInt(impl->num_nodes, block_size);
209 
210     switch (rstr_type) {
211       case CEED_RESTRICTION_STRIDED: {
212         void *args[] = {&d_u, &d_v};
213 
214         CeedCallBackend(CeedRunKernel_Hip(ceed, impl->ApplyTranspose, grid, block_size, args));
215       } break;
216       case CEED_RESTRICTION_STANDARD: {
217         if (!is_deterministic) {
218           void *args[] = {&impl->d_ind, &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_ind, &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_ind, &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_ind, &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_ind, &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_ind, &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       case CEED_RESTRICTION_POINTS: {
284         // LCOV_EXCL_START
285         return CeedError(ceed, CEED_ERROR_UNSUPPORTED, "Backend does not implement restriction CeedElemRestrictionAtPoints");
286         // LCOV_EXCL_STOP
287       } break;
288     }
289   }
290 
291   if (request != CEED_REQUEST_IMMEDIATE && request != CEED_REQUEST_ORDERED) *request = NULL;
292 
293   // Restore arrays
294   CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
295   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
296   return CEED_ERROR_SUCCESS;
297 }
298 
299 //------------------------------------------------------------------------------
300 // Apply restriction
301 //------------------------------------------------------------------------------
302 static int CeedElemRestrictionApply_Hip(CeedElemRestriction rstr, CeedTransposeMode t_mode, CeedVector u, CeedVector v, CeedRequest *request) {
303   return CeedElemRestrictionApply_Hip_Core(rstr, t_mode, true, true, u, v, request);
304 }
305 
306 //------------------------------------------------------------------------------
307 // Apply unsigned restriction
308 //------------------------------------------------------------------------------
309 static int CeedElemRestrictionApplyUnsigned_Hip(CeedElemRestriction rstr, CeedTransposeMode t_mode, CeedVector u, CeedVector v,
310                                                 CeedRequest *request) {
311   return CeedElemRestrictionApply_Hip_Core(rstr, t_mode, false, true, u, v, request);
312 }
313 
314 //------------------------------------------------------------------------------
315 // Apply unoriented restriction
316 //------------------------------------------------------------------------------
317 static int CeedElemRestrictionApplyUnoriented_Hip(CeedElemRestriction rstr, CeedTransposeMode t_mode, CeedVector u, CeedVector v,
318                                                   CeedRequest *request) {
319   return CeedElemRestrictionApply_Hip_Core(rstr, t_mode, false, false, u, v, request);
320 }
321 
322 //------------------------------------------------------------------------------
323 // Get offsets
324 //------------------------------------------------------------------------------
325 static int CeedElemRestrictionGetOffsets_Hip(CeedElemRestriction rstr, CeedMemType mem_type, const CeedInt **offsets) {
326   CeedElemRestriction_Hip *impl;
327 
328   CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
329   switch (mem_type) {
330     case CEED_MEM_HOST:
331       *offsets = impl->h_ind;
332       break;
333     case CEED_MEM_DEVICE:
334       *offsets = impl->d_ind;
335       break;
336   }
337   return CEED_ERROR_SUCCESS;
338 }
339 
340 //------------------------------------------------------------------------------
341 // Get orientations
342 //------------------------------------------------------------------------------
343 static int CeedElemRestrictionGetOrientations_Hip(CeedElemRestriction rstr, CeedMemType mem_type, const bool **orients) {
344   CeedElemRestriction_Hip *impl;
345   CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
346 
347   switch (mem_type) {
348     case CEED_MEM_HOST:
349       *orients = impl->h_orients;
350       break;
351     case CEED_MEM_DEVICE:
352       *orients = impl->d_orients;
353       break;
354   }
355   return CEED_ERROR_SUCCESS;
356 }
357 
358 //------------------------------------------------------------------------------
359 // Get curl-conforming orientations
360 //------------------------------------------------------------------------------
361 static int CeedElemRestrictionGetCurlOrientations_Hip(CeedElemRestriction rstr, CeedMemType mem_type, const CeedInt8 **curl_orients) {
362   CeedElemRestriction_Hip *impl;
363   CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
364 
365   switch (mem_type) {
366     case CEED_MEM_HOST:
367       *curl_orients = impl->h_curl_orients;
368       break;
369     case CEED_MEM_DEVICE:
370       *curl_orients = impl->d_curl_orients;
371       break;
372   }
373   return CEED_ERROR_SUCCESS;
374 }
375 
376 //------------------------------------------------------------------------------
377 // Destroy restriction
378 //------------------------------------------------------------------------------
379 static int CeedElemRestrictionDestroy_Hip(CeedElemRestriction rstr) {
380   Ceed                     ceed;
381   CeedElemRestriction_Hip *impl;
382 
383   CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
384   CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed));
385   if (impl->module) {
386     CeedCallHip(ceed, hipModuleUnload(impl->module));
387   }
388   CeedCallBackend(CeedFree(&impl->h_ind_allocated));
389   CeedCallHip(ceed, hipFree(impl->d_ind_allocated));
390   CeedCallHip(ceed, hipFree(impl->d_t_offsets));
391   CeedCallHip(ceed, hipFree(impl->d_t_indices));
392   CeedCallHip(ceed, hipFree(impl->d_l_vec_indices));
393   CeedCallBackend(CeedFree(&impl->h_orients_allocated));
394   CeedCallHip(ceed, hipFree(impl->d_orients_allocated));
395   CeedCallBackend(CeedFree(&impl->h_curl_orients_allocated));
396   CeedCallHip(ceed, hipFree(impl->d_curl_orients_allocated));
397   CeedCallBackend(CeedFree(&impl));
398   return CEED_ERROR_SUCCESS;
399 }
400 
401 //------------------------------------------------------------------------------
402 // Create transpose offsets and indices
403 //------------------------------------------------------------------------------
404 static int CeedElemRestrictionOffset_Hip(const CeedElemRestriction rstr, const CeedInt *indices) {
405   Ceed                     ceed;
406   bool                    *is_node;
407   CeedSize                 l_size;
408   CeedInt                  num_elem, elem_size, num_comp, num_nodes = 0;
409   CeedInt                 *ind_to_offset, *l_vec_indices, *t_offsets, *t_indices;
410   CeedElemRestriction_Hip *impl;
411 
412   CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed));
413   CeedCallBackend(CeedElemRestrictionGetData(rstr, &impl));
414   CeedCallBackend(CeedElemRestrictionGetNumElements(rstr, &num_elem));
415   CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elem_size));
416   CeedCallBackend(CeedElemRestrictionGetLVectorSize(rstr, &l_size));
417   CeedCallBackend(CeedElemRestrictionGetNumComponents(rstr, &num_comp));
418   const CeedInt size_indices = num_elem * elem_size;
419 
420   // Count num_nodes
421   CeedCallBackend(CeedCalloc(l_size, &is_node));
422 
423   for (CeedInt i = 0; i < size_indices; i++) is_node[indices[i]] = 1;
424   for (CeedInt i = 0; i < l_size; i++) num_nodes += is_node[i];
425   impl->num_nodes = num_nodes;
426 
427   // L-vector offsets array
428   CeedCallBackend(CeedCalloc(l_size, &ind_to_offset));
429   CeedCallBackend(CeedCalloc(num_nodes, &l_vec_indices));
430   for (CeedInt i = 0, j = 0; i < l_size; i++) {
431     if (is_node[i]) {
432       l_vec_indices[j] = i;
433       ind_to_offset[i] = j++;
434     }
435   }
436   CeedCallBackend(CeedFree(&is_node));
437 
438   // Compute transpose offsets and indices
439   const CeedInt size_offsets = num_nodes + 1;
440 
441   CeedCallBackend(CeedCalloc(size_offsets, &t_offsets));
442   CeedCallBackend(CeedMalloc(size_indices, &t_indices));
443   // Count node multiplicity
444   for (CeedInt e = 0; e < num_elem; ++e) {
445     for (CeedInt i = 0; i < elem_size; ++i) ++t_offsets[ind_to_offset[indices[elem_size * e + i]] + 1];
446   }
447   // Convert to running sum
448   for (CeedInt i = 1; i < size_offsets; ++i) t_offsets[i] += t_offsets[i - 1];
449   // List all E-vec indices associated with L-vec node
450   for (CeedInt e = 0; e < num_elem; ++e) {
451     for (CeedInt i = 0; i < elem_size; ++i) {
452       const CeedInt lid = elem_size * e + i;
453       const CeedInt gid = indices[lid];
454 
455       t_indices[t_offsets[ind_to_offset[gid]]++] = lid;
456     }
457   }
458   // Reset running sum
459   for (int i = size_offsets - 1; i > 0; --i) t_offsets[i] = t_offsets[i - 1];
460   t_offsets[0] = 0;
461 
462   // Copy data to device
463   // -- L-vector indices
464   CeedCallHip(ceed, hipMalloc((void **)&impl->d_l_vec_indices, num_nodes * sizeof(CeedInt)));
465   CeedCallHip(ceed, hipMemcpy(impl->d_l_vec_indices, l_vec_indices, num_nodes * sizeof(CeedInt), hipMemcpyHostToDevice));
466   // -- Transpose offsets
467   CeedCallHip(ceed, hipMalloc((void **)&impl->d_t_offsets, size_offsets * sizeof(CeedInt)));
468   CeedCallHip(ceed, hipMemcpy(impl->d_t_offsets, t_offsets, size_offsets * sizeof(CeedInt), hipMemcpyHostToDevice));
469   // -- Transpose indices
470   CeedCallHip(ceed, hipMalloc((void **)&impl->d_t_indices, size_indices * sizeof(CeedInt)));
471   CeedCallHip(ceed, hipMemcpy(impl->d_t_indices, t_indices, size_indices * sizeof(CeedInt), hipMemcpyHostToDevice));
472 
473   // Cleanup
474   CeedCallBackend(CeedFree(&ind_to_offset));
475   CeedCallBackend(CeedFree(&l_vec_indices));
476   CeedCallBackend(CeedFree(&t_offsets));
477   CeedCallBackend(CeedFree(&t_indices));
478   return CEED_ERROR_SUCCESS;
479 }
480 
481 //------------------------------------------------------------------------------
482 // Create restriction
483 //------------------------------------------------------------------------------
484 int CeedElemRestrictionCreate_Hip(CeedMemType mem_type, CeedCopyMode copy_mode, const CeedInt *indices, const bool *orients,
485                                   const CeedInt8 *curl_orients, CeedElemRestriction rstr) {
486   Ceed                     ceed, ceed_parent;
487   bool                     is_deterministic;
488   CeedInt                  num_elem, elem_size;
489   CeedRestrictionType      rstr_type;
490   CeedElemRestriction_Hip *impl;
491 
492   CeedCallBackend(CeedElemRestrictionGetCeed(rstr, &ceed));
493   CeedCallBackend(CeedGetParent(ceed, &ceed_parent));
494   CeedCallBackend(CeedIsDeterministic(ceed_parent, &is_deterministic));
495   CeedCallBackend(CeedElemRestrictionGetNumElements(rstr, &num_elem));
496   CeedCallBackend(CeedElemRestrictionGetElementSize(rstr, &elem_size));
497   CeedCallBackend(CeedElemRestrictionGetType(rstr, &rstr_type));
498   const CeedInt size = num_elem * elem_size;
499 
500   CeedCallBackend(CeedCalloc(1, &impl));
501   impl->num_nodes                = size;
502   impl->h_ind                    = NULL;
503   impl->h_ind_allocated          = NULL;
504   impl->d_ind                    = NULL;
505   impl->d_ind_allocated          = NULL;
506   impl->d_t_indices              = NULL;
507   impl->d_t_offsets              = NULL;
508   impl->h_orients                = NULL;
509   impl->h_orients_allocated      = NULL;
510   impl->d_orients                = NULL;
511   impl->d_orients_allocated      = NULL;
512   impl->h_curl_orients           = NULL;
513   impl->h_curl_orients_allocated = NULL;
514   impl->d_curl_orients           = NULL;
515   impl->d_curl_orients_allocated = NULL;
516   CeedCallBackend(CeedElemRestrictionSetData(rstr, impl));
517 
518   // Set layouts
519   {
520     bool    has_backend_strides;
521     CeedInt layout[3] = {1, size, elem_size};
522 
523     CeedCallBackend(CeedElemRestrictionSetELayout(rstr, layout));
524     if (rstr_type == CEED_RESTRICTION_STRIDED) {
525       CeedCallBackend(CeedElemRestrictionHasBackendStrides(rstr, &has_backend_strides));
526       if (has_backend_strides) {
527         CeedCallBackend(CeedElemRestrictionSetLLayout(rstr, layout));
528       }
529     }
530   }
531 
532   // Set up device offset/orientation arrays
533   if (rstr_type != CEED_RESTRICTION_STRIDED) {
534     switch (mem_type) {
535       case CEED_MEM_HOST: {
536         switch (copy_mode) {
537           case CEED_OWN_POINTER:
538             impl->h_ind_allocated = (CeedInt *)indices;
539             impl->h_ind           = (CeedInt *)indices;
540             break;
541           case CEED_USE_POINTER:
542             impl->h_ind = (CeedInt *)indices;
543             break;
544           case CEED_COPY_VALUES:
545             CeedCallBackend(CeedMalloc(size, &impl->h_ind_allocated));
546             memcpy(impl->h_ind_allocated, indices, size * sizeof(CeedInt));
547             impl->h_ind = impl->h_ind_allocated;
548             break;
549         }
550         CeedCallHip(ceed, hipMalloc((void **)&impl->d_ind, size * sizeof(CeedInt)));
551         impl->d_ind_allocated = impl->d_ind;  // We own the device memory
552         CeedCallHip(ceed, hipMemcpy(impl->d_ind, indices, size * sizeof(CeedInt), hipMemcpyHostToDevice));
553         if (is_deterministic) CeedCallBackend(CeedElemRestrictionOffset_Hip(rstr, indices));
554       } break;
555       case CEED_MEM_DEVICE: {
556         switch (copy_mode) {
557           case CEED_COPY_VALUES:
558             CeedCallHip(ceed, hipMalloc((void **)&impl->d_ind, size * sizeof(CeedInt)));
559             impl->d_ind_allocated = impl->d_ind;  // We own the device memory
560             CeedCallHip(ceed, hipMemcpy(impl->d_ind, indices, size * sizeof(CeedInt), hipMemcpyDeviceToDevice));
561             break;
562           case CEED_OWN_POINTER:
563             impl->d_ind           = (CeedInt *)indices;
564             impl->d_ind_allocated = impl->d_ind;
565             break;
566           case CEED_USE_POINTER:
567             impl->d_ind = (CeedInt *)indices;
568             break;
569         }
570         CeedCallBackend(CeedMalloc(size, &impl->h_ind_allocated));
571         CeedCallHip(ceed, hipMemcpy(impl->h_ind_allocated, impl->d_ind, size * sizeof(CeedInt), hipMemcpyDeviceToHost));
572         impl->h_ind = impl->h_ind_allocated;
573         if (is_deterministic) CeedCallBackend(CeedElemRestrictionOffset_Hip(rstr, indices));
574       } break;
575     }
576 
577     // Orientation data
578     if (rstr_type == CEED_RESTRICTION_ORIENTED) {
579       switch (mem_type) {
580         case CEED_MEM_HOST: {
581           switch (copy_mode) {
582             case CEED_OWN_POINTER:
583               impl->h_orients_allocated = (bool *)orients;
584               impl->h_orients           = (bool *)orients;
585               break;
586             case CEED_USE_POINTER:
587               impl->h_orients = (bool *)orients;
588               break;
589             case CEED_COPY_VALUES:
590               CeedCallBackend(CeedMalloc(size, &impl->h_orients_allocated));
591               memcpy(impl->h_orients_allocated, orients, size * sizeof(bool));
592               impl->h_orients = impl->h_orients_allocated;
593               break;
594           }
595           CeedCallHip(ceed, hipMalloc((void **)&impl->d_orients, size * sizeof(bool)));
596           impl->d_orients_allocated = impl->d_orients;  // We own the device memory
597           CeedCallHip(ceed, hipMemcpy(impl->d_orients, orients, size * sizeof(bool), hipMemcpyHostToDevice));
598         } break;
599         case CEED_MEM_DEVICE: {
600           switch (copy_mode) {
601             case CEED_COPY_VALUES:
602               CeedCallHip(ceed, hipMalloc((void **)&impl->d_orients, size * sizeof(bool)));
603               impl->d_orients_allocated = impl->d_orients;  // We own the device memory
604               CeedCallHip(ceed, hipMemcpy(impl->d_orients, orients, size * sizeof(bool), hipMemcpyDeviceToDevice));
605               break;
606             case CEED_OWN_POINTER:
607               impl->d_orients           = (bool *)orients;
608               impl->d_orients_allocated = impl->d_orients;
609               break;
610             case CEED_USE_POINTER:
611               impl->d_orients = (bool *)orients;
612               break;
613           }
614           CeedCallBackend(CeedMalloc(size, &impl->h_orients_allocated));
615           CeedCallHip(ceed, hipMemcpy(impl->h_orients_allocated, impl->d_orients, size * sizeof(bool), hipMemcpyDeviceToHost));
616           impl->h_orients = impl->h_orients_allocated;
617         } break;
618       }
619     } else if (rstr_type == CEED_RESTRICTION_CURL_ORIENTED) {
620       switch (mem_type) {
621         case CEED_MEM_HOST: {
622           switch (copy_mode) {
623             case CEED_OWN_POINTER:
624               impl->h_curl_orients_allocated = (CeedInt8 *)curl_orients;
625               impl->h_curl_orients           = (CeedInt8 *)curl_orients;
626               break;
627             case CEED_USE_POINTER:
628               impl->h_curl_orients = (CeedInt8 *)curl_orients;
629               break;
630             case CEED_COPY_VALUES:
631               CeedCallBackend(CeedMalloc(3 * size, &impl->h_curl_orients_allocated));
632               memcpy(impl->h_curl_orients_allocated, curl_orients, 3 * size * sizeof(CeedInt8));
633               impl->h_curl_orients = impl->h_curl_orients_allocated;
634               break;
635           }
636           CeedCallHip(ceed, hipMalloc((void **)&impl->d_curl_orients, 3 * size * sizeof(CeedInt8)));
637           impl->d_curl_orients_allocated = impl->d_curl_orients;  // We own the device memory
638           CeedCallHip(ceed, hipMemcpy(impl->d_curl_orients, curl_orients, 3 * size * sizeof(CeedInt8), hipMemcpyHostToDevice));
639         } break;
640         case CEED_MEM_DEVICE: {
641           switch (copy_mode) {
642             case CEED_COPY_VALUES:
643               CeedCallHip(ceed, hipMalloc((void **)&impl->d_curl_orients, 3 * size * sizeof(CeedInt8)));
644               impl->d_curl_orients_allocated = impl->d_curl_orients;  // We own the device memory
645               CeedCallHip(ceed, hipMemcpy(impl->d_curl_orients, curl_orients, 3 * size * sizeof(CeedInt8), hipMemcpyDeviceToDevice));
646               break;
647             case CEED_OWN_POINTER:
648               impl->d_curl_orients           = (CeedInt8 *)curl_orients;
649               impl->d_curl_orients_allocated = impl->d_curl_orients;
650               break;
651             case CEED_USE_POINTER:
652               impl->d_curl_orients = (CeedInt8 *)curl_orients;
653               break;
654           }
655           CeedCallBackend(CeedMalloc(3 * size, &impl->h_curl_orients_allocated));
656           CeedCallHip(ceed, hipMemcpy(impl->h_curl_orients_allocated, impl->d_curl_orients, 3 * size * sizeof(CeedInt8), hipMemcpyDeviceToHost));
657           impl->h_curl_orients = impl->h_curl_orients_allocated;
658         } break;
659       }
660     }
661   }
662 
663   // Register backend functions
664   CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "Apply", CeedElemRestrictionApply_Hip));
665   CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "ApplyUnsigned", CeedElemRestrictionApplyUnsigned_Hip));
666   CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "ApplyUnoriented", CeedElemRestrictionApplyUnoriented_Hip));
667   CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetOffsets", CeedElemRestrictionGetOffsets_Hip));
668   CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetOrientations", CeedElemRestrictionGetOrientations_Hip));
669   CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "GetCurlOrientations", CeedElemRestrictionGetCurlOrientations_Hip));
670   CeedCallBackend(CeedSetBackendFunction(ceed, "ElemRestriction", rstr, "Destroy", CeedElemRestrictionDestroy_Hip));
671   return CEED_ERROR_SUCCESS;
672 }
673 
674 //------------------------------------------------------------------------------
675