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