1 // Copyright (c) 2017-2024, 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