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