xref: /libCEED/backends/hip-ref/ceed-hip-ref-qfunction-load.cpp (revision 4753b775a3a8f79e2dd83c2aab10890a6a04e913)
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 #include <ceed.h>
9 #include <ceed/backend.h>
10 #include <ceed/jit-tools.h>
11 #include <string.h>
12 
13 #include <iostream>
14 #include <sstream>
15 
16 #include "../hip/ceed-hip-common.h"
17 #include "../hip/ceed-hip-compile.h"
18 #include "ceed-hip-ref.h"
19 
20 //------------------------------------------------------------------------------
21 // Build QFunction kernel
22 //------------------------------------------------------------------------------
23 extern "C" int CeedQFunctionBuildKernel_Hip_ref(CeedQFunction qf) {
24   using std::ostringstream;
25   using std::string;
26 
27   Ceed                ceed;
28   Ceed_Hip           *ceed_Hip;
29   CeedInt             num_input_fields, num_output_fields, size;
30   CeedQFunctionField *input_fields, *output_fields;
31   CeedQFunction_Hip  *data;
32 
33   CeedCallBackend(CeedQFunctionGetCeed(qf, &ceed));
34   CeedCallBackend(CeedGetData(ceed, &ceed_Hip));
35   CeedCallBackend(CeedQFunctionGetData(qf, (void **)&data));
36 
37   // QFunction is built
38   if (data->QFunction) return CEED_ERROR_SUCCESS;
39 
40   // QFunction kernel generation
41   CeedCallBackend(CeedQFunctionGetFields(qf, &num_input_fields, &input_fields, &num_output_fields, &output_fields));
42 
43   // Build strings for final kernel
44   string        qfunction_name(data->qfunction_name);
45   string        kernel_name = "CeedKernelHipRefQFunction_" + qfunction_name;
46   ostringstream code;
47 
48   // Definitions
49   code << "// QFunction source\n";
50   code << "#include <ceed/jit-source/hip/hip-ref-qfunction.h>\n\n";
51   {
52     const char *source_path;
53 
54     CeedCallBackend(CeedQFunctionGetSourcePath(qf, &source_path));
55     CeedCheck(source_path, ceed, CEED_ERROR_BACKEND, "No QFunction source or hipFunction_t provided.");
56 
57     code << "// User QFunction source\n";
58     code << "#include \"" << source_path << "\"\n\n";
59   }
60   code << "extern \"C\" __launch_bounds__(BLOCK_SIZE)\n";
61   code << "__global__ void " << kernel_name << "(void *ctx, CeedInt Q, Fields_Hip fields) {\n";
62 
63   // Inputs
64   code << "  // Input fields\n";
65   for (CeedInt i = 0; i < num_input_fields; i++) {
66     CeedCallBackend(CeedQFunctionFieldGetSize(input_fields[i], &size));
67     code << "  const CeedInt size_input_" << i << " = " << size << ";\n";
68     code << "  CeedScalar input_" << i << "[size_input_" << i << "];\n";
69   }
70   code << "  const CeedScalar *inputs[" << CeedIntMax(num_input_fields, 1) << "];\n";
71   for (CeedInt i = 0; i < num_input_fields; i++) {
72     code << "  inputs[" << i << "] = input_" << i << ";\n";
73   }
74   code << "\n";
75 
76   // Outputs
77   code << "  // Output fields\n";
78   for (CeedInt i = 0; i < num_output_fields; i++) {
79     CeedCallBackend(CeedQFunctionFieldGetSize(output_fields[i], &size));
80     code << "  const CeedInt size_output_" << i << " = " << size << ";\n";
81     code << "  CeedScalar output_" << i << "[size_output_" << i << "];\n";
82   }
83   code << "  CeedScalar *outputs[" << CeedIntMax(num_output_fields, 1) << "];\n";
84   for (CeedInt i = 0; i < num_output_fields; i++) {
85     code << "  outputs[" << i << "] = output_" << i << ";\n";
86   }
87   code << "\n";
88 
89   // Loop over quadrature points
90   code << "  // Loop over quadrature points\n";
91   code << "  for (CeedInt q = blockIdx.x * blockDim.x + threadIdx.x; q < Q; q += blockDim.x * gridDim.x) {\n";
92 
93   // Load inputs
94   code << "    // -- Load inputs\n";
95   for (CeedInt i = 0; i < num_input_fields; i++) {
96     code << "    readQuads<size_input_" << i << ">(q, Q, fields.inputs[" << i << "], input_" << i << ");\n";
97   }
98   code << "\n";
99 
100   // QFunction
101   code << "    // -- Call QFunction\n";
102   code << "    " << qfunction_name << "(ctx, 1, inputs, outputs);\n\n";
103 
104   // Write outputs
105   code << "    // -- Write outputs\n";
106   for (CeedInt i = 0; i < num_output_fields; i++) {
107     code << "    writeQuads<size_output_" << i << ">(q, Q, output_" << i << ", fields.outputs[" << i << "]);\n";
108   }
109   code << "  }\n";
110   code << "}\n";
111 
112   // View kernel for debugging
113   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "Generated QFunction Kernels:\n");
114   CeedDebug(ceed, code.str().c_str());
115 
116   // Compile kernel
117   CeedCallBackend(CeedCompile_Hip(ceed, code.str().c_str(), &data->module, 1, "BLOCK_SIZE", ceed_Hip->opt_block_size));
118   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, kernel_name.c_str(), &data->QFunction));
119   return CEED_ERROR_SUCCESS;
120 }
121 
122 //------------------------------------------------------------------------------
123