xref: /libCEED/rust/libceed-sys/c-src/include/ceed/jit-source/magma/magma-basis-weight-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 //////////////////////////////////////////////////////////////////////////////////////////
9 // weight basis action -- 2D
10 template <typename T, int DIM_, int NCOMP_, int Q_, int iDIM, int iCOMP>
11 __device__ __inline__ void magma_weight_2d_device(const T *sTweight, T rV[DIM_][NCOMP_][Q_], const int tx) {
12   // Assumptions
13   // 1. 1D thread configuration of size Q_
14   // 2. rV[][][] matches the storage used in other actions (interp, grad, ... etc)
15   // 3. iDIM and iCOMP specify which indexes to use in rV,
16   //    since the output per thread is a register array of size Q_
17   // 4. Sync is recommended after the call (to make sure sTweight can be overwritten)
18 
19   if (tx < Q_) {
20     // x sTweight[j]  for first update
21     // x sTweight[tx] for second update
22     for (int j = 0; j < Q_; j++) {
23       rV[iDIM][iCOMP][j] = sTweight[j] * sTweight[tx];
24     }
25   }
26 }
27 
28 //////////////////////////////////////////////////////////////////////////////////////////
29 extern "C" __launch_bounds__(MAGMA_BASIS_BOUNDS(Q, MAGMA_MAXTHREADS_2D)) __global__
30     void magma_weight_2d_kernel(const CeedScalar *dqweight1d, CeedScalar *dV, const int v_stride, const int nelem) {
31   MAGMA_DEVICE_SHARED(CeedScalar, shared_data)
32 
33   const int tx      = threadIdx.x;
34   const int ty      = threadIdx.y;
35   const int elem_id = (blockIdx.x * blockDim.y) + ty;
36 
37   if (elem_id >= nelem) return;
38 
39   CeedScalar rV[1][1][Q];  // allocate with DIM=NCOMP=1, but sizes may differ for a fused operator
40   // global memory pointers
41   dV += elem_id * v_stride;
42 
43   // shared memory pointers
44   CeedScalar *sTweight = (CeedScalar *)shared_data;
45 
46   // read dqweight_1d
47   if (ty == 0 && tx < Q) {
48     sTweight[tx] = dqweight1d[tx];
49   }
50 
51   __syncthreads();
52   magma_weight_2d_device<CeedScalar, 1, 1, Q, 0, 0>(sTweight, rV, tx);
53 
54   // write V
55   if (tx < Q) {
56     for (int j = 0; j < Q; j++) {
57       dV[j * Q + tx] = rV[0][0][j];
58     }
59   }
60 }
61