xref: /libCEED/include/ceed/jit-source/magma/magma-basis-grad-1d.h (revision 49a40c8a2d720db341b0b117b89656b473cbebfb)
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 /// @file
9 /// Internal header for MAGMA tensor basis gradient in 1D
10 #ifndef CEED_MAGMA_BASIS_GRAD_1D_H
11 #define CEED_MAGMA_BASIS_GRAD_1D_H
12 
13 #include "magma-common-tensor.h"
14 
15 // macros to abstract access of shared memory and reg. file
16 #define sT(i, j) sT[(j)*P + (i)]
17 
18 ////////////////////////////////////////////////////////////////////////////////
19 // grad basis action (1D)
20 template <typename T, int DIM, int NUM_COMP, int P, int Q>
21 static __device__ __inline__ void magma_grad_1d_device(const T *sT, T *sU[NUM_COMP], T *sV[NUM_COMP], const int tx) {
22   // Assumptions
23   // 1. 1D threads of size max(P,Q)
24   // 2. sU[i] is 1xP: in shared memory
25   // 3. sV[i] is 1xQ: in shared memory
26   // 4. P_roduct per component is one row (1xP) times T matrix (PxQ) => one row (1xQ)
27   // 5. Each thread computes one entry in sV[i]
28   // 6. Must sync before and after call
29   // 7. Note that the layout for U and V is different from 2D/3D problem
30 
31   if (tx < Q) {
32     for (int comp = 0; comp < NUM_COMP; comp++) {
33       T rv = 0.0;
34       for (int i = 0; i < P; i++) {
35         rv += sU[comp][i] * sT(i, tx);
36       }
37       sV[comp][tx] = rv;
38     }
39   }
40 }
41 
42 ////////////////////////////////////////////////////////////////////////////////
43 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_1D)) __global__
44     void magma_gradn_1d_kernel(const CeedScalar *dTinterp, const CeedScalar *dTgrad, const CeedScalar *dU, const int estrdU, const int cstrdU,
45                                const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) {
46   MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
47 
48   const int tx      = threadIdx.x;
49   const int ty      = threadIdx.y;
50   const int elem_id = (blockIdx.x * blockDim.y) + ty;
51 
52   if (elem_id >= nelem) return;
53 
54   CeedScalar *sU[BASIS_NUM_COMP];
55   CeedScalar *sV[BASIS_NUM_COMP];
56 
57   // shift global memory pointers by elem stride
58   dU += elem_id * estrdU;
59   dV += elem_id * estrdV;
60 
61   // assign shared memory pointers
62   CeedScalar *sT = (CeedScalar *)shared_data;
63   CeedScalar *sW = sT + BASIS_P * BASIS_Q;
64   sU[0]          = sW + ty * BASIS_NUM_COMP * (BASIS_P + BASIS_Q);
65   sV[0]          = sU[0] + (BASIS_NUM_COMP * 1 * BASIS_P);
66   for (int comp = 1; comp < BASIS_NUM_COMP; comp++) {
67     sU[comp] = sU[comp - 1] + (1 * BASIS_P);
68     sV[comp] = sV[comp - 1] + (1 * BASIS_Q);
69   }
70 
71   // read T
72   if (ty == 0) {
73     read_T_notrans_gm2sm<BASIS_P, BASIS_Q>(tx, dTgrad, sT);
74   }
75 
76   // read U
77   read_1d<CeedScalar, BASIS_P, BASIS_NUM_COMP>(dU, cstrdU, sU, tx);
78 
79   __syncthreads();
80   magma_grad_1d_device<CeedScalar, BASIS_DIM, BASIS_NUM_COMP, BASIS_P, BASIS_Q>(sT, sU, sV, tx);
81   __syncthreads();
82 
83   // write V
84   write_1d<CeedScalar, BASIS_Q, BASIS_NUM_COMP>(sV, dV, cstrdV, tx);
85 }
86 
87 ////////////////////////////////////////////////////////////////////////////////
88 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(BASIS_MAX_P_Q, MAGMA_MAXTHREADS_1D)) __global__
89     void magma_gradt_1d_kernel(const CeedScalar *dTinterp, const CeedScalar *dTgrad, const CeedScalar *dU, const int estrdU, const int cstrdU,
90                                const int dstrdU, CeedScalar *dV, const int estrdV, const int cstrdV, const int dstrdV, const int nelem) {
91   MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
92 
93   const int tx      = threadIdx.x;
94   const int ty      = threadIdx.y;
95   const int elem_id = (blockIdx.x * blockDim.y) + ty;
96 
97   if (elem_id >= nelem) return;
98 
99   CeedScalar *sU[BASIS_NUM_COMP];
100   CeedScalar *sV[BASIS_NUM_COMP];
101 
102   // shift global memory pointers by elem stride
103   dU += elem_id * estrdU;
104   dV += elem_id * estrdV;
105 
106   // assign shared memory pointers
107   CeedScalar *sT = (CeedScalar *)shared_data;
108   CeedScalar *sW = sT + BASIS_Q * BASIS_P;
109   sU[0]          = sW + ty * BASIS_NUM_COMP * (BASIS_Q + BASIS_P);
110   sV[0]          = sU[0] + (BASIS_NUM_COMP * 1 * BASIS_Q);
111   for (int comp = 1; comp < BASIS_NUM_COMP; comp++) {
112     sU[comp] = sU[comp - 1] + (1 * BASIS_Q);
113     sV[comp] = sV[comp - 1] + (1 * BASIS_P);
114   }
115 
116   // read T
117   if (ty == 0) {
118     read_T_trans_gm2sm<BASIS_Q, BASIS_P>(tx, dTgrad, sT);
119   }
120 
121   // read U
122   read_1d<CeedScalar, BASIS_Q, BASIS_NUM_COMP>(dU, cstrdU, sU, tx);
123 
124   __syncthreads();
125   magma_grad_1d_device<CeedScalar, BASIS_DIM, BASIS_NUM_COMP, BASIS_Q, BASIS_P>(sT, sU, sV, tx);
126   __syncthreads();
127 
128   // write V
129   write_1d<CeedScalar, BASIS_P, BASIS_NUM_COMP>(sV, dV, cstrdV, tx);
130 }
131 
132 #endif  // CEED_MAGMA_BASIS_GRAD_1D_H
133