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