xref: /libCEED/include/ceed/jit-source/hip/hip-ref-restriction-strided.h (revision 5ebd836c59d60a2e5e1cb67f6731404c7da26f85)
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 /// @file
9 /// Internal header for HIP strided element restriction kernels
10 
11 #include <ceed.h>
12 
13 //------------------------------------------------------------------------------
14 // L-vector -> E-vector, strided
15 //------------------------------------------------------------------------------
16 extern "C" __global__ void StridedNoTranspose(const CeedScalar *__restrict__ u, CeedScalar *__restrict__ v) {
17   for (CeedInt node = blockIdx.x * blockDim.x + threadIdx.x; node < RSTR_NUM_ELEM * RSTR_ELEM_SIZE; node += blockDim.x * gridDim.x) {
18     const CeedInt loc_node = node % RSTR_ELEM_SIZE;
19     const CeedInt elem     = node / RSTR_ELEM_SIZE;
20 
21     for (CeedInt comp = 0; comp < RSTR_NUM_COMP; comp++) {
22       v[loc_node + comp * RSTR_ELEM_SIZE * RSTR_NUM_ELEM + elem * RSTR_ELEM_SIZE] =
23           u[loc_node * RSTR_STRIDE_NODES + comp * RSTR_STRIDE_COMP + elem * RSTR_STRIDE_ELEM];
24     }
25   }
26 }
27 
28 //------------------------------------------------------------------------------
29 // E-vector -> L-vector, strided
30 //------------------------------------------------------------------------------
31 extern "C" __global__ void StridedTranspose(const CeedScalar *__restrict__ u, CeedScalar *__restrict__ v) {
32   for (CeedInt node = blockIdx.x * blockDim.x + threadIdx.x; node < RSTR_NUM_ELEM * RSTR_ELEM_SIZE; node += blockDim.x * gridDim.x) {
33     const CeedInt loc_node = node % RSTR_ELEM_SIZE;
34     const CeedInt elem     = node / RSTR_ELEM_SIZE;
35 
36     for (CeedInt comp = 0; comp < RSTR_NUM_COMP; comp++) {
37       v[loc_node * RSTR_STRIDE_NODES + comp * RSTR_STRIDE_COMP + elem * RSTR_STRIDE_ELEM] +=
38           u[loc_node + comp * RSTR_ELEM_SIZE * RSTR_NUM_ELEM + elem * RSTR_ELEM_SIZE];
39     }
40   }
41 }
42