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