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