xref: /libCEED/include/ceed/jit-source/cuda/cuda-ref-operator-assemble.h (revision caee03026e6576cbf7a399c2fc51bb918c77f451)
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 #include <ceed.h>
9 
10 //------------------------------------------------------------------------------
11 // Matrix assembly kernel for low-order elements (2D thread block)
12 //------------------------------------------------------------------------------
13 extern "C" __launch_bounds__(BLOCK_SIZE) __global__
14     void linearAssemble(const CeedScalar *B_in, const CeedScalar *B_out, const CeedScalar *__restrict__ qf_array,
15                         CeedScalar *__restrict__ values_array) {
16   // This kernel assumes B_in and B_out have the same number of quadrature points and
17   // basis points.
18   // TODO: expand to more general cases
19   const int i = threadIdx.x;  // The output row index of each B^TDB operation
20   const int l = threadIdx.y;  // The output column index of each B^TDB operation
21                               // such that we have (Bout^T)_ij D_jk Bin_kl = C_il
22 
23   // Strides for final output ordering, determined by the reference (interface) implementation of
24   // the symbolic assembly, slowest --> fastest: element, comp_in, comp_out, node_row, node_col
25   const CeedInt comp_out_stride = NNODES * NNODES;
26   const CeedInt comp_in_stride  = comp_out_stride * NCOMP;
27   const CeedInt e_stride        = comp_in_stride * NCOMP;
28   // Strides for QF array, slowest --> fastest:  emode_in, comp_in, emode_out, comp_out, elem, qpt
29   const CeedInt qe_stride         = NQPTS;
30   const CeedInt qcomp_out_stride  = NELEM * qe_stride;
31   const CeedInt qemode_out_stride = qcomp_out_stride * NCOMP;
32   const CeedInt qcomp_in_stride   = qemode_out_stride * NUMEMODEOUT;
33   const CeedInt qemode_in_stride  = qcomp_in_stride * NCOMP;
34 
35   // Loop over each element (if necessary)
36   for (CeedInt e = blockIdx.x * blockDim.z + threadIdx.z; e < NELEM; e += gridDim.x * blockDim.z) {
37     for (CeedInt comp_in = 0; comp_in < NCOMP; comp_in++) {
38       for (CeedInt comp_out = 0; comp_out < NCOMP; comp_out++) {
39         CeedScalar result        = 0.0;
40         CeedInt    qf_index_comp = qcomp_in_stride * comp_in + qcomp_out_stride * comp_out + qe_stride * e;
41         for (CeedInt emode_in = 0; emode_in < NUMEMODEIN; emode_in++) {
42           CeedInt b_in_index = emode_in * NQPTS * NNODES;
43           for (CeedInt emode_out = 0; emode_out < NUMEMODEOUT; emode_out++) {
44             CeedInt b_out_index = emode_out * NQPTS * NNODES;
45             CeedInt qf_index    = qf_index_comp + qemode_out_stride * emode_out + qemode_in_stride * emode_in;
46             // Perform the B^T D B operation for this 'chunk' of D (the qf_array)
47             for (CeedInt j = 0; j < NQPTS; j++) {
48               result += B_out[b_out_index + j * NNODES + i] * qf_array[qf_index + j] * B_in[b_in_index + j * NNODES + l];
49             }
50           }  // end of emode_out
51         }    // end of emode_in
52         CeedInt val_index       = comp_in_stride * comp_in + comp_out_stride * comp_out + e_stride * e + NNODES * i + l;
53         values_array[val_index] = result;
54       }  // end of out component
55     }    // end of in component
56   }      // end of element loop
57 }
58 
59 //------------------------------------------------------------------------------
60 // Fallback kernel for larger orders (1D thread block)
61 //------------------------------------------------------------------------------
62 extern "C" __launch_bounds__(BLOCK_SIZE) __global__
63     void linearAssembleFallback(const CeedScalar *B_in, const CeedScalar *B_out, const CeedScalar *__restrict__ qf_array,
64                                 CeedScalar *__restrict__ values_array) {
65   // This kernel assumes B_in and B_out have the same number of quadrature points and
66   // basis points.
67   // TODO: expand to more general cases
68   const int l = threadIdx.x;  // The output column index of each B^TDB operation
69                               // such that we have (Bout^T)_ij D_jk Bin_kl = C_il
70 
71   // Strides for final output ordering, determined by the reference (interface) implementation of
72   // the symbolic assembly, slowest --> fastest: element, comp_in, comp_out, node_row, node_col
73   const CeedInt comp_out_stride = NNODES * NNODES;
74   const CeedInt comp_in_stride  = comp_out_stride * NCOMP;
75   const CeedInt e_stride        = comp_in_stride * NCOMP;
76   // Strides for QF array, slowest --> fastest:  emode_in, comp_in, emode_out, comp_out, elem, qpt
77   const CeedInt qe_stride         = NQPTS;
78   const CeedInt qcomp_out_stride  = NELEM * qe_stride;
79   const CeedInt qemode_out_stride = qcomp_out_stride * NCOMP;
80   const CeedInt qcomp_in_stride   = qemode_out_stride * NUMEMODEOUT;
81   const CeedInt qemode_in_stride  = qcomp_in_stride * NCOMP;
82 
83   // Loop over each element (if necessary)
84   for (CeedInt e = blockIdx.x * blockDim.z + threadIdx.z; e < NELEM; e += gridDim.x * blockDim.z) {
85     for (CeedInt comp_in = 0; comp_in < NCOMP; comp_in++) {
86       for (CeedInt comp_out = 0; comp_out < NCOMP; comp_out++) {
87         for (CeedInt i = 0; i < NNODES; i++) {
88           CeedScalar result        = 0.0;
89           CeedInt    qf_index_comp = qcomp_in_stride * comp_in + qcomp_out_stride * comp_out + qe_stride * e;
90           for (CeedInt emode_in = 0; emode_in < NUMEMODEIN; emode_in++) {
91             CeedInt b_in_index = emode_in * NQPTS * NNODES;
92             for (CeedInt emode_out = 0; emode_out < NUMEMODEOUT; emode_out++) {
93               CeedInt b_out_index = emode_out * NQPTS * NNODES;
94               CeedInt qf_index    = qf_index_comp + qemode_out_stride * emode_out + qemode_in_stride * emode_in;
95               // Perform the B^T D B operation for this 'chunk' of D (the qf_array)
96               for (CeedInt j = 0; j < NQPTS; j++) {
97                 result += B_out[b_out_index + j * NNODES + i] * qf_array[qf_index + j] * B_in[b_in_index + j * NNODES + l];
98               }
99             }  // end of emode_out
100           }    // end of emode_in
101           CeedInt val_index       = comp_in_stride * comp_in + comp_out_stride * comp_out + e_stride * e + NNODES * i + l;
102           values_array[val_index] = result;
103         }  // end of loop over element node index, i
104       }    // end of out component
105     }      // end of in component
106   }        // end of element loop
107 }
108 
109 //------------------------------------------------------------------------------
110