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 #define CEED_DEBUG_COLOR 12 9 10 #include <ceed.h> 11 #include <ceed/backend.h> 12 #include <ceed/jit-tools.h> 13 #include <cuda_runtime.h> 14 15 #include <iostream> 16 #include <sstream> 17 #include <string> 18 19 #include "../cuda-ref/ceed-cuda-ref.h" 20 #include "../cuda-shared/ceed-cuda-shared.h" 21 #include "../cuda/ceed-cuda-common.h" 22 #include "../cuda/ceed-cuda-compile.h" 23 #include "ceed-cuda-gen.h" 24 25 //------------------------------------------------------------------------------ 26 // Build single operator kernel 27 //------------------------------------------------------------------------------ 28 extern "C" int CeedOperatorBuildKernel_Cuda_gen(CeedOperator op) { 29 using std::ostringstream; 30 using std::string; 31 32 bool is_setup_done, is_identity_qf; 33 struct cudaDeviceProp prop; 34 Ceed ceed; 35 Ceed_Cuda *ceed_data; 36 CeedSize l_size; 37 CeedInt Q, P_1d = 0, Q_1d = 0, elem_size, num_input_fields, num_output_fields, num_comp, dim = 1; 38 CeedEvalMode eval_mode; 39 CeedElemRestriction elem_rstr; 40 CeedElemRestriction_Cuda *rstr_data; 41 CeedBasis basis; 42 CeedBasis_Cuda_shared *basis_data; 43 CeedQFunctionField *qf_input_fields, *qf_output_fields; 44 CeedQFunction_Cuda_gen *qf_data; 45 CeedQFunction qf; 46 CeedOperatorField *op_input_fields, *op_output_fields; 47 CeedOperator_Cuda_gen *data; 48 49 CeedCallBackend(CeedOperatorIsSetupDone(op, &is_setup_done)); 50 if (is_setup_done) return CEED_ERROR_SUCCESS; 51 52 CeedCallBackend(CeedOperatorGetCeed(op, &ceed)); 53 CeedCallBackend(CeedOperatorGetData(op, &data)); 54 CeedCallBackend(CeedOperatorGetQFunction(op, &qf)); 55 CeedCallBackend(CeedQFunctionGetData(qf, &qf_data)); 56 CeedCallBackend(CeedOperatorGetNumQuadraturePoints(op, &Q)); 57 Q_1d = Q; 58 CeedCallBackend(CeedOperatorGetFields(op, &num_input_fields, &op_input_fields, &num_output_fields, &op_output_fields)); 59 CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_input_fields, NULL, &qf_output_fields)); 60 61 // TODO: put in a function? 62 // Check for restriction only identity operator 63 CeedCallBackend(CeedQFunctionIsIdentity(qf, &is_identity_qf)); 64 if (is_identity_qf) { 65 CeedEvalMode eval_mode_in, eval_mode_out; 66 67 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[0], &eval_mode_in)); 68 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[0], &eval_mode_out)); 69 CeedCheck(eval_mode_in != CEED_EVAL_NONE || eval_mode_out != CEED_EVAL_NONE, ceed, CEED_ERROR_BACKEND, 70 "Backend does not implement restriction only identity operators"); 71 } 72 73 ostringstream code; 74 75 // TODO: put in a function? 76 // Add atomicAdd function for old NVidia architectures 77 CeedCallBackend(CeedGetData(ceed, &ceed_data)); 78 CeedCallBackend(cudaGetDeviceProperties(&prop, ceed_data->device_id)); 79 if ((prop.major < 6) && (CEED_SCALAR_TYPE != CEED_SCALAR_FP32)) { 80 char *atomic_add_source; 81 const char *atomic_add_path; 82 83 CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/cuda/cuda-atomic-add-fallback.h", &atomic_add_path)); 84 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Atomic Add Source -----\n"); 85 CeedCallBackend(CeedLoadSourceToBuffer(ceed, atomic_add_path, &atomic_add_source)); 86 code << atomic_add_source; 87 CeedCallBackend(CeedFree(&atomic_add_path)); 88 CeedCallBackend(CeedFree(&atomic_add_source)); 89 } 90 91 // Load basis source files 92 // TODO: generalize to accept different device functions? 93 { 94 char *tensor_basis_kernel_source; 95 const char *tensor_basis_kernel_path; 96 97 CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/cuda/cuda-shared-basis-tensor-templates.h", &tensor_basis_kernel_path)); 98 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Tensor Basis Kernel Source -----\n"); 99 CeedCallBackend(CeedLoadSourceToBuffer(ceed, tensor_basis_kernel_path, &tensor_basis_kernel_source)); 100 code << tensor_basis_kernel_source; 101 CeedCallBackend(CeedFree(&tensor_basis_kernel_path)); 102 CeedCallBackend(CeedFree(&tensor_basis_kernel_source)); 103 } 104 { 105 char *cuda_gen_template_source; 106 const char *cuda_gen_template_path; 107 108 CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/cuda/cuda-gen-templates.h", &cuda_gen_template_path)); 109 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Cuda-Gen Template Source -----\n"); 110 CeedCallBackend(CeedLoadSourceToBuffer(ceed, cuda_gen_template_path, &cuda_gen_template_source)); 111 code << cuda_gen_template_source; 112 CeedCallBackend(CeedFree(&cuda_gen_template_path)); 113 CeedCallBackend(CeedFree(&cuda_gen_template_source)); 114 } 115 116 // Get QFunction source and name 117 string qfunction_source(qf_data->qfunction_source); 118 string qfunction_name(qf_data->qfunction_name); 119 string operator_name; 120 operator_name = "CeedKernelCudaGenOperator_" + qfunction_name; 121 122 // Find dim, P_1d, Q_1d 123 data->max_P_1d = 0; 124 for (CeedInt i = 0; i < num_input_fields; i++) { 125 CeedCallBackend(CeedOperatorFieldGetBasis(op_input_fields[i], &basis)); 126 if (basis != CEED_BASIS_NONE) { 127 bool is_tensor; 128 129 CeedCallBackend(CeedBasisGetData(basis, &basis_data)); 130 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode)); 131 132 // Collect dim, P_1d, and Q_1d 133 CeedCallBackend(CeedBasisGetDimension(basis, &dim)); 134 CeedCallBackend(CeedBasisIsTensor(basis, &is_tensor)); 135 CeedCheck(is_tensor, ceed, CEED_ERROR_BACKEND, "Backend does not implement operators with non-tensor basis"); 136 CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d)); 137 CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d)); 138 data->max_P_1d = CeedIntMax(data->max_P_1d, P_1d); 139 } 140 } 141 // Check output bases for Q_1d, dim as well 142 // The only input basis might be CEED_BASIS_NONE 143 for (CeedInt i = 0; i < num_output_fields; i++) { 144 CeedCallBackend(CeedOperatorFieldGetBasis(op_output_fields[i], &basis)); 145 if (basis != CEED_BASIS_NONE) { 146 bool is_tensor; 147 148 CeedCallBackend(CeedBasisGetData(basis, &basis_data)); 149 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode)); 150 151 // Collect Q_1d 152 CeedCallBackend(CeedBasisGetDimension(basis, &dim)); 153 CeedCallBackend(CeedBasisIsTensor(basis, &is_tensor)); 154 CeedCheck(is_tensor, ceed, CEED_ERROR_BACKEND, "Backend does not implement operators with non-tensor basis"); 155 CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d)); 156 } 157 } 158 data->dim = dim; 159 data->Q_1d = Q_1d; 160 161 // Only use 3D collocated gradient parallelization strategy when gradient is computed 162 // TODO: put in a function? 163 bool use_collograd_parallelization = false; 164 165 if (dim == 3) { 166 bool was_grad_found = false; 167 168 for (CeedInt i = 0; i < num_input_fields; i++) { 169 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode)); 170 if (eval_mode == CEED_EVAL_GRAD) { 171 CeedCallBackend(CeedOperatorFieldGetBasis(op_input_fields[i], &basis)); 172 CeedCallBackend(CeedBasisGetData(basis, &basis_data)); 173 use_collograd_parallelization = basis_data->d_collo_grad_1d && (was_grad_found ? use_collograd_parallelization : true); 174 was_grad_found = true; 175 } 176 } 177 for (CeedInt i = 0; i < num_output_fields; i++) { 178 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode)); 179 if (eval_mode == CEED_EVAL_GRAD) { 180 CeedCallBackend(CeedOperatorFieldGetBasis(op_output_fields[i], &basis)); 181 CeedCallBackend(CeedBasisGetData(basis, &basis_data)); 182 use_collograd_parallelization = basis_data->d_collo_grad_1d && (was_grad_found ? use_collograd_parallelization : true); 183 was_grad_found = true; 184 } 185 } 186 } 187 188 // Define CEED_Q_VLA 189 code << "\n#undef CEED_Q_VLA\n"; 190 if (dim != 3 || use_collograd_parallelization) { 191 code << "#define CEED_Q_VLA 1\n\n"; 192 } else { 193 code << "#define CEED_Q_VLA " << Q_1d << "\n\n"; 194 } 195 196 code << qfunction_source; 197 198 // Setup 199 code << "\n// -----------------------------------------------------------------------------\n"; 200 code << "\nextern \"C\" __global__ void " << operator_name 201 << "(CeedInt num_elem, void* ctx, FieldsInt_Cuda indices, Fields_Cuda fields, Fields_Cuda B, Fields_Cuda G, CeedScalar* W) {\n"; 202 for (CeedInt i = 0; i < num_input_fields; i++) { 203 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode)); 204 if (eval_mode != CEED_EVAL_WEIGHT) { // Skip CEED_EVAL_WEIGHT 205 code << " const CeedScalar* d_u_" << i << " = fields.inputs[" << i << "];\n"; 206 } 207 } 208 209 for (CeedInt i = 0; i < num_output_fields; i++) { 210 code << " CeedScalar* d_v_" << i << " = fields.outputs[" << i << "];\n"; 211 } 212 213 code << " const CeedInt dim = " << dim << ";\n"; 214 code << " const CeedInt Q_1d = " << Q_1d << ";\n"; 215 216 code << " extern __shared__ CeedScalar slice[];\n"; 217 // TODO put in a function? InitSharedData_Cuda? 218 code << " SharedData_Cuda data;\n"; 219 code << " data.t_id_x = threadIdx.x;\n"; 220 code << " data.t_id_y = threadIdx.y;\n"; 221 code << " data.t_id_z = threadIdx.z;\n"; 222 code << " data.t_id = threadIdx.x + threadIdx.y*blockDim.x + threadIdx.z*blockDim.y*blockDim.x;\n"; 223 code << " data.slice = slice+data.t_id_z*T_1D" << (dim > 1 ? "*T_1D" : "") << ";\n"; 224 225 code << "\n // -- Input field constants and basis data --\n"; 226 // TODO: Put in a function? 227 // Initialize constants, and matrices B and G 228 for (CeedInt i = 0; i < num_input_fields; i++) { 229 code << " // ---- Input field " << i << " ----\n"; 230 // Get elem_size, eval_mode, num_comp 231 CeedCallBackend(CeedOperatorFieldGetElemRestriction(op_input_fields[i], &elem_rstr)); 232 CeedCallBackend(CeedElemRestrictionGetElementSize(elem_rstr, &elem_size)); 233 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode)); 234 CeedCallBackend(CeedElemRestrictionGetNumComponents(elem_rstr, &num_comp)); 235 236 // Set field constants 237 if (eval_mode != CEED_EVAL_WEIGHT) { 238 CeedCallBackend(CeedOperatorFieldGetBasis(op_input_fields[i], &basis)); 239 if (basis != CEED_BASIS_NONE) { 240 CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d)); 241 code << " const CeedInt P_in_" << i << " = " << P_1d << ";\n"; 242 } else { 243 code << " const CeedInt P_in_" << i << " = " << Q_1d << ";\n"; 244 } 245 code << " const CeedInt num_comp_in_" << i << " = " << num_comp << ";\n"; 246 } 247 248 // Load basis data 249 code << " // EvalMode: " << CeedEvalModes[eval_mode] << "\n"; 250 switch (eval_mode) { 251 case CEED_EVAL_NONE: 252 break; 253 case CEED_EVAL_INTERP: 254 CeedCallBackend(CeedBasisGetData(basis, &basis_data)); 255 data->B.inputs[i] = basis_data->d_interp_1d; 256 code << " __shared__ CeedScalar s_B_in_" << i << "[" << P_1d * Q_1d << "];\n"; 257 code << " loadMatrix<P_in_" << i << ",Q_1d>(data, B.inputs[" << i << "], s_B_in_" << i << ");\n"; 258 break; 259 case CEED_EVAL_GRAD: 260 CeedCallBackend(CeedBasisGetData(basis, &basis_data)); 261 data->B.inputs[i] = basis_data->d_interp_1d; 262 code << " __shared__ CeedScalar s_B_in_" << i << "[" << P_1d * Q_1d << "];\n"; 263 code << " loadMatrix<P_in_" << i << ",Q_1d>(data, B.inputs[" << i << "], s_B_in_" << i << ");\n"; 264 if (use_collograd_parallelization) { 265 data->G.inputs[i] = basis_data->d_collo_grad_1d; 266 code << " __shared__ CeedScalar s_G_in_" << i << "[" << Q_1d * Q_1d << "];\n"; 267 code << " loadMatrix<Q_1d,Q_1d>(data, G.inputs[" << i << "], s_G_in_" << i << ");\n"; 268 } else { 269 bool has_collo_grad = basis_data->d_collo_grad_1d; 270 data->G.inputs[i] = has_collo_grad ? basis_data->d_collo_grad_1d : basis_data->d_grad_1d; 271 code << " __shared__ CeedScalar s_G_in_" << i << "[" << Q_1d * (has_collo_grad ? Q_1d : P_1d) << "];\n"; 272 code << " loadMatrix<" << (has_collo_grad ? "Q_1d" : ("P_in_" + std::to_string(i))) << ",Q_1d>(data, G.inputs[" << i << "], s_G_in_" << i 273 << ");\n"; 274 } 275 break; 276 case CEED_EVAL_WEIGHT: 277 break; // No action 278 case CEED_EVAL_DIV: 279 break; // TODO: Not implemented 280 case CEED_EVAL_CURL: 281 break; // TODO: Not implemented 282 } 283 } 284 285 code << "\n // -- Output field constants and basis data --\n"; 286 for (CeedInt i = 0; i < num_output_fields; i++) { 287 code << " // ---- Output field " << i << " ----\n"; 288 // Get elem_size, eval_mode, num_comp 289 CeedCallBackend(CeedOperatorFieldGetElemRestriction(op_output_fields[i], &elem_rstr)); 290 CeedCallBackend(CeedElemRestrictionGetElementSize(elem_rstr, &elem_size)); 291 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode)); 292 CeedCallBackend(CeedElemRestrictionGetNumComponents(elem_rstr, &num_comp)); 293 294 // Set field constants 295 CeedCallBackend(CeedOperatorFieldGetBasis(op_output_fields[i], &basis)); 296 if (basis != CEED_BASIS_NONE) { 297 CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d)); 298 code << " const CeedInt P_out_" << i << " = " << P_1d << ";\n"; 299 } else { 300 code << " const CeedInt P_out_" << i << " = " << Q_1d << ";\n"; 301 } 302 code << " const CeedInt num_comp_out_" << i << " = " << num_comp << ";\n"; 303 304 // Load basis data 305 code << " // EvalMode: " << CeedEvalModes[eval_mode] << "\n"; 306 switch (eval_mode) { 307 case CEED_EVAL_NONE: 308 break; // No action 309 case CEED_EVAL_INTERP: 310 CeedCallBackend(CeedBasisGetData(basis, &basis_data)); 311 data->B.outputs[i] = basis_data->d_interp_1d; 312 code << " __shared__ CeedScalar s_B_out_" << i << "[" << P_1d * Q_1d << "];\n"; 313 code << " loadMatrix<P_out_" << i << ",Q_1d>(data, B.outputs[" << i << "], s_B_out_" << i << ");\n"; 314 break; 315 case CEED_EVAL_GRAD: 316 CeedCallBackend(CeedBasisGetData(basis, &basis_data)); 317 data->B.outputs[i] = basis_data->d_interp_1d; 318 code << " __shared__ CeedScalar s_B_out_" << i << "[" << P_1d * Q_1d << "];\n"; 319 code << " loadMatrix<P_out_" << i << ",Q_1d>(data, B.outputs[" << i << "], s_B_out_" << i << ");\n"; 320 if (use_collograd_parallelization) { 321 data->G.outputs[i] = basis_data->d_collo_grad_1d; 322 code << " __shared__ CeedScalar s_G_out_" << i << "[" << Q_1d * Q_1d << "];\n"; 323 code << " loadMatrix<Q_1d,Q_1d>(data, G.outputs[" << i << "], s_G_out_" << i << ");\n"; 324 } else { 325 bool has_collo_grad = basis_data->d_collo_grad_1d; 326 data->G.outputs[i] = has_collo_grad ? basis_data->d_collo_grad_1d : basis_data->d_grad_1d; 327 code << " __shared__ CeedScalar s_G_out_" << i << "[" << Q_1d * (has_collo_grad ? Q_1d : P_1d) << "];\n"; 328 code << " loadMatrix<" << (has_collo_grad ? "Q_1d" : ("P_out_" + std::to_string(i))) << ",Q_1d>(data, G.outputs[" << i << "], s_G_out_" 329 << i << ");\n"; 330 } 331 break; 332 // LCOV_EXCL_START 333 case CEED_EVAL_WEIGHT: { 334 Ceed ceed; 335 336 CeedCallBackend(CeedOperatorGetCeed(op, &ceed)); 337 return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT cannot be an output evaluation mode"); 338 break; // Should not occur 339 } 340 case CEED_EVAL_DIV: 341 case CEED_EVAL_CURL: { 342 Ceed ceed; 343 344 CeedCallBackend(CeedOperatorGetCeed(op, &ceed)); 345 return CeedError(ceed, CEED_ERROR_BACKEND, "%s not supported", CeedEvalModes[eval_mode]); 346 break; // Should not occur 347 } 348 // LCOV_EXCL_STOP 349 } 350 } 351 code << "\n // -- Element loop --\n"; 352 code << " __syncthreads();\n"; 353 code << " for (CeedInt elem = blockIdx.x*blockDim.z + threadIdx.z; elem < num_elem; elem += gridDim.x*blockDim.z) {\n"; 354 // Input basis apply if needed 355 // Generate the correct eval mode code for each input 356 code << " // -- Input field restrictions and basis actions --\n"; 357 for (CeedInt i = 0; i < num_input_fields; i++) { 358 code << " // ---- Input field " << i << " ----\n"; 359 // Get elem_size, eval_mode, num_comp 360 CeedCallBackend(CeedOperatorFieldGetElemRestriction(op_input_fields[i], &elem_rstr)); 361 CeedCallBackend(CeedElemRestrictionGetElementSize(elem_rstr, &elem_size)); 362 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode)); 363 CeedCallBackend(CeedElemRestrictionGetNumComponents(elem_rstr, &num_comp)); 364 365 // TODO: put in a function? 366 // Restriction 367 if (eval_mode != CEED_EVAL_WEIGHT && !((eval_mode == CEED_EVAL_NONE) && use_collograd_parallelization)) { 368 code << " CeedScalar r_u_" << i << "[num_comp_in_" << i << "*P_in_" << i << "];\n"; 369 370 bool is_strided; 371 372 CeedCallBackend(CeedElemRestrictionIsStrided(elem_rstr, &is_strided)); 373 if (!is_strided) { 374 CeedInt comp_stride; 375 376 CeedCallBackend(CeedElemRestrictionGetLVectorSize(elem_rstr, &l_size)); 377 code << " const CeedInt l_size_in_" << i << " = " << l_size << ";\n"; 378 CeedCallBackend(CeedElemRestrictionGetCompStride(elem_rstr, &comp_stride)); 379 code << " // CompStride: " << comp_stride << "\n"; 380 CeedCallBackend(CeedElemRestrictionGetData(elem_rstr, &rstr_data)); 381 data->indices.inputs[i] = rstr_data->d_ind; 382 code << " readDofsOffset" << dim << "d<num_comp_in_" << i << ", " << comp_stride << ", P_in_" << i << ">(data, l_size_in_" << i 383 << ", elem, indices.inputs[" << i << "], d_u_" << i << ", r_u_" << i << ");\n"; 384 } else { 385 bool has_backend_strides; 386 CeedInt num_elem; 387 388 CeedCallBackend(CeedElemRestrictionHasBackendStrides(elem_rstr, &has_backend_strides)); 389 CeedCallBackend(CeedElemRestrictionGetNumElements(elem_rstr, &num_elem)); 390 CeedInt strides[3] = {1, elem_size * num_elem, elem_size}; 391 392 if (!has_backend_strides) { 393 CeedCallBackend(CeedElemRestrictionGetStrides(elem_rstr, strides)); 394 } 395 code << " // Strides: {" << strides[0] << ", " << strides[1] << ", " << strides[2] << "}\n"; 396 code << " readDofsStrided" << dim << "d<num_comp_in_" << i << ",P_in_" << i << "," << strides[0] << "," << strides[1] << "," << strides[2] 397 << ">(data, elem, d_u_" << i << ", r_u_" << i << ");\n"; 398 } 399 } 400 401 // TODO: put in a function? 402 // Basis action 403 code << " // EvalMode: " << CeedEvalModes[eval_mode] << "\n"; 404 switch (eval_mode) { 405 case CEED_EVAL_NONE: 406 if (!use_collograd_parallelization) { 407 code << " CeedScalar* r_t_" << i << " = r_u_" << i << ";\n"; 408 } 409 break; 410 case CEED_EVAL_INTERP: 411 code << " CeedScalar r_t_" << i << "[num_comp_in_" << i << "*Q_1d];\n"; 412 code << " Interp" << (dim > 1 ? "Tensor" : "") << dim << "d<num_comp_in_" << i << ",P_in_" << i << ",Q_1d>(data, r_u_" << i << ", s_B_in_" 413 << i << ", r_t_" << i << ");\n"; 414 break; 415 case CEED_EVAL_GRAD: 416 if (use_collograd_parallelization) { 417 code << " CeedScalar r_t_" << i << "[num_comp_in_" << i << "*Q_1d];\n"; 418 code << " Interp" << (dim > 1 ? "Tensor" : "") << dim << "d<num_comp_in_" << i << ",P_in_" << i << ",Q_1d>(data, r_u_" << i 419 << ", s_B_in_" << i << ", r_t_" << i << ");\n"; 420 } else { 421 CeedInt P_1d; 422 423 CeedCallBackend(CeedOperatorFieldGetBasis(op_input_fields[i], &basis)); 424 CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d)); 425 code << " CeedScalar r_t_" << i << "[num_comp_in_" << i << "*dim*Q_1d];\n"; 426 code << " Grad" << (dim > 1 ? "Tensor" : "") << (dim == 3 && Q_1d >= P_1d ? "Collocated" : "") << dim << "d<num_comp_in_" << i 427 << ",P_in_" << i << ",Q_1d>(data, r_u_" << i << ", s_B_in_" << i << ", s_G_in_" << i << ", r_t_" << i << ");\n"; 428 } 429 break; 430 case CEED_EVAL_WEIGHT: 431 code << " CeedScalar r_t_" << i << "[Q_1d];\n"; 432 CeedCallBackend(CeedOperatorFieldGetBasis(op_input_fields[i], &basis)); 433 CeedCallBackend(CeedBasisGetData(basis, &basis_data)); 434 data->W = basis_data->d_q_weight_1d; 435 code << " Weight" << (dim > 1 ? "Tensor" : "") << dim << "d<Q_1d>(data, W, r_t_" << i << ");\n"; 436 break; // No action 437 case CEED_EVAL_DIV: 438 break; // TODO: Not implemented 439 case CEED_EVAL_CURL: 440 break; // TODO: Not implemented 441 } 442 } 443 444 // TODO: put in a function + separate collograd logic 445 // Q function 446 code << "\n // -- Output field setup --\n"; 447 for (CeedInt i = 0; i < num_output_fields; i++) { 448 code << "\n // ---- Output field " << i << " ----\n"; 449 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode)); 450 if (eval_mode == CEED_EVAL_GRAD) { 451 if (use_collograd_parallelization) { 452 // Accumulator for gradient slices 453 code << " CeedScalar r_tt_" << i << "[num_comp_out_" << i << "*Q_1d];\n"; 454 code << " for (CeedInt i = 0; i < num_comp_out_" << i << "; i++) {\n"; 455 code << " for (CeedInt j = 0; j < Q_1d; ++j) {\n"; 456 code << " r_tt_" << i << "[j + i*Q_1d] = 0.0;\n"; 457 code << " }\n"; 458 code << " }\n"; 459 } else { 460 code << " CeedScalar r_tt_" << i << "[num_comp_out_" << i << "*dim*Q_1d];\n"; 461 } 462 } 463 if (eval_mode == CEED_EVAL_NONE || eval_mode == CEED_EVAL_INTERP) { 464 code << " CeedScalar r_tt_" << i << "[num_comp_out_" << i << "*Q_1d];\n"; 465 } 466 } 467 // We treat quadrature points per slice in 3d to save registers 468 if (use_collograd_parallelization) { 469 code << "\n // Note: Using planes of 3D elements\n"; 470 code << "#pragma unroll\n"; 471 code << " for (CeedInt q = 0; q < Q_1d; q++) {\n"; 472 code << " // -- Input fields --\n"; 473 for (CeedInt i = 0; i < num_input_fields; i++) { 474 code << " // ---- Input field " << i << " ----\n"; 475 // Get elem_size, eval_mode, num_comp 476 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode)); 477 // Basis action 478 code << " // EvalMode: " << CeedEvalModes[eval_mode] << "\n"; 479 switch (eval_mode) { 480 case CEED_EVAL_NONE: 481 bool is_strided; 482 483 code << " CeedScalar r_q_" << i << "[num_comp_in_" << i << "];\n"; 484 485 CeedCallBackend(CeedOperatorFieldGetElemRestriction(op_input_fields[i], &elem_rstr)); 486 CeedCallBackend(CeedElemRestrictionIsStrided(elem_rstr, &is_strided)); 487 if (!is_strided) { 488 CeedInt comp_stride; 489 490 CeedCallBackend(CeedElemRestrictionGetLVectorSize(elem_rstr, &l_size)); 491 code << " const CeedInt l_size_in_" << i << " = " << l_size << ";\n"; 492 CeedCallBackend(CeedElemRestrictionGetCompStride(elem_rstr, &comp_stride)); 493 code << " // CompStride: " << comp_stride << "\n"; 494 CeedCallBackend(CeedElemRestrictionGetData(elem_rstr, &rstr_data)); 495 data->indices.inputs[i] = rstr_data->d_ind; 496 code << " readSliceQuadsOffset" 497 << "3d<num_comp_in_" << i << ", " << comp_stride << ", Q_1d>(data, l_size_in_" << i << ", elem, q, indices.inputs[" << i << "], d_u_" 498 << i << ", r_q_" << i << ");\n"; 499 } else { 500 bool has_backend_strides; 501 CeedInt num_elem; 502 503 CeedCallBackend(CeedElemRestrictionGetElementSize(elem_rstr, &elem_size)); 504 CeedCallBackend(CeedElemRestrictionHasBackendStrides(elem_rstr, &has_backend_strides)); 505 CeedCallBackend(CeedElemRestrictionGetNumElements(elem_rstr, &num_elem)); 506 CeedInt strides[3] = {1, elem_size * num_elem, elem_size}; 507 508 if (!has_backend_strides) { 509 CeedCallBackend(CeedElemRestrictionGetStrides(elem_rstr, strides)); 510 } 511 code << " // Strides: {" << strides[0] << ", " << strides[1] << ", " << strides[2] << "}\n"; 512 code << " readSliceQuadsStrided" 513 << "3d<num_comp_in_" << i 514 << ",Q_1d" 515 "," 516 << strides[0] << "," << strides[1] << "," << strides[2] << ">(data, elem, q, d_u_" << i << ", r_q_" << i << ");\n"; 517 } 518 break; 519 case CEED_EVAL_INTERP: 520 code << " CeedScalar r_q_" << i << "[num_comp_in_" << i << "];\n"; 521 code << " for (CeedInt j = 0; j < num_comp_in_" << i << " ; ++j) {\n"; 522 code << " r_q_" << i << "[j] = r_t_" << i << "[q + j*Q_1d];\n"; 523 code << " }\n"; 524 break; 525 case CEED_EVAL_GRAD: 526 code << " CeedScalar r_q_" << i << "[num_comp_in_" << i << "*dim];\n"; 527 code << " gradCollo3d<num_comp_in_" << i << ",Q_1d>(data, q, r_t_" << i << ", s_G_in_" << i << ", r_q_" << i << ");\n"; 528 break; 529 case CEED_EVAL_WEIGHT: 530 code << " CeedScalar r_q_" << i << "[1];\n"; 531 code << " r_q_" << i << "[0] = r_t_" << i << "[q];\n"; 532 break; // No action 533 case CEED_EVAL_DIV: 534 break; // TODO: Not implemented 535 case CEED_EVAL_CURL: 536 break; // TODO: Not implemented 537 } 538 } 539 code << "\n // -- Output fields --\n"; 540 for (CeedInt i = 0; i < num_output_fields; i++) { 541 code << " // ---- Output field " << i << " ----\n"; 542 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode)); 543 // Basis action 544 switch (eval_mode) { 545 case CEED_EVAL_NONE: 546 code << " CeedScalar r_qq_" << i << "[num_comp_out_" << i << "];\n"; 547 break; // No action 548 case CEED_EVAL_INTERP: 549 code << " CeedScalar r_qq_" << i << "[num_comp_out_" << i << "];\n"; 550 break; 551 case CEED_EVAL_GRAD: 552 code << " CeedScalar r_qq_" << i << "[num_comp_out_" << i << "*dim];\n"; 553 break; 554 case CEED_EVAL_WEIGHT: 555 break; // Should not occur 556 case CEED_EVAL_DIV: 557 break; // TODO: Not implemented 558 case CEED_EVAL_CURL: 559 break; // TODO: Not implemented 560 } 561 } 562 } else { 563 code << "\n // Note: Using full elements\n"; 564 code << " // -- Input fields --\n"; 565 for (CeedInt i = 0; i < num_input_fields; i++) { 566 code << " // ---- Input field " << i << " ----\n"; 567 code << " CeedScalar* r_q_" << i << " = r_t_" << i << ";\n"; 568 } 569 code << " // -- Output fields --\n"; 570 for (CeedInt i = 0; i < num_output_fields; i++) { 571 code << " // ---- Output field " << i << " ----\n"; 572 code << " CeedScalar* r_qq_" << i << " = r_tt_" << i << ";\n"; 573 } 574 } 575 code << "\n // -- QFunction Inputs and outputs --\n"; 576 code << " CeedScalar* in[" << num_input_fields << "];\n"; 577 for (CeedInt i = 0; i < num_input_fields; i++) { 578 code << " // ---- Input field " << i << " ----\n"; 579 code << " in[" << i << "] = r_q_" << i << ";\n"; 580 } 581 code << " CeedScalar* out[" << num_output_fields << "];\n"; 582 for (CeedInt i = 0; i < num_output_fields; i++) { 583 code << " // ---- Output field " << i << " ----\n"; 584 code << " out[" << i << "] = r_qq_" << i << ";\n"; 585 } 586 code << "\n // -- Apply QFunction --\n"; 587 code << " " << qfunction_name << "(ctx, "; 588 if (dim != 3 || use_collograd_parallelization) { 589 code << "1"; 590 } else { 591 code << "Q_1d"; 592 } 593 code << ", in, out);\n"; 594 if (use_collograd_parallelization) { 595 code << " // -- Output fields --\n"; 596 for (CeedInt i = 0; i < num_output_fields; i++) { 597 code << " // ---- Output field " << i << " ----\n"; 598 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode)); 599 // Basis action 600 code << " // EvalMode: " << CeedEvalModes[eval_mode] << "\n"; 601 switch (eval_mode) { 602 case CEED_EVAL_NONE: 603 code << " for (CeedInt j = 0; j < num_comp_out_" << i << " ; ++j) {\n"; 604 code << " r_tt_" << i << "[q + j*Q_1d] = r_qq_" << i << "[j];\n"; 605 code << " }\n"; 606 break; // No action 607 case CEED_EVAL_INTERP: 608 code << " for (CeedInt j = 0; j < num_comp_out_" << i << " ; ++j) {\n"; 609 code << " r_tt_" << i << "[q + j*Q_1d] = r_qq_" << i << "[j];\n"; 610 code << " }\n"; 611 break; 612 case CEED_EVAL_GRAD: 613 code << " gradColloTranspose3d<num_comp_out_" << i << ",Q_1d>(data, q, r_qq_" << i << ", s_G_out_" << i << ", r_tt_" << i << ");\n"; 614 break; 615 case CEED_EVAL_WEIGHT: 616 break; // Should not occur 617 case CEED_EVAL_DIV: 618 break; // TODO: Not implemented 619 case CEED_EVAL_CURL: 620 break; // TODO: Not implemented 621 } 622 } 623 code << " }\n"; 624 } 625 626 // Output basis apply if needed 627 // Generate the correct eval mode code for each output 628 code << "\n // -- Output field basis action and restrictions --\n"; 629 for (CeedInt i = 0; i < num_output_fields; i++) { 630 code << " // ---- Output field " << i << " ----\n"; 631 // Get elem_size, eval_mode, num_comp 632 CeedCallBackend(CeedOperatorFieldGetElemRestriction(op_output_fields[i], &elem_rstr)); 633 CeedCallBackend(CeedElemRestrictionGetElementSize(elem_rstr, &elem_size)); 634 CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode)); 635 CeedCallBackend(CeedElemRestrictionGetNumComponents(elem_rstr, &num_comp)); 636 // TODO put in a function 637 // Basis action 638 code << " // EvalMode: " << CeedEvalModes[eval_mode] << "\n"; 639 switch (eval_mode) { 640 case CEED_EVAL_NONE: 641 code << " CeedScalar* r_v_" << i << " = r_tt_" << i << ";\n"; 642 break; // No action 643 case CEED_EVAL_INTERP: 644 code << " CeedScalar r_v_" << i << "[num_comp_out_" << i << "*P_out_" << i << "];\n"; 645 code << " InterpTranspose" << (dim > 1 ? "Tensor" : "") << dim << "d<num_comp_out_" << i << ",P_out_" << i << ",Q_1d>(data, r_tt_" << i 646 << ", s_B_out_" << i << ", r_v_" << i << ");\n"; 647 break; 648 case CEED_EVAL_GRAD: 649 code << " CeedScalar r_v_" << i << "[num_comp_out_" << i << "*P_out_" << i << "];\n"; 650 if (use_collograd_parallelization) { 651 code << " InterpTranspose" << (dim > 1 ? "Tensor" : "") << dim << "d<num_comp_out_" << i << ",P_out_" << i << ",Q_1d>(data, r_tt_" << i 652 << ", s_B_out_" << i << ", r_v_" << i << ");\n"; 653 } else { 654 CeedInt P_1d; 655 CeedCallBackend(CeedOperatorFieldGetBasis(op_output_fields[i], &basis)); 656 CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d)); 657 code << " GradTranspose" << (dim > 1 ? "Tensor" : "") << (dim == 3 && Q_1d >= P_1d ? "Collocated" : "") << dim << "d<num_comp_out_" << i 658 << ",P_out_" << i << ",Q_1d>(data, r_tt_" << i << ", s_B_out_" << i << ", s_G_out_" << i << ", r_v_" << i << ");\n"; 659 } 660 break; 661 // LCOV_EXCL_START 662 case CEED_EVAL_WEIGHT: { 663 Ceed ceed; 664 665 CeedCallBackend(CeedOperatorGetCeed(op, &ceed)); 666 return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT cannot be an output evaluation mode"); 667 break; // Should not occur 668 } 669 case CEED_EVAL_DIV: 670 case CEED_EVAL_CURL: { 671 Ceed ceed; 672 673 CeedCallBackend(CeedOperatorGetCeed(op, &ceed)); 674 return CeedError(ceed, CEED_ERROR_BACKEND, "%s not supported", CeedEvalModes[eval_mode]); 675 break; // Should not occur 676 } 677 // LCOV_EXCL_STOP 678 } 679 // TODO put in a function 680 // Restriction 681 bool is_strided; 682 CeedCallBackend(CeedElemRestrictionIsStrided(elem_rstr, &is_strided)); 683 if (!is_strided) { 684 CeedInt comp_stride; 685 686 CeedCallBackend(CeedElemRestrictionGetLVectorSize(elem_rstr, &l_size)); 687 code << " const CeedInt l_size_out_" << i << " = " << l_size << ";\n"; 688 CeedCallBackend(CeedElemRestrictionGetCompStride(elem_rstr, &comp_stride)); 689 code << " // CompStride: " << comp_stride << "\n"; 690 CeedCallBackend(CeedElemRestrictionGetData(elem_rstr, &rstr_data)); 691 data->indices.outputs[i] = rstr_data->d_ind; 692 code << " writeDofsOffset" << dim << "d<num_comp_out_" << i << ", " << comp_stride << ", P_out_" << i << ">(data, l_size_out_" << i 693 << ", elem, indices.outputs[" << i << "], r_v_" << i << ", d_v_" << i << ");\n"; 694 } else { 695 bool has_backend_strides; 696 CeedInt num_elem; 697 698 CeedCallBackend(CeedElemRestrictionHasBackendStrides(elem_rstr, &has_backend_strides)); 699 CeedCallBackend(CeedElemRestrictionGetNumElements(elem_rstr, &num_elem)); 700 CeedInt strides[3] = {1, elem_size * num_elem, elem_size}; 701 702 if (!has_backend_strides) { 703 CeedCallBackend(CeedElemRestrictionGetStrides(elem_rstr, strides)); 704 } 705 code << " // Strides: {" << strides[0] << ", " << strides[1] << ", " << strides[2] << "}\n"; 706 code << " writeDofsStrided" << dim << "d<num_comp_out_" << i << ",P_out_" << i << "," << strides[0] << "," << strides[1] << "," << strides[2] 707 << ">(data, elem, r_v_" << i << ", d_v_" << i << ");\n"; 708 } 709 } 710 711 code << " }\n"; 712 code << "}\n"; 713 code << "// -----------------------------------------------------------------------------\n\n"; 714 715 // View kernel for debugging 716 CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "Generated Operator Kernels:\n"); 717 CeedDebug(ceed, code.str().c_str()); 718 719 CeedCallBackend(CeedCompile_Cuda(ceed, code.str().c_str(), &data->module, 1, "T_1D", CeedIntMax(Q_1d, data->max_P_1d))); 720 CeedCallBackend(CeedGetKernel_Cuda(ceed, data->module, operator_name.c_str(), &data->op)); 721 722 CeedCallBackend(CeedOperatorSetSetupDone(op)); 723 return CEED_ERROR_SUCCESS; 724 } 725 726 //------------------------------------------------------------------------------ 727