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