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