xref: /libCEED/rust/libceed-sys/c-src/include/ceed/jit-source/magma/magma-basis-interp-2d.h (revision f80f4a748154eed4bc661c135f695b92b1bc45b9)
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 // macros to abstract access of shared memory and reg. file
9 #define sT(i, j) sT[(j)*P_ + (i)]
10 #define sTmp(i, j, ldw) sTmp[(j) * (ldw) + (i)]
11 
12 //////////////////////////////////////////////////////////////////////////////////////////
13 // interp basis action (2D)
14 template <typename T, int DIM_U, int DIM_V, int NCOMP_, int P_, int Q_, int rUsize, int rVsize>
15 static __device__ __inline__ void magma_interp_2d_device(const T *sT, magma_trans_t transT, T rU[DIM_U][NCOMP_][rUsize], T rV[DIM_V][NCOMP_][rVsize],
16                                                          const int tx, T rTmp, T *swork) {
17   // Assumptions
18   // 1. 1D threads of size max(P_,Q_)
19   // 2. input:  rU[DIM_U x NCOMP_ x rUsize] in registers (per thread)
20   // 3. output: rV[DIM_V x NCOMP_ x rVsize] in registers (per thread)
21   // 4. Two products per component
22   //  4.1 Batch P_ of (1xP_) matrices times (P_xQ_) matrix => Batch P_ of (1xQ_) matrices
23   //  4.2 Batch 1 of (Q_xP_) matrix   times (P_xQ_) matrix => (Q_xQ_) matrix
24   // 5. Each thread computes one row of the output of each product
25   // 6. Sync is recommended before and after the call
26 
27   for (int icomp = 0; icomp < NCOMP_; icomp++) {
28     // 1st product -- Batch P_ of (1xP_) matrices [reg] x (P_xQ_) [shmem] => Batch P_ of (1xQ_) matrices
29     // the batch output P_ x (1xQ_) is written on the fly to shmem
30     if (tx < P_) {
31       const int batchid = tx;
32       const int sld     = 1;
33       T        *sTmp    = swork + batchid * (1 * Q_);
34       for (int j = 0; j < Q_; j++) {
35         rTmp = 0.0;
36         for (int i = 0; i < P_; i++) {
37           rTmp += rU[0][icomp][i] * sT(i, j);
38         }
39         sTmp(0, j, sld) = rTmp;
40       }
41     }  // end of: if (tx < P_)
42     __syncthreads();
43 
44     // 2nd product -- Batch 1 of a (Q_xP_) matrix [shmem] x (P_xQ_) [shmem] => (Q_xQ_) matrix [reg]
45     if (tx < Q_) {
46       const int batchid = 0;
47       const int sld     = Q_;
48       T        *sTmp    = swork + batchid * (Q_ * P_);
49       for (int j = 0; j < Q_; j++) {
50         rTmp = 0.0;
51         for (int i = 0; i < P_; i++) {
52           rTmp += sTmp(tx, i, sld) * sT(i, j);
53         }
54         rV[0][icomp][j] += rTmp;
55       }
56     }
57     __syncthreads();
58   }
59 }
60 
61 //////////////////////////////////////////////////////////////////////////////////////////
62 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(MAXPQ, MAGMA_MAXTHREADS_2D)) __global__
63     void magma_interpn_2d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV,
64                                  const int cstrdV, const int nelem) {
65   MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
66 
67   const int     tx      = threadIdx.x;
68   const int     ty      = threadIdx.y;
69   const int     elem_id = (blockIdx.x * blockDim.y) + ty;
70   magma_trans_t transT  = MagmaNoTrans;
71 
72   if (elem_id >= nelem) return;
73 
74   CeedScalar rU[1][NCOMP][P] = {0.0};  // for a non fused operator DIM is always 1
75   CeedScalar rV[1][NCOMP][Q] = {0.0};  // for a non fused operator DIM is always 1
76   CeedScalar rTmp            = 0.0;
77 
78   // shift global memory pointers by elem stride
79   dU += elem_id * estrdU;
80   dV += elem_id * estrdV;
81 
82   // assign shared memory pointers
83   CeedScalar *sT   = (CeedScalar *)(shared_data);
84   CeedScalar *sTmp = sT + P * Q;
85   sTmp += ty * (P * MAXPQ);
86 
87   // read T
88   if (ty == 0) {
89     dread_T_gm2sm<P, Q>(tx, transT, dT, sT);
90   }
91 
92   // read U -- there is a sync at the end of this function
93   readU_2d<CeedScalar, P, 1, NCOMP, P, 0>(dU, cstrdU, rU, sTmp, tx);
94 
95   // no sync needed here -- readU_2d already syncs at the end
96   magma_interp_2d_device<CeedScalar, 1, 1, NCOMP, P, Q, P, Q>(sT, transT, rU, rV, tx, rTmp, sTmp);
97   __syncthreads();
98 
99   // write V
100   writeV_2d<CeedScalar, Q, 1, NCOMP, Q, 0>(dV, cstrdV, rV, tx);
101 }
102 
103 //////////////////////////////////////////////////////////////////////////////////////////
104 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(MAXPQ, MAGMA_MAXTHREADS_2D)) __global__
105     void magma_interpt_2d_kernel(const CeedScalar *dT, const CeedScalar *dU, const int estrdU, const int cstrdU, CeedScalar *dV, const int estrdV,
106                                  const int cstrdV, const int nelem) {
107   MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
108 
109   const int     tx      = threadIdx.x;
110   const int     ty      = threadIdx.y;
111   const int     elem_id = (blockIdx.x * blockDim.y) + ty;
112   magma_trans_t transT  = MagmaTrans;
113 
114   if (elem_id >= nelem) return;
115 
116   CeedScalar rU[1][NCOMP][Q] = {0.0};  // for a non fused operator DIM is always 1
117   CeedScalar rV[1][NCOMP][P] = {0.0};  // for a non fused operator DIM is always 1
118   CeedScalar rTmp            = 0.0;
119 
120   // shift global memory pointers by elem stride
121   dU += elem_id * estrdU;
122   dV += elem_id * estrdV;
123 
124   // assign shared memory pointers
125   CeedScalar *sT   = (CeedScalar *)(shared_data);
126   CeedScalar *sTmp = sT + Q * P;
127   sTmp += ty * (Q * MAXPQ);
128 
129   // read T
130   if (ty == 0) {
131     dread_T_gm2sm<Q, P>(tx, transT, dT, sT);
132   }
133 
134   // read V
135   readV_2d<CeedScalar, P, 1, NCOMP, P, 0>(dV, cstrdV, rV, tx);
136 
137   // read U -- there is a sync at the end of this function
138   readU_2d<CeedScalar, Q, 1, NCOMP, Q, 0>(dU, cstrdU, rU, sTmp, tx);
139 
140   // no sync needed here -- readU_2d already syncs at the end
141   magma_interp_2d_device<CeedScalar, 1, 1, NCOMP, Q, P, Q, P>(sT, transT, rU, rV, tx, rTmp, sTmp);
142   __syncthreads();
143 
144   // write V
145   writeV_2d<CeedScalar, P, 1, NCOMP, P, 0>(dV, cstrdV, rV, tx);
146 }
147