From eee2e895d7dc65262238a24f1625aa37d2757e7a Mon Sep 17 00:00:00 2001 From: Lin Jiang Date: Tue, 13 Sep 2022 16:26:58 +0800 Subject: [PATCH] [llvm] [refactor] Split LLVMCompiledData of kernels and tasks (#6019) Related issue = #5511 The compiler first compiles every offloaded task to a `LLVMCompiledTask` which contains an LLVM module, the name of the offloaded task function inside the module, and some extra information for linking. Then, The compiler links the modules in the `LLVMCompiledTask`s, the runtime modules and the struct modules used in the kernel together, and creates a `LLVMCompiledKernel` that contain the linked LLVM module and the names of the offloaded tasks. Both `LLVMCompiledTask` and `LLVMCompiledKernel` need to store the generated LLVM module and the names of the functions of the offloaded tasks inside the module. `LLVMCompiledTask` also stores additional information which will be used when linking like which SNodeTrees are used and the sizes of TLS buffers used in parallel struct for. Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- taichi/codegen/codegen.cpp | 18 ++++++------- taichi/codegen/codegen.h | 14 +++++----- taichi/codegen/cpu/codegen_cpu.cpp | 8 +++--- taichi/codegen/cpu/codegen_cpu.h | 4 +-- taichi/codegen/cuda/codegen_cuda.cpp | 7 +++-- taichi/codegen/cuda/codegen_cuda.h | 4 +-- taichi/codegen/llvm/codegen_llvm.cpp | 8 ++++-- taichi/codegen/llvm/codegen_llvm.h | 4 +-- taichi/codegen/llvm/llvm_compiled_data.h | 26 ++++++++++++++----- taichi/codegen/wasm/codegen_wasm.cpp | 12 ++++----- taichi/codegen/wasm/codegen_wasm.h | 4 +-- .../runtime/cpu/aot_module_builder_impl.cpp | 2 +- taichi/runtime/cpu/aot_module_builder_impl.h | 2 +- .../runtime/cuda/aot_module_builder_impl.cpp | 2 +- taichi/runtime/cuda/aot_module_builder_impl.h | 2 +- taichi/runtime/llvm/llvm_aot_module_builder.h | 2 +- taichi/runtime/llvm/llvm_context.cpp | 6 ++--- taichi/runtime/llvm/llvm_context.h | 4 +-- taichi/runtime/llvm/llvm_offline_cache.cpp | 2 +- taichi/runtime/llvm/llvm_offline_cache.h | 4 +-- .../program_impls/llvm/llvm_program.cpp | 2 +- .../runtime/program_impls/llvm/llvm_program.h | 2 +- tests/cpp/codegen/refine_coordinates_test.cpp | 6 ++--- 23 files changed, 80 insertions(+), 65 deletions(-) diff --git a/taichi/codegen/codegen.cpp b/taichi/codegen/codegen.cpp index e9956154b9ad9..9c0647e4ee225 100644 --- a/taichi/codegen/codegen.cpp +++ b/taichi/codegen/codegen.cpp @@ -57,7 +57,7 @@ std::unique_ptr KernelCodeGen::create(Arch arch, } #ifdef TI_WITH_LLVM -std::optional +std::optional KernelCodeGen::maybe_read_compilation_from_cache( const std::string &kernel_key) { TI_AUTO_PROF; @@ -79,13 +79,13 @@ KernelCodeGen::maybe_read_compilation_from_cache( return {std::move(cache_data.compiled_data)}; } -void KernelCodeGen::cache_module(const std::string &kernel_key, - const LLVMCompiledData &data) { +void KernelCodeGen::cache_kernel(const std::string &kernel_key, + const LLVMCompiledKernel &data) { get_llvm_program(prog)->cache_kernel(kernel_key, data, infer_launch_args(kernel)); } -LLVMCompiledData KernelCodeGen::compile_kernel_to_module() { +LLVMCompiledKernel KernelCodeGen::compile_kernel_to_module() { auto *llvm_prog = get_llvm_program(prog); auto *tlctx = llvm_prog->get_llvm_context(kernel->arch); auto &config = prog->config; @@ -97,7 +97,7 @@ LLVMCompiledData KernelCodeGen::compile_kernel_to_module() { if (res) { TI_DEBUG("Create kernel '{}' from cache (key='{}')", kernel->get_name(), kernel_key); - cache_module(kernel_key, *res); + cache_kernel(kernel_key, *res); return std::move(*res); } } @@ -110,7 +110,7 @@ LLVMCompiledData KernelCodeGen::compile_kernel_to_module() { TI_ASSERT(block); auto &offloads = block->statements; - std::vector> data(offloads.size()); + std::vector> data(offloads.size()); using TaskFunc = int32 (*)(void *); std::vector task_funcs(offloads.size()); for (int i = 0; i < offloads.size(); i++) { @@ -120,7 +120,7 @@ LLVMCompiledData KernelCodeGen::compile_kernel_to_module() { irpass::analysis::clone(offloads[i].get(), offloads[i]->get_kernel()); irpass::re_id(offload.get()); auto new_data = this->compile_task(nullptr, offload->as()); - data[i] = std::make_unique(std::move(new_data)); + data[i] = std::make_unique(std::move(new_data)); }; if (kernel->is_evaluator) { compile_func(); @@ -135,7 +135,7 @@ LLVMCompiledData KernelCodeGen::compile_kernel_to_module() { if (!kernel->is_evaluator) { TI_DEBUG("Cache kernel '{}' (key='{}')", kernel->get_name(), kernel_key); - cache_module(kernel_key, linked); + cache_kernel(kernel_key, linked); } return linked; } @@ -147,7 +147,7 @@ ModuleToFunctionConverter::ModuleToFunctionConverter( } FunctionType ModuleToFunctionConverter::convert(const Kernel *kernel, - LLVMCompiledData data) const { + LLVMCompiledKernel data) const { return convert(kernel->name, infer_launch_args(kernel), std::move(data)); } diff --git a/taichi/codegen/codegen.h b/taichi/codegen/codegen.h index bff0eb3729aa4..aafe0e04aa82d 100644 --- a/taichi/codegen/codegen.h +++ b/taichi/codegen/codegen.h @@ -33,16 +33,16 @@ class KernelCodeGen { } #ifdef TI_WITH_LLVM - virtual LLVMCompiledData compile_kernel_to_module(); + virtual LLVMCompiledKernel compile_kernel_to_module(); - virtual LLVMCompiledData compile_task( + virtual LLVMCompiledTask compile_task( std::unique_ptr &&module = nullptr, OffloadedStmt *stmt = nullptr){TI_NOT_IMPLEMENTED} - std::optional maybe_read_compilation_from_cache( + std::optional maybe_read_compilation_from_cache( const std::string &kernel_key); - void cache_module(const std::string &kernel_key, - const LLVMCompiledData &data); + void cache_kernel(const std::string &kernel_key, + const LLVMCompiledKernel &data); #endif }; @@ -57,10 +57,10 @@ class ModuleToFunctionConverter { virtual FunctionType convert(const std::string &kernel_name, const std::vector &args, - LLVMCompiledData data) const = 0; + LLVMCompiledKernel data) const = 0; virtual FunctionType convert(const Kernel *kernel, - LLVMCompiledData data) const; + LLVMCompiledKernel data) const; protected: TaichiLLVMContext *tlctx_{nullptr}; diff --git a/taichi/codegen/cpu/codegen_cpu.cpp b/taichi/codegen/cpu/codegen_cpu.cpp index 986eb43706b14..feda1ea612e70 100644 --- a/taichi/codegen/cpu/codegen_cpu.cpp +++ b/taichi/codegen/cpu/codegen_cpu.cpp @@ -234,7 +234,7 @@ std::unique_ptr KernelCodeGenCPU::make_codegen_llvm( FunctionType CPUModuleToFunctionConverter::convert( const std::string &kernel_name, const std::vector &args, - LLVMCompiledData data) const { + LLVMCompiledKernel data) const { TI_AUTO_PROF; auto jit_module = tlctx_->create_jit_module(std::move(data.module)); using TaskFunc = int32 (*)(void *); @@ -271,7 +271,7 @@ FunctionType CPUModuleToFunctionConverter::convert( }; } -LLVMCompiledData KernelCodeGenCPU::compile_task( +LLVMCompiledTask KernelCodeGenCPU::compile_task( std::unique_ptr &&module, OffloadedStmt *stmt) { TaskCodeGenCPU gen(kernel, stmt); @@ -284,10 +284,8 @@ FunctionType KernelCodeGenCPU::compile_to_function() { auto *llvm_prog = get_llvm_program(prog); auto *tlctx = llvm_prog->get_llvm_context(kernel->arch); - LLVMCompiledData data = compile_kernel_to_module(); - CPUModuleToFunctionConverter converter( tlctx, get_llvm_program(prog)->get_runtime_executor()); - return converter.convert(kernel, std::move(data)); + return converter.convert(kernel, compile_kernel_to_module()); } TLANG_NAMESPACE_END diff --git a/taichi/codegen/cpu/codegen_cpu.h b/taichi/codegen/cpu/codegen_cpu.h index 880029149468c..773a19a2f9dd8 100644 --- a/taichi/codegen/cpu/codegen_cpu.h +++ b/taichi/codegen/cpu/codegen_cpu.h @@ -23,7 +23,7 @@ class KernelCodeGenCPU : public KernelCodeGen { bool supports_offline_cache() const override { return true; } - LLVMCompiledData compile_task( + LLVMCompiledTask compile_task( std::unique_ptr &&module = nullptr, OffloadedStmt *stmt = nullptr) override; @@ -45,7 +45,7 @@ class CPUModuleToFunctionConverter : public ModuleToFunctionConverter { FunctionType convert(const std::string &kernel_name, const std::vector &args, - LLVMCompiledData data) const override; + LLVMCompiledKernel data) const override; }; #endif diff --git a/taichi/codegen/cuda/codegen_cuda.cpp b/taichi/codegen/cuda/codegen_cuda.cpp index f792d84535173..606c715dca55b 100644 --- a/taichi/codegen/cuda/codegen_cuda.cpp +++ b/taichi/codegen/cuda/codegen_cuda.cpp @@ -694,7 +694,7 @@ std::unique_ptr KernelCodeGenCUDA::make_codegen_llvm( } #endif // TI_WITH_LLVM -LLVMCompiledData KernelCodeGenCUDA::compile_task( +LLVMCompiledTask KernelCodeGenCUDA::compile_task( std::unique_ptr &&module, OffloadedStmt *stmt) { TaskCodeGenCUDA gen(kernel, stmt); @@ -706,17 +706,16 @@ FunctionType KernelCodeGenCUDA::compile_to_function() { auto *llvm_prog = get_llvm_program(prog); auto *tlctx = llvm_prog->get_llvm_context(kernel->arch); - LLVMCompiledData data = compile_kernel_to_module(); CUDAModuleToFunctionConverter converter{tlctx, llvm_prog->get_runtime_executor()}; - return converter.convert(this->kernel, std::move(data)); + return converter.convert(this->kernel, compile_kernel_to_module()); } FunctionType CUDAModuleToFunctionConverter::convert( const std::string &kernel_name, const std::vector &args, - LLVMCompiledData data) const { + LLVMCompiledKernel data) const { auto &mod = data.module; auto &tasks = data.tasks; #ifdef TI_WITH_CUDA diff --git a/taichi/codegen/cuda/codegen_cuda.h b/taichi/codegen/cuda/codegen_cuda.h index 54b618b0ad997..8fe6e1beb2b0b 100644 --- a/taichi/codegen/cuda/codegen_cuda.h +++ b/taichi/codegen/cuda/codegen_cuda.h @@ -17,7 +17,7 @@ class KernelCodeGenCUDA : public KernelCodeGen { #ifdef TI_WITH_LLVM static std::unique_ptr make_codegen_llvm(Kernel *kernel, IRNode *ir); - LLVMCompiledData compile_task( + LLVMCompiledTask compile_task( std::unique_ptr &&module = nullptr, OffloadedStmt *stmt = nullptr) override; #endif // TI_WITH_LLVM @@ -39,7 +39,7 @@ class CUDAModuleToFunctionConverter : public ModuleToFunctionConverter { FunctionType convert(const std::string &kernel_name, const std::vector &args, - LLVMCompiledData data) const override; + LLVMCompiledKernel data) const override; }; TLANG_NAMESPACE_END diff --git a/taichi/codegen/llvm/codegen_llvm.cpp b/taichi/codegen/llvm/codegen_llvm.cpp index 92c7d65141cbe..e1aa3a486397b 100644 --- a/taichi/codegen/llvm/codegen_llvm.cpp +++ b/taichi/codegen/llvm/codegen_llvm.cpp @@ -2647,7 +2647,7 @@ void TaskCodeGenLLVM::emit_to_module() { ir->accept(this); } -LLVMCompiledData TaskCodeGenLLVM::run_compilation() { +LLVMCompiledTask TaskCodeGenLLVM::run_compilation() { // Final lowering auto config = kernel->program->config; @@ -2732,11 +2732,15 @@ void TaskCodeGenLLVM::visit(FuncCallStmt *stmt) { } } -LLVMCompiledData LLVMCompiledData::clone() const { +LLVMCompiledTask LLVMCompiledTask::clone() const { return {tasks, llvm::CloneModule(*module), used_tree_ids, struct_for_tls_sizes}; } +LLVMCompiledKernel LLVMCompiledKernel::clone() const { + return {tasks, llvm::CloneModule(*module)}; +} + TLANG_NAMESPACE_END #endif // #ifdef TI_WITH_LLVM diff --git a/taichi/codegen/llvm/codegen_llvm.h b/taichi/codegen/llvm/codegen_llvm.h index 63a254c289a57..a02fcbb2c81e6 100644 --- a/taichi/codegen/llvm/codegen_llvm.h +++ b/taichi/codegen/llvm/codegen_llvm.h @@ -115,9 +115,9 @@ class TaskCodeGenLLVM : public IRVisitor, public LLVMModuleBuilder { * * After this call, `module` and `tasks` will be moved. * - * @return LLVMCompiledData + * @return LLVMCompiledTask */ - virtual LLVMCompiledData run_compilation(); + virtual LLVMCompiledTask run_compilation(); // For debugging only virtual llvm::Value *create_print(std::string tag, DataType dt, diff --git a/taichi/codegen/llvm/llvm_compiled_data.h b/taichi/codegen/llvm/llvm_compiled_data.h index a98fa7fb0e719..ccb3d66e5b8b8 100644 --- a/taichi/codegen/llvm/llvm_compiled_data.h +++ b/taichi/codegen/llvm/llvm_compiled_data.h @@ -21,15 +21,15 @@ class OffloadedTask { TI_IO_DEF(name, block_dim, grid_dim); }; -struct LLVMCompiledData { +struct LLVMCompiledTask { std::vector tasks; std::unique_ptr module{nullptr}; std::unordered_set used_tree_ids; std::unordered_set struct_for_tls_sizes; - LLVMCompiledData() = default; - LLVMCompiledData(LLVMCompiledData &&) = default; - LLVMCompiledData &operator=(LLVMCompiledData &&) = default; - LLVMCompiledData(std::vector tasks, + LLVMCompiledTask() = default; + LLVMCompiledTask(LLVMCompiledTask &&) = default; + LLVMCompiledTask &operator=(LLVMCompiledTask &&) = default; + LLVMCompiledTask(std::vector tasks, std::unique_ptr module, std::unordered_set used_tree_ids, std::unordered_set struct_for_tls_sizes) @@ -38,7 +38,21 @@ struct LLVMCompiledData { used_tree_ids(std::move(used_tree_ids)), struct_for_tls_sizes(std::move(struct_for_tls_sizes)) { } - LLVMCompiledData clone() const; + LLVMCompiledTask clone() const; + TI_IO_DEF(tasks); +}; + +struct LLVMCompiledKernel { + std::vector tasks; + std::unique_ptr module{nullptr}; + LLVMCompiledKernel() = default; + LLVMCompiledKernel(LLVMCompiledKernel &&) = default; + LLVMCompiledKernel &operator=(LLVMCompiledKernel &&) = default; + LLVMCompiledKernel(std::vector tasks, + std::unique_ptr module) + : tasks(std::move(tasks)), module(std::move(module)) { + } + LLVMCompiledKernel clone() const; TI_IO_DEF(tasks); }; diff --git a/taichi/codegen/wasm/codegen_wasm.cpp b/taichi/codegen/wasm/codegen_wasm.cpp index e5c3934366555..bb70a2debb286 100644 --- a/taichi/codegen/wasm/codegen_wasm.cpp +++ b/taichi/codegen/wasm/codegen_wasm.cpp @@ -215,7 +215,7 @@ class TaskCodeGenWASM : public TaskCodeGenLLVM { TI_ASSERT(!llvm::verifyFunction(*func, &llvm::errs())); } - LLVMCompiledData run_compilation() override { + LLVMCompiledTask run_compilation() override { // lower kernel if (!kernel->lowered()) { kernel->lower(); @@ -235,7 +235,7 @@ class TaskCodeGenWASM : public TaskCodeGenLLVM { } return func_name == offloaded_task_name; }); - LLVMCompiledData res; + LLVMCompiledTask res; res.tasks.emplace_back(offloaded_task_name); res.module = std::move(this->module); return res; @@ -255,7 +255,7 @@ FunctionType KernelCodeGenWASM::compile_to_function() { }; } -LLVMCompiledData KernelCodeGenWASM::compile_task( +LLVMCompiledTask KernelCodeGenWASM::compile_task( std::unique_ptr &&module, OffloadedStmt *stmt) { kernel->offload_to_executable(ir); @@ -281,14 +281,14 @@ LLVMCompiledData KernelCodeGenWASM::compile_task( return {name_list, std::move(gen->module), {}, {}}; } -LLVMCompiledData KernelCodeGenWASM::compile_kernel_to_module() { +LLVMCompiledKernel KernelCodeGenWASM::compile_kernel_to_module() { auto *tlctx = get_llvm_program(prog)->get_llvm_context(kernel->arch); if (!kernel->lowered()) { kernel->lower(/*to_executable=*/false); } auto res = compile_task(); - std::vector> data; - data.push_back(std::make_unique(std::move(res))); + std::vector> data; + data.push_back(std::make_unique(std::move(res))); return tlctx->link_compiled_tasks(std::move(data)); } diff --git a/taichi/codegen/wasm/codegen_wasm.h b/taichi/codegen/wasm/codegen_wasm.h index 6ae99ac291b71..f5e076dca5da2 100644 --- a/taichi/codegen/wasm/codegen_wasm.h +++ b/taichi/codegen/wasm/codegen_wasm.h @@ -20,11 +20,11 @@ class KernelCodeGenWASM : public KernelCodeGen { FunctionType compile_to_function() override; #ifdef TI_WITH_LLVM - LLVMCompiledData compile_task( + LLVMCompiledTask compile_task( std::unique_ptr &&module = nullptr, OffloadedStmt *stmt = nullptr) override; // AOT Module Gen - LLVMCompiledData compile_kernel_to_module() override; + LLVMCompiledKernel compile_kernel_to_module() override; #endif }; diff --git a/taichi/runtime/cpu/aot_module_builder_impl.cpp b/taichi/runtime/cpu/aot_module_builder_impl.cpp index 15b48e540855a..b2c81c71792a5 100644 --- a/taichi/runtime/cpu/aot_module_builder_impl.cpp +++ b/taichi/runtime/cpu/aot_module_builder_impl.cpp @@ -9,7 +9,7 @@ namespace taichi { namespace lang { namespace cpu { -LLVMCompiledData AotModuleBuilderImpl::compile_kernel(Kernel *kernel) { +LLVMCompiledKernel AotModuleBuilderImpl::compile_kernel(Kernel *kernel) { auto cgen = KernelCodeGenCPU(kernel); return cgen.compile_kernel_to_module(); } diff --git a/taichi/runtime/cpu/aot_module_builder_impl.h b/taichi/runtime/cpu/aot_module_builder_impl.h index 280812c0474e6..9bf2d3a30136b 100644 --- a/taichi/runtime/cpu/aot_module_builder_impl.h +++ b/taichi/runtime/cpu/aot_module_builder_impl.h @@ -15,7 +15,7 @@ class AotModuleBuilderImpl : public LlvmAotModuleBuilder { } private: - LLVMCompiledData compile_kernel(Kernel *kernel) override; + LLVMCompiledKernel compile_kernel(Kernel *kernel) override; }; } // namespace cpu diff --git a/taichi/runtime/cuda/aot_module_builder_impl.cpp b/taichi/runtime/cuda/aot_module_builder_impl.cpp index 42a1ab3e03db0..6669f899745be 100644 --- a/taichi/runtime/cuda/aot_module_builder_impl.cpp +++ b/taichi/runtime/cuda/aot_module_builder_impl.cpp @@ -9,7 +9,7 @@ namespace taichi { namespace lang { namespace cuda { -LLVMCompiledData AotModuleBuilderImpl::compile_kernel(Kernel *kernel) { +LLVMCompiledKernel AotModuleBuilderImpl::compile_kernel(Kernel *kernel) { auto cgen = KernelCodeGenCUDA(kernel); return cgen.compile_kernel_to_module(); } diff --git a/taichi/runtime/cuda/aot_module_builder_impl.h b/taichi/runtime/cuda/aot_module_builder_impl.h index f1050dead609a..278c403b0d7d5 100644 --- a/taichi/runtime/cuda/aot_module_builder_impl.h +++ b/taichi/runtime/cuda/aot_module_builder_impl.h @@ -15,7 +15,7 @@ class AotModuleBuilderImpl : public LlvmAotModuleBuilder { } private: - LLVMCompiledData compile_kernel(Kernel *kernel) override; + LLVMCompiledKernel compile_kernel(Kernel *kernel) override; }; } // namespace cuda diff --git a/taichi/runtime/llvm/llvm_aot_module_builder.h b/taichi/runtime/llvm/llvm_aot_module_builder.h index f3ae7b8e7d194..2efa9b496b88f 100644 --- a/taichi/runtime/llvm/llvm_aot_module_builder.h +++ b/taichi/runtime/llvm/llvm_aot_module_builder.h @@ -17,7 +17,7 @@ class LlvmAotModuleBuilder : public AotModuleBuilder { protected: void add_per_backend(const std::string &identifier, Kernel *kernel) override; - virtual LLVMCompiledData compile_kernel(Kernel *kernel) = 0; + virtual LLVMCompiledKernel compile_kernel(Kernel *kernel) = 0; void add_field_per_backend(const std::string &identifier, const SNode *rep_snode, diff --git a/taichi/runtime/llvm/llvm_context.cpp b/taichi/runtime/llvm/llvm_context.cpp index 52eeb803857d6..d8a13cff66b82 100644 --- a/taichi/runtime/llvm/llvm_context.cpp +++ b/taichi/runtime/llvm/llvm_context.cpp @@ -889,9 +889,9 @@ TaichiLLVMContext::ThreadLocalData::~ThreadLocalData() { thread_safe_llvm_context.reset(); } -LLVMCompiledData TaichiLLVMContext::link_compiled_tasks( - std::vector> data_list) { - LLVMCompiledData linked; +LLVMCompiledKernel TaichiLLVMContext::link_compiled_tasks( + std::vector> data_list) { + LLVMCompiledKernel linked; std::unordered_set used_tree_ids; std::unordered_set tls_sizes; std::unordered_set offloaded_names; diff --git a/taichi/runtime/llvm/llvm_context.h b/taichi/runtime/llvm/llvm_context.h index de915d0b77f63..9735a18acf503 100644 --- a/taichi/runtime/llvm/llvm_context.h +++ b/taichi/runtime/llvm/llvm_context.h @@ -142,8 +142,8 @@ class TaichiLLVMContext { static std::string get_struct_for_func_name(int tls_size); - LLVMCompiledData link_compiled_tasks( - std::vector> data_list); + LLVMCompiledKernel link_compiled_tasks( + std::vector> data_list); private: std::unique_ptr clone_module_to_context( diff --git a/taichi/runtime/llvm/llvm_offline_cache.cpp b/taichi/runtime/llvm/llvm_offline_cache.cpp index ed4dbab2160dc..fed16a0911d71 100644 --- a/taichi/runtime/llvm/llvm_offline_cache.cpp +++ b/taichi/runtime/llvm/llvm_offline_cache.cpp @@ -384,7 +384,7 @@ void LlvmOfflineCacheFileWriter::merge_with(LlvmOfflineCache &&data) { void LlvmOfflineCacheFileWriter::mangle_offloaded_task_name( const std::string &kernel_key, - LLVMCompiledData &compiled_data) { + LLVMCompiledKernel &compiled_data) { if (!mangled_) { for (auto &offload : compiled_data.tasks) { std::string mangled_name = diff --git a/taichi/runtime/llvm/llvm_offline_cache.h b/taichi/runtime/llvm/llvm_offline_cache.h index 0499fad6d6b43..c4eac6cc4fc07 100644 --- a/taichi/runtime/llvm/llvm_offline_cache.h +++ b/taichi/runtime/llvm/llvm_offline_cache.h @@ -25,7 +25,7 @@ struct LlvmOfflineCache { struct KernelCacheData { std::string kernel_key; std::vector args; - LLVMCompiledData compiled_data; + LLVMCompiledKernel compiled_data; // For cache cleaning std::size_t size{0}; // byte @@ -171,7 +171,7 @@ class LlvmOfflineCacheFileWriter { void merge_with(LlvmOfflineCache &&data); void mangle_offloaded_task_name(const std::string &kernel_key, - LLVMCompiledData &compiled_data); + LLVMCompiledKernel &compiled_data); LlvmOfflineCache data_; bool mangled_{false}; diff --git a/taichi/runtime/program_impls/llvm/llvm_program.cpp b/taichi/runtime/program_impls/llvm/llvm_program.cpp index 8bce3d2b1d9a7..9aea9a0801ea2 100644 --- a/taichi/runtime/program_impls/llvm/llvm_program.cpp +++ b/taichi/runtime/program_impls/llvm/llvm_program.cpp @@ -124,7 +124,7 @@ std::unique_ptr LlvmProgramImpl::make_aot_kernel(Kernel &kernel) { } void LlvmProgramImpl::cache_kernel(const std::string &kernel_key, - const LLVMCompiledData &data, + const LLVMCompiledKernel &data, std::vector &&args) { if (cache_data_->kernels.find(kernel_key) != cache_data_->kernels.end()) { return; diff --git a/taichi/runtime/program_impls/llvm/llvm_program.h b/taichi/runtime/program_impls/llvm/llvm_program.h index b48a22054439c..2c5ab68e07faf 100644 --- a/taichi/runtime/program_impls/llvm/llvm_program.h +++ b/taichi/runtime/program_impls/llvm/llvm_program.h @@ -51,7 +51,7 @@ class LlvmProgramImpl : public ProgramImpl { void materialize_snode_tree(SNodeTree *tree, uint64 *result_buffer) override; void cache_kernel(const std::string &kernel_key, - const LLVMCompiledData &data, + const LLVMCompiledKernel &data, std::vector &&args); ; diff --git a/tests/cpp/codegen/refine_coordinates_test.cpp b/tests/cpp/codegen/refine_coordinates_test.cpp index 0ca338b2088a4..8812ec64a6e4f 100644 --- a/tests/cpp/codegen/refine_coordinates_test.cpp +++ b/tests/cpp/codegen/refine_coordinates_test.cpp @@ -33,12 +33,12 @@ class InvokeRefineCoordinatesBuilder : public LLVMModuleBuilder { static FuncType build(const SNode *snode, TaichiLLVMContext *tlctx) { InvokeRefineCoordinatesBuilder mb{tlctx}; mb.run_jit(snode); - LLVMCompiledData data; + LLVMCompiledTask data; data.module = std::move(mb.module); data.used_tree_ids = std::move(mb.used_snode_tree_ids); data.tasks.emplace_back(kFuncName); - std::vector> data_list; - data_list.push_back(std::make_unique(std::move(data))); + std::vector> data_list; + data_list.push_back(std::make_unique(std::move(data))); auto linked_data = tlctx->link_compiled_tasks(std::move(data_list)); auto *jit = tlctx->create_jit_module(std::move(linked_data.module)); auto *fn = jit->lookup_function(kFuncName);