#include <ceed.h>
#include <ceed/backend.h>
#include <ceed/jit-tools.h>
#include <string.h>
#include <iostream>
#include <sstream>
#include "../hip/ceed-hip-common.h"
#include "../hip/ceed-hip-compile.h"
#include "ceed-hip-ref.h"
extern "C" int CeedQFunctionBuildKernel_Hip_ref(CeedQFunction qf) {
using std::ostringstream;
using std::string;
Ceed ceed;
char *read_write_kernel_path, *read_write_kernel_source;
Ceed_Hip *ceed_Hip;
CeedInt num_input_fields, num_output_fields, size;
CeedQFunctionField *input_fields, *output_fields;
CeedQFunction_Hip *data;
CeedQFunctionGetCeed(qf, &ceed);
CeedCallBackend(CeedGetData(ceed, &ceed_Hip));
CeedCallBackend(CeedQFunctionGetData(qf, (void **)&data));
if (data->QFunction) return CEED_ERROR_SUCCESS;
CeedCheck(data->qfunction_source, ceed, CEED_ERROR_BACKEND, "No QFunction source or hipFunction_t provided.");
CeedCallBackend(CeedQFunctionGetFields(qf, &num_input_fields, &input_fields, &num_output_fields, &output_fields));
CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-qfunction.h", &read_write_kernel_path));
CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading QFunction Read/Write Kernel Source -----\n");
CeedCallBackend(CeedLoadSourceToBuffer(ceed, read_write_kernel_path, &read_write_kernel_source));
CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading QFunction Read/Write Kernel Source Complete! -----\n");
string qfunction_source(data->qfunction_source);
string qfunction_name(data->qfunction_name);
string read_write(read_write_kernel_source);
string kernel_name = "CeedKernelHipRefQFunction_" + qfunction_name;
ostringstream code;
code << read_write;
code << qfunction_source;
code << "\n";
code << "extern \"C\" __launch_bounds__(BLOCK_SIZE)\n";
code << "__global__ void " << kernel_name << "(void *ctx, CeedInt Q, Fields_Hip fields) {\n";
code << " // Input fields\n";
for (CeedInt i = 0; i < num_input_fields; i++) {
CeedCallBackend(CeedQFunctionFieldGetSize(input_fields[i], &size));
code << " const CeedInt size_input_" << i << " = " << size << ";\n";
code << " CeedScalar input_" << i << "[size_input_" << i << "];\n";
}
code << " const CeedScalar* inputs[" << num_input_fields << "];\n";
for (CeedInt i = 0; i < num_input_fields; i++) {
code << " inputs[" << i << "] = input_" << i << ";\n";
}
code << "\n";
code << " // Output fields\n";
for (CeedInt i = 0; i < num_output_fields; i++) {
CeedCallBackend(CeedQFunctionFieldGetSize(output_fields[i], &size));
code << " const CeedInt size_output_" << i << " = " << size << ";\n";
code << " CeedScalar output_" << i << "[size_output_" << i << "];\n";
}
code << " CeedScalar* outputs[" << num_output_fields << "];\n";
for (CeedInt i = 0; i < num_output_fields; i++) {
code << " outputs[" << i << "] = output_" << i << ";\n";
}
code << "\n";
code << " // Loop over quadrature points\n";
code << " for (CeedInt q = blockIdx.x * blockDim.x + threadIdx.x; q < Q; q += blockDim.x * gridDim.x) {\n";
code << " // -- Load inputs\n";
for (CeedInt i = 0; i < num_input_fields; i++) {
code << " readQuads<size_input_" << i << ">(q, Q, fields.inputs[" << i << "], input_" << i << ");\n";
}
code << "\n";
code << " // -- Call QFunction\n";
code << " " << qfunction_name << "(ctx, 1, inputs, outputs);\n\n";
code << " // -- Write outputs\n";
for (CeedInt i = 0; i < num_output_fields; i++) {
code << " writeQuads<size_output_" << i << ">(q, Q, output_" << i << ", fields.outputs[" << i << "]);\n";
}
code << " }\n";
code << "}\n";
CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "Generated QFunction Kernels:\n");
CeedDebug(ceed, code.str().c_str());
CeedCallBackend(CeedCompile_Hip(ceed, code.str().c_str(), &data->module, 1, "BLOCK_SIZE", ceed_Hip->opt_block_size));
CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, kernel_name.c_str(), &data->QFunction));
CeedCallBackend(CeedFree(&data->qfunction_source));
CeedCallBackend(CeedFree(&read_write_kernel_path));
CeedCallBackend(CeedFree(&read_write_kernel_source));
return CEED_ERROR_SUCCESS;
}