xref: /libCEED/include/ceed/jit-source/cuda/cuda-ref-operator-assemble.h (revision f8608ea82c72806ce37b46314218ab211cba2735)
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 CUDA operator full assembly
10 #ifndef CEED_CUDA_REF_OPERATOR_ASSEMBLE_H
11 #define CEED_CUDA_REF_OPERATOR_ASSEMBLE_H
12 
13 #include <ceed.h>
14 
15 #if USE_CEEDSIZE
16 typedef CeedSize IndexType;
17 #else
18 typedef CeedInt IndexType;
19 #endif
20 
21 //------------------------------------------------------------------------------
22 // Matrix assembly kernel
23 //------------------------------------------------------------------------------
24 extern "C" __launch_bounds__(BLOCK_SIZE) __global__
25     void LinearAssemble(const CeedInt num_elem, const CeedScalar *B_in, const CeedScalar *B_out, const bool *orients_in,
26                         const CeedInt8 *curl_orients_in, const bool *orients_out, const CeedInt8 *curl_orients_out,
27                         const CeedScalar *__restrict__ qf_array, CeedScalar *__restrict__ values_array) {
28   extern __shared__ CeedScalar s_CT[];
29   CeedScalar                  *s_C = s_CT + NUM_NODES_OUT * NUM_NODES_IN;
30 
31   const int l = threadIdx.x;  // The output column index of each B^T D B operation
32                               // such that we have (Bout^T)_ij D_jk Bin_kl = C_il
33 
34   // Strides for final output ordering, determined by the reference (interface) implementation of the symbolic assembly, slowest --> fastest: e,
35   // comp_in, comp_out, node_row, node_col
36   const IndexType comp_out_stride = NUM_NODES_OUT * NUM_NODES_IN;
37   const IndexType comp_in_stride  = comp_out_stride * NUM_COMP_OUT;
38   const IndexType e_stride        = comp_in_stride * NUM_COMP_IN;
39 
40   // Strides for QF array, slowest --> fastest: e_in, comp_in, e_out, comp_out, e, q
41   const IndexType q_e_stride             = NUM_QPTS;
42   const IndexType q_comp_out_stride      = num_elem * q_e_stride;
43   const IndexType q_eval_mode_out_stride = q_comp_out_stride * NUM_COMP_OUT;
44   const IndexType q_comp_in_stride       = q_eval_mode_out_stride * NUM_EVAL_MODES_OUT;
45   const IndexType q_eval_mode_in_stride  = q_comp_in_stride * NUM_COMP_IN;
46 
47   // Loop over each element (if necessary)
48   for (IndexType e = blockIdx.x * blockDim.z + threadIdx.z; e < num_elem; e += gridDim.x * blockDim.z) {
49     for (IndexType comp_in = 0; comp_in < NUM_COMP_IN; comp_in++) {
50       for (IndexType comp_out = 0; comp_out < NUM_COMP_OUT; comp_out++) {
51         for (IndexType i = threadIdx.y; i < NUM_NODES_OUT; i += BLOCK_SIZE_Y) {
52           CeedScalar result        = 0.0;
53           IndexType  qf_index_comp = q_comp_in_stride * comp_in + q_comp_out_stride * comp_out + q_e_stride * e;
54 
55           for (IndexType e_in = 0; e_in < NUM_EVAL_MODES_IN; e_in++) {
56             IndexType b_in_index = e_in * NUM_QPTS * NUM_NODES_IN;
57 
58             for (IndexType e_out = 0; e_out < NUM_EVAL_MODES_OUT; e_out++) {
59               IndexType b_out_index = e_out * NUM_QPTS * NUM_NODES_OUT;
60               IndexType qf_index    = qf_index_comp + q_eval_mode_out_stride * e_out + q_eval_mode_in_stride * e_in;
61 
62               // Perform the B^T D B operation for this 'chunk' of D (the qf_array)
63               for (IndexType j = 0; j < NUM_QPTS; j++) {
64                 result += B_out[b_out_index + j * NUM_NODES_OUT + i] * qf_array[qf_index + j] * B_in[b_in_index + j * NUM_NODES_IN + l];
65               }
66             }  // end of out eval mode
67           }    // end of in eval mode
68           if (orients_in) {
69             result *= orients_in[NUM_NODES_IN * e + l] ? -1.0 : 1.0;
70           }
71           if (orients_out) {
72             result *= orients_out[NUM_NODES_OUT * e + i] ? -1.0 : 1.0;
73           }
74           if (!curl_orients_in && !curl_orients_out) {
75             IndexType val_index = e_stride * e + comp_in_stride * comp_in + comp_out_stride * comp_out + NUM_NODES_IN * i + l;
76 
77             values_array[val_index] = result;
78           } else if (curl_orients_in) {
79             s_C[NUM_NODES_IN * threadIdx.y + l] = result;
80             __syncthreads();
81             s_CT[NUM_NODES_IN * i + l] =
82                 (l > 0 ? s_C[NUM_NODES_IN * threadIdx.y + l - 1] * curl_orients_in[3 * NUM_NODES_IN * e + 3 * l - 1] : 0.0) +
83                 s_C[NUM_NODES_IN * threadIdx.y + l] * curl_orients_in[3 * NUM_NODES_IN * e + 3 * l + 1] +
84                 (l < (NUM_NODES_IN - 1) ? s_C[NUM_NODES_IN * threadIdx.y + l + 1] * curl_orients_in[3 * NUM_NODES_IN * e + 3 * l + 3] : 0.0);
85           } else {
86             s_CT[NUM_NODES_IN * i + l] = result;
87           }
88         }  // end of loop over element node index, i
89         if (curl_orients_in || curl_orients_out) {
90           // Compute and store the final T^T (B^T D B T) using the fully computed C T product in shared memory
91           if (curl_orients_out) __syncthreads();
92           for (IndexType i = threadIdx.y; i < NUM_NODES_OUT; i += BLOCK_SIZE_Y) {
93             IndexType val_index = e_stride * e + comp_in_stride * comp_in + comp_out_stride * comp_out + NUM_NODES_IN * i + l;
94 
95             if (curl_orients_out) {
96               values_array[val_index] =
97                   (i > 0 ? s_CT[NUM_NODES_IN * (i - 1) + l] * curl_orients_out[3 * NUM_NODES_OUT * e + 3 * i - 1] : 0.0) +
98                   s_CT[NUM_NODES_IN * i + l] * curl_orients_out[3 * NUM_NODES_OUT * e + 3 * i + 1] +
99                   (i < (NUM_NODES_OUT - 1) ? s_CT[NUM_NODES_IN * (i + 1) + l] * curl_orients_out[3 * NUM_NODES_OUT * e + 3 * i + 3] : 0.0);
100             } else {
101               values_array[val_index] = s_CT[NUM_NODES_IN * i + l];
102             }
103           }
104         }
105       }  // end of out component
106     }    // end of in component
107   }      // end of element loop
108 }
109 
110 //------------------------------------------------------------------------------
111 
112 #endif  // CEED_CUDA_REF_OPERATOR_ASSEMBLE_H
113