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