From 186ff52a9592d51e8f93163cc4cad083bc691922 Mon Sep 17 00:00:00 2001 From: jim19930609 Date: Wed, 15 Jun 2022 18:08:38 +0800 Subject: [PATCH 1/7] [llvm] [refactor] LLVMProgramImpl code clean up: part-1 --- taichi/backends/cpu/jit_cpu.cpp | 21 +++++++++++---------- taichi/backends/cuda/jit_cuda.cpp | 18 ++++++++++-------- taichi/backends/cuda/jit_cuda.h | 9 ++++++--- taichi/jit/jit_session.cpp | 16 ++++++++++------ taichi/jit/jit_session.h | 19 ++++++++----------- taichi/llvm/llvm_context.cpp | 3 ++- taichi/llvm/llvm_context.h | 1 + taichi/llvm/llvm_program.cpp | 6 ++++-- taichi/program/kernel.cpp | 6 +----- taichi/program/program.cpp | 7 ------- 10 files changed, 53 insertions(+), 53 deletions(-) diff --git a/taichi/backends/cpu/jit_cpu.cpp b/taichi/backends/cpu/jit_cpu.cpp index 46366778d2f27..051c17fef92f9 100644 --- a/taichi/backends/cpu/jit_cpu.cpp +++ b/taichi/backends/cpu/jit_cpu.cpp @@ -40,7 +40,7 @@ #include "taichi/program/program.h" #include "taichi/jit/jit_session.h" #include "taichi/util/file_sequence_writer.h" -#include "taichi/llvm/llvm_program.h" +#include "taichi/llvm/llvm_context.h" TLANG_NAMESPACE_BEGIN @@ -102,10 +102,11 @@ class JITSessionCPU : public JITSession { SectionMemoryManager *memory_manager_; public: - JITSessionCPU(LlvmProgramImpl *llvm_prog, + JITSessionCPU(TaichiLLVMContext *tlctx, + CompileConfig *config, JITTargetMachineBuilder JTMB, DataLayout DL) - : JITSession(llvm_prog), + : JITSession(tlctx, config), object_layer_(es_, [&]() { auto smgr = std::make_unique(); @@ -148,9 +149,8 @@ class JITSessionCPU : public JITSession { dylib.addGenerator( cantFail(llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess( dl_.getGlobalPrefix()))); - auto *thread_safe_context = this->llvm_prog() - ->get_llvm_context(host_arch()) - ->get_this_thread_thread_safe_context(); + auto *thread_safe_context = + this->tlctx_->get_this_thread_thread_safe_context(); cantFail(compile_layer_.add( dylib, llvm::orc::ThreadSafeModule(std::move(M), *thread_safe_context))); @@ -210,7 +210,7 @@ void JITSessionCPU::global_optimize_module_cpu(llvm::Module *module) { TargetOptions options; options.PrintMachineCode = false; - if (this->llvm_prog()->config->fast_math) { + if (this->config_->fast_math) { options.AllowFPOpFusion = FPOpFusion::Fast; options.UnsafeFPMath = 1; options.NoInfsFPMath = 1; @@ -268,7 +268,7 @@ void JITSessionCPU::global_optimize_module_cpu(llvm::Module *module) { module_pass_manager.run(*module); } - if (this->llvm_prog()->config->print_kernel_llvm_ir_optimized) { + if (this->config_->print_kernel_llvm_ir_optimized) { if (false) { TI_INFO("Functions with > 100 instructions in optimized LLVM IR:"); TaichiLLVMContext::print_huge_functions(module); @@ -281,11 +281,12 @@ void JITSessionCPU::global_optimize_module_cpu(llvm::Module *module) { } std::unique_ptr create_llvm_jit_session_cpu( - LlvmProgramImpl *llvm_prog, + TaichiLLVMContext *tlctx, + CompileConfig *config, Arch arch) { TI_ASSERT(arch_is_cpu(arch)); auto target_info = get_host_target_info(); - return std::make_unique(llvm_prog, target_info.first, + return std::make_unique(tlctx, config, target_info.first, target_info.second); } diff --git a/taichi/backends/cuda/jit_cuda.cpp b/taichi/backends/cuda/jit_cuda.cpp index 1e336f75dead2..773dfda8e3408 100644 --- a/taichi/backends/cuda/jit_cuda.cpp +++ b/taichi/backends/cuda/jit_cuda.cpp @@ -1,5 +1,5 @@ #include "taichi/backends/cuda/jit_cuda.h" -#include "taichi/llvm/llvm_program.h" +#include "taichi/llvm/llvm_context.h" TLANG_NAMESPACE_BEGIN @@ -8,7 +8,7 @@ TLANG_NAMESPACE_BEGIN JITModule *JITSessionCUDA ::add_module(std::unique_ptr M, int max_reg) { auto ptx = compile_module_to_ptx(M); - if (this->llvm_prog()->config->print_kernel_nvptx) { + if (this->config_->print_kernel_nvptx) { static FileSequenceWriter writer("taichi_kernel_nvptx_{:04d}.ptx", "module NVPTX"); writer.write(ptx); @@ -82,7 +82,7 @@ std::string JITSessionCUDA::compile_module_to_ptx( using namespace llvm; - if (this->llvm_prog()->config->print_kernel_llvm_ir) { + if (this->config_->print_kernel_llvm_ir) { static FileSequenceWriter writer("taichi_kernel_cuda_llvm_ir_{:04d}.ll", "unoptimized LLVM IR (CUDA)"); writer.write(module.get()); @@ -104,7 +104,7 @@ std::string JITSessionCUDA::compile_module_to_ptx( TargetOptions options; options.PrintMachineCode = 0; - if (this->llvm_prog()->config->fast_math) { + if (this->config_->fast_math) { options.AllowFPOpFusion = FPOpFusion::Fast; // See NVPTXISelLowering.cpp // Setting UnsafeFPMath true will result in approximations such as @@ -207,7 +207,7 @@ std::string JITSessionCUDA::compile_module_to_ptx( module_pass_manager.run(*module); } - if (this->llvm_prog()->config->print_kernel_llvm_ir_optimized) { + if (this->config_->print_kernel_llvm_ir_optimized) { static FileSequenceWriter writer( "taichi_kernel_cuda_llvm_ir_optimized_{:04d}.ll", "optimized LLVM IR (CUDA)"); @@ -222,18 +222,20 @@ std::string JITSessionCUDA::compile_module_to_ptx( } std::unique_ptr create_llvm_jit_session_cuda( - LlvmProgramImpl *llvm_prog, + TaichiLLVMContext *tlctx, + CompileConfig *config, Arch arch) { TI_ASSERT(arch == Arch::cuda); // https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html#data-layout auto data_layout = llvm::DataLayout( "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-" "f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"); - return std::make_unique(llvm_prog, data_layout); + return std::make_unique(tlctx, config, data_layout); } #else std::unique_ptr create_llvm_jit_session_cuda( - LlvmProgramImpl *llvm_prog, + TaichiLLVMContext *tlctx, + CompileConfig *config, Arch arch) { TI_NOT_IMPLEMENTED } diff --git a/taichi/backends/cuda/jit_cuda.h b/taichi/backends/cuda/jit_cuda.h index 4c188add7763d..b34374e0df021 100644 --- a/taichi/backends/cuda/jit_cuda.h +++ b/taichi/backends/cuda/jit_cuda.h @@ -83,8 +83,10 @@ class JITSessionCUDA : public JITSession { public: llvm::DataLayout data_layout; - JITSessionCUDA(LlvmProgramImpl *llvm_prog, llvm::DataLayout data_layout) - : JITSession(llvm_prog), data_layout(data_layout) { + JITSessionCUDA(TaichiLLVMContext *tlctx, + CompileConfig *config, + llvm::DataLayout data_layout) + : JITSession(tlctx, config), data_layout(data_layout) { } JITModule *add_module(std::unique_ptr M, int max_reg) override; @@ -100,7 +102,8 @@ class JITSessionCUDA : public JITSession { #endif std::unique_ptr create_llvm_jit_session_cuda( - LlvmProgramImpl *llvm_prog, + TaichiLLVMContext *tlctx, + CompileConfig *config, Arch arch); TLANG_NAMESPACE_END diff --git a/taichi/jit/jit_session.cpp b/taichi/jit/jit_session.cpp index 83e021a288008..aa7890ad63256 100644 --- a/taichi/jit/jit_session.cpp +++ b/taichi/jit/jit_session.cpp @@ -8,25 +8,29 @@ TLANG_NAMESPACE_BEGIN #ifdef TI_WITH_LLVM std::unique_ptr create_llvm_jit_session_cpu( - LlvmProgramImpl *llvm_prog, + TaichiLLVMContext *tlctx, + CompileConfig *config, Arch arch); std::unique_ptr create_llvm_jit_session_cuda( - LlvmProgramImpl *llvm_prog, + TaichiLLVMContext *tlctx, + CompileConfig *config, Arch arch); #endif -JITSession::JITSession(LlvmProgramImpl *llvm_prog) : llvm_prog_(llvm_prog) { +JITSession::JITSession(TaichiLLVMContext *tlctx, CompileConfig *config) + : tlctx_(tlctx), config_(config) { } -std::unique_ptr JITSession::create(LlvmProgramImpl *llvm_prog, +std::unique_ptr JITSession::create(TaichiLLVMContext *tlctx, + CompileConfig *config, Arch arch) { #ifdef TI_WITH_LLVM if (arch_is_cpu(arch)) { - return create_llvm_jit_session_cpu(llvm_prog, arch); + return create_llvm_jit_session_cpu(tlctx, config, arch); } else if (arch == Arch::cuda) { #if defined(TI_WITH_CUDA) - return create_llvm_jit_session_cuda(llvm_prog, arch); + return create_llvm_jit_session_cuda(tlctx, config, arch); #else TI_NOT_IMPLEMENTED #endif diff --git a/taichi/jit/jit_session.h b/taichi/jit/jit_session.h index 7a5e141d5bd3a..ec62b5c844d03 100644 --- a/taichi/jit/jit_session.h +++ b/taichi/jit/jit_session.h @@ -11,17 +11,18 @@ TLANG_NAMESPACE_BEGIN // Backend JIT compiler for all archs -class LlvmProgramImpl; +class TaichiLLVMContext; +struct CompileConfig; class JITSession { - private: - LlvmProgramImpl *llvm_prog_; - protected: + TaichiLLVMContext *tlctx_; + CompileConfig *config_; + std::vector> modules; public: - JITSession(LlvmProgramImpl *llvm_prog); + JITSession(TaichiLLVMContext *tlctx, CompileConfig *config); virtual JITModule *add_module(std::unique_ptr M, int max_reg = 0) = 0; @@ -34,18 +35,14 @@ class JITSession { virtual llvm::DataLayout get_data_layout() = 0; - static std::unique_ptr create(LlvmProgramImpl *llvm_prog, + static std::unique_ptr create(TaichiLLVMContext *tlctx, + CompileConfig *config, Arch arch); virtual void global_optimize_module(llvm::Module *module) { } virtual ~JITSession() = default; - - protected: - LlvmProgramImpl *llvm_prog() const { - return llvm_prog_; - } }; TLANG_NAMESPACE_END diff --git a/taichi/llvm/llvm_context.cpp b/taichi/llvm/llvm_context.cpp index fa98d12245694..872763722395d 100644 --- a/taichi/llvm/llvm_context.cpp +++ b/taichi/llvm/llvm_context.cpp @@ -41,6 +41,7 @@ #include "taichi/common/task.h" #include "taichi/util/environ_config.h" #include "llvm_context.h" +#include "taichi/llvm/llvm_program.h" #ifdef _WIN32 // Travis CI seems doesn't support ... @@ -93,7 +94,7 @@ TaichiLLVMContext::TaichiLLVMContext(LlvmProgramImpl *llvm_prog, Arch arch) TI_NOT_IMPLEMENTED #endif } - jit = JITSession::create(llvm_prog, arch); + jit = JITSession::create(this, llvm_prog->config, arch); TI_TRACE("Taichi llvm context created."); } diff --git a/taichi/llvm/llvm_context.h b/taichi/llvm/llvm_context.h index 20dc820af46a3..61a1b05ef3f51 100644 --- a/taichi/llvm/llvm_context.h +++ b/taichi/llvm/llvm_context.h @@ -18,6 +18,7 @@ namespace taichi { namespace lang { class JITSessionCPU; +class LlvmProgramImpl; /** * Manages an LLVMContext for Taichi's usage. diff --git a/taichi/llvm/llvm_program.cpp b/taichi/llvm/llvm_program.cpp index a805ade265e37..9ac8844b8f44f 100644 --- a/taichi/llvm/llvm_program.cpp +++ b/taichi/llvm/llvm_program.cpp @@ -66,6 +66,7 @@ LlvmProgramImpl::LlvmProgramImpl(CompileConfig &config_, preallocated_device_buffer_ = nullptr; llvm_runtime_ = nullptr; llvm_context_host_ = std::make_unique(this, host_arch()); + if (config_.arch == Arch::cuda) { #if defined(TI_WITH_CUDA) int num_SMs{1}; @@ -114,8 +115,11 @@ LlvmProgramImpl::LlvmProgramImpl(CompileConfig &config_, } CUDAContext::get_instance().set_debug(config->debug); device_ = std::make_shared(); + + this->maybe_initialize_cuda_llvm_context(); } #endif + this->initialize_host(); } void LlvmProgramImpl::initialize_host() { @@ -366,8 +370,6 @@ std::unique_ptr LlvmProgramImpl::make_aot_module_builder() { void LlvmProgramImpl::materialize_runtime(MemoryPool *memory_pool, KernelProfilerBase *profiler, uint64 **result_buffer_ptr) { - maybe_initialize_cuda_llvm_context(); - std::size_t prealloc_size = 0; TaichiLLVMContext *tlctx = nullptr; if (config->arch == Arch::cuda) { diff --git a/taichi/program/kernel.cpp b/taichi/program/kernel.cpp index 0a16f76ae43a3..2e06995aa9987 100644 --- a/taichi/program/kernel.cpp +++ b/taichi/program/kernel.cpp @@ -415,11 +415,7 @@ void Kernel::init(Program &program, this->autodiff_mode = autodiff_mode; this->lowered_ = false; this->program = &program; -#ifdef TI_WITH_LLVM - if (auto *llvm_program_impl = program.get_llvm_program_impl()) { - llvm_program_impl->maybe_initialize_cuda_llvm_context(); - } -#endif + is_accessor = false; is_evaluator = false; compiled_ = nullptr; diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp index e6149271d38c0..6de93705647dc 100644 --- a/taichi/program/program.cpp +++ b/taichi/program/program.cpp @@ -125,13 +125,6 @@ Program::Program(Arch desired_arch) total_compilation_time_ = 0; num_instances_ += 1; SNode::counter = 0; - if (arch_uses_llvm(config.arch)) { -#if TI_WITH_LLVM - static_cast(program_impl_.get())->initialize_host(); -#else - TI_NOT_IMPLEMENTED -#endif - } result_buffer = nullptr; current_callable = nullptr; From c2e0dd06f576dcab959c8e0435ef176844806765 Mon Sep 17 00:00:00 2001 From: jim19930609 Date: Thu, 16 Jun 2022 14:43:45 +0800 Subject: [PATCH 2/7] [llvm] [refactor] LLVMProgramImpl code clean up: part-2 --- taichi/backends/cuda/codegen_cuda.cpp | 5 +-- taichi/codegen/codegen_llvm.cpp | 24 +++++++------- taichi/llvm/llvm_program.h | 8 +++-- taichi/program/kernel.cpp | 7 +---- taichi/program/ndarray.h | 1 - taichi/program/program.cpp | 31 ++++++------------- taichi/program/program.h | 7 +++-- taichi/program/program_impl.h | 17 ++++++++++ tests/cpp/codegen/refine_coordinates_test.cpp | 5 ++- tests/cpp/llvm/llvm_offline_cache_test.cpp | 4 ++- 10 files changed, 62 insertions(+), 47 deletions(-) diff --git a/taichi/backends/cuda/codegen_cuda.cpp b/taichi/backends/cuda/codegen_cuda.cpp index 9ebe952e8b978..65650ded7840d 100644 --- a/taichi/backends/cuda/codegen_cuda.cpp +++ b/taichi/backends/cuda/codegen_cuda.cpp @@ -39,8 +39,9 @@ class CodeGenLLVMCUDA : public CodeGenLLVM { FunctionType gen() override { auto compiled_res = run_compilation(); - CUDAModuleToFunctionConverter converter{ - tlctx, this->kernel->program->get_llvm_program_impl()}; + auto *llvm_prog = static_cast( + this->kernel->program->get_program_impl()); + CUDAModuleToFunctionConverter converter{tlctx, llvm_prog}; return converter.convert(this->kernel, std::move(compiled_res.llvm_module), std::move(compiled_res.offloaded_tasks)); diff --git a/taichi/codegen/codegen_llvm.cpp b/taichi/codegen/codegen_llvm.cpp index ed811d756a983..659d131249fd8 100644 --- a/taichi/codegen/codegen_llvm.cpp +++ b/taichi/codegen/codegen_llvm.cpp @@ -305,12 +305,13 @@ CodeGenLLVM::CodeGenLLVM(Kernel *kernel, std::unique_ptr &&module) // TODO: simplify LLVMModuleBuilder ctor input : LLVMModuleBuilder( - module == nullptr ? kernel->program->get_llvm_program_impl() + module == nullptr ? static_cast( + kernel->program->get_program_impl()) ->get_llvm_context(kernel->arch) ->clone_struct_module() : std::move(module), - kernel->program->get_llvm_program_impl()->get_llvm_context( - kernel->arch)), + static_cast(kernel->program->get_program_impl()) + ->get_llvm_context(kernel->arch)), kernel(kernel), ir(ir), prog(kernel->program) { @@ -2292,7 +2293,8 @@ FunctionCreationGuard CodeGenLLVM::get_function_creation_guard( } void CodeGenLLVM::initialize_context() { - tlctx = prog->get_llvm_program_impl()->get_llvm_context(kernel->arch); + tlctx = static_cast(prog->get_program_impl()) + ->get_llvm_context(kernel->arch); llvm_context = tlctx->get_this_thread_context(); builder = std::make_unique>(*llvm_context); } @@ -2399,8 +2401,8 @@ bool CodeGenLLVM::maybe_read_compilation_from_cache( } LlvmOfflineCache::KernelCacheData cache_data; - auto *tlctx = - this->prog->get_llvm_program_impl()->get_llvm_context(config.arch); + auto *tlctx = static_cast(prog->get_program_impl()) + ->get_llvm_context(config.arch); auto &llvm_ctx = *tlctx->get_this_thread_context(); if (!reader->get_kernel_cache(cache_data, kernel_key, llvm_ctx)) { @@ -2422,8 +2424,8 @@ bool CodeGenLLVM::maybe_read_compilation_from_cache( FunctionType CodeGenLLVM::gen() { auto compiled_res = run_compilation(); - ModuleToFunctionConverter converter{tlctx, - kernel->program->get_llvm_program_impl()}; + ModuleToFunctionConverter converter{ + tlctx, static_cast(prog->get_program_impl())}; return converter.convert(kernel, std::move(compiled_res.llvm_module), std::move(compiled_res.offloaded_tasks)); } @@ -2504,9 +2506,9 @@ void CodeGenLLVM::cache_module(const std::string &kernel_key) { task_cache.block_dim = task.block_dim; task_cache.grid_dim = task.grid_dim; } - prog->get_llvm_program_impl()->cache_kernel(kernel_key, this->module.get(), - infer_launch_args(kernel), - std::move(offloaded_task_list)); + static_cast(prog->get_program_impl()) + ->cache_kernel(kernel_key, this->module.get(), infer_launch_args(kernel), + std::move(offloaded_task_list)); } ModuleToFunctionConverter::ModuleToFunctionConverter(TaichiLLVMContext *tlctx, diff --git a/taichi/llvm/llvm_program.h b/taichi/llvm/llvm_program.h index 2eec64dd8e7bd..3f9576b9507db 100644 --- a/taichi/llvm/llvm_program.h +++ b/taichi/llvm/llvm_program.h @@ -64,6 +64,10 @@ class LlvmProgramImpl : public ProgramImpl { return static_cast(llvm_runtime_); } + void prepare_runtime_context(RuntimeContext *ctx) override { + ctx->runtime = get_llvm_runtime(); + } + FunctionType compile(Kernel *kernel, OffloadedStmt *offloaded) override; void compile_snode_tree_types(SNodeTree *tree) override; @@ -106,11 +110,11 @@ class LlvmProgramImpl : public ProgramImpl { DeviceAllocation allocate_memory_ndarray(std::size_t alloc_size, uint64 *result_buffer) override; - uint64_t *get_ndarray_alloc_info_ptr(const DeviceAllocation &alloc); + uint64_t *get_ndarray_alloc_info_ptr(const DeviceAllocation &alloc) override; void fill_ndarray(const DeviceAllocation &alloc, std::size_t size, - uint32_t data); + uint32_t data) override; void cache_kernel(const std::string &kernel_key, llvm::Module *module, diff --git a/taichi/program/kernel.cpp b/taichi/program/kernel.cpp index 2e06995aa9987..c6c697c01369e 100644 --- a/taichi/program/kernel.cpp +++ b/taichi/program/kernel.cpp @@ -303,12 +303,7 @@ void Kernel::LaunchContextBuilder::set_arg_raw(int arg_id, uint64 d) { } RuntimeContext &Kernel::LaunchContextBuilder::get_context() { -#ifdef TI_WITH_LLVM - if (auto *llvm_program_impl = kernel_->program->get_llvm_program_impl()) { - ctx_->runtime = llvm_program_impl->get_llvm_runtime(); - } -#endif - ctx_->result_buffer = kernel_->program->result_buffer; + kernel_->program->prepare_runtime_context(ctx_); return *ctx_; } diff --git a/taichi/program/ndarray.h b/taichi/program/ndarray.h index dae837153cfd6..5088d35da0cd6 100644 --- a/taichi/program/ndarray.h +++ b/taichi/program/ndarray.h @@ -11,7 +11,6 @@ namespace taichi { namespace lang { class Program; -class LlvmProgramImpl; class NdarrayRwAccessorsBank; class TI_DLL_EXPORT Ndarray { diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp index 6de93705647dc..d97052000d531 100644 --- a/taichi/program/program.cpp +++ b/taichi/program/program.cpp @@ -577,29 +577,21 @@ Texture *Program::create_texture(const DataType type, intptr_t Program::get_ndarray_data_ptr_as_int(const Ndarray *ndarray) { uint64_t *data_ptr{nullptr}; -#ifdef TI_WITH_LLVM if (arch_is_cpu(config.arch) || config.arch == Arch::cuda) { // For the LLVM backends, device allocation is a physical pointer. - data_ptr = get_llvm_program_impl()->get_ndarray_alloc_info_ptr( - ndarray->ndarray_alloc_); + data_ptr = + program_impl_->get_ndarray_alloc_info_ptr(ndarray->ndarray_alloc_); } -#else - TI_ERROR("Llvm disabled"); -#endif return reinterpret_cast(data_ptr); } void Program::fill_ndarray_fast(Ndarray *ndarray, uint32_t val) { -// This is a temporary solution to bypass device api. -// Should be moved to CommandList once available in CUDA. -#ifdef TI_WITH_LLVM - get_llvm_program_impl()->fill_ndarray( + // This is a temporary solution to bypass device api. + // Should be moved to CommandList once available in CUDA. + program_impl_->fill_ndarray( ndarray->ndarray_alloc_, ndarray->get_nelement() * ndarray->get_element_size(), val); -#else - TI_ERROR("Not supported"); -#endif } Program::~Program() { @@ -626,14 +618,6 @@ std::unique_ptr Program::make_aot_module_builder(Arch arch) { return nullptr; } -LlvmProgramImpl *Program::get_llvm_program_impl() { -#ifdef TI_WITH_LLVM - return static_cast(program_impl_.get()); -#else - TI_ERROR("Llvm disabled"); -#endif -} - int Program::allocate_snode_tree_id() { if (free_snode_tree_ids_.empty()) { return snode_trees_.size(); @@ -644,5 +628,10 @@ int Program::allocate_snode_tree_id() { } } +void Program::prepare_runtime_context(RuntimeContext *ctx) { + ctx->result_buffer = result_buffer; + program_impl_->prepare_runtime_context(ctx); +} + } // namespace lang } // namespace taichi diff --git a/taichi/program/program.h b/taichi/program/program.h index 5fdacd9f2f179..598a891ae8cf3 100644 --- a/taichi/program/program.h +++ b/taichi/program/program.h @@ -77,7 +77,6 @@ namespace taichi { namespace lang { class StructCompiler; -class LlvmProgramImpl; class AsyncEngine; /** @@ -301,7 +300,9 @@ class TI_DLL_EXPORT Program { std::unique_ptr make_aot_module_builder(Arch arch); - LlvmProgramImpl *get_llvm_program_impl(); + ProgramImpl *get_program_impl() { + return program_impl_.get(); + } DevicePtr get_snode_tree_device_ptr(int tree_id) { return program_impl_->get_snode_tree_device_ptr(tree_id); @@ -343,6 +344,8 @@ class TI_DLL_EXPORT Program { return Identifier(global_id_counter_++, name); } + void prepare_runtime_context(RuntimeContext *ctx); + private: uint64 ndarray_writer_counter_{0}; uint64 ndarray_reader_counter_{0}; diff --git a/taichi/program/program_impl.h b/taichi/program/program_impl.h index f0ae8d243c69a..714b128114a17 100644 --- a/taichi/program/program_impl.h +++ b/taichi/program/program_impl.h @@ -13,6 +13,8 @@ namespace taichi { namespace lang { +struct RuntimeContext; + class ProgramImpl { public: // TODO: Make it safer, we exposed it for now as it's directly accessed @@ -99,6 +101,21 @@ class ProgramImpl { virtual ~ProgramImpl() { } + virtual uint64_t *get_ndarray_alloc_info_ptr(const DeviceAllocation &alloc) { + TI_ERROR( + "get_ndarray_alloc_info_ptr() not implemented on the current backend"); + return nullptr; + } + + virtual void fill_ndarray(const DeviceAllocation &alloc, + std::size_t size, + uint32_t data) { + TI_ERROR("fill_ndarray() not implemented on the current backend"); + } + + virtual void prepare_runtime_context(RuntimeContext *ctx) { + } + private: }; diff --git a/tests/cpp/codegen/refine_coordinates_test.cpp b/tests/cpp/codegen/refine_coordinates_test.cpp index f405d8727d41c..cee7a13091f47 100644 --- a/tests/cpp/codegen/refine_coordinates_test.cpp +++ b/tests/cpp/codegen/refine_coordinates_test.cpp @@ -12,6 +12,7 @@ #include "taichi/llvm/llvm_codegen_utils.h" #include "taichi/program/compile_config.h" #include "taichi/program/program.h" +#include "taichi/llvm/llvm_program.h" #include "taichi/struct/struct_llvm.h" namespace taichi { @@ -107,7 +108,9 @@ class RefineCoordinatesTest : public ::testing::Test { config_.packed = false; config_.print_kernel_llvm_ir = false; prog_ = std::make_unique(arch_); - tlctx_ = prog_->get_llvm_program_impl()->get_llvm_context(arch_); + auto *llvm_prog_ = + dynamic_cast(prog_->get_program_impl()); + tlctx_ = llvm_prog_->get_llvm_context(arch_); root_snode_ = std::make_unique(/*depth=*/0, /*t=*/SNodeType::root); const std::vector axes = {Axis{0}}; diff --git a/tests/cpp/llvm/llvm_offline_cache_test.cpp b/tests/cpp/llvm/llvm_offline_cache_test.cpp index ab8282bc71bc7..09b61d3e1bd28 100644 --- a/tests/cpp/llvm/llvm_offline_cache_test.cpp +++ b/tests/cpp/llvm/llvm_offline_cache_test.cpp @@ -46,7 +46,9 @@ class LlvmOfflineCacheTest : public testing::TestWithParam { config_.packed = false; config_.print_kernel_llvm_ir = false; prog_ = std::make_unique(arch); - tlctx_ = prog_->get_llvm_program_impl()->get_llvm_context(arch); + auto *llvm_prog_ = + dynamic_cast(prog_->get_program_impl()); + tlctx_ = llvm_prog_->get_llvm_context(arch); } static std::unique_ptr make_module( From d5df81a6112e1e25cca996884682aaf1abe80cc1 Mon Sep 17 00:00:00 2001 From: jim19930609 Date: Thu, 16 Jun 2022 15:59:25 +0800 Subject: [PATCH 3/7] [llvm] [refactor] LLVMProgramImpl code clean up: part-3 --- taichi/llvm/llvm_program.cpp | 2 ++ taichi/llvm/llvm_program.h | 8 ++++---- taichi/program/program.cpp | 34 ++++------------------------------ taichi/program/program_impl.h | 21 +++++++++++++++++++++ 4 files changed, 31 insertions(+), 34 deletions(-) diff --git a/taichi/llvm/llvm_program.cpp b/taichi/llvm/llvm_program.cpp index 9ac8844b8f44f..a1d36533752ac 100644 --- a/taichi/llvm/llvm_program.cpp +++ b/taichi/llvm/llvm_program.cpp @@ -474,6 +474,8 @@ void LlvmProgramImpl::materialize_runtime(MemoryPool *memory_pool, } void LlvmProgramImpl::check_runtime_error(uint64 *result_buffer) { + TI_ASSERT(arch_uses_llvm(config->arch)); + synchronize(); auto tlctx = llvm_context_host_.get(); if (llvm_context_device_) { diff --git a/taichi/llvm/llvm_program.h b/taichi/llvm/llvm_program.h index 3f9576b9507db..69d306983bcd9 100644 --- a/taichi/llvm/llvm_program.h +++ b/taichi/llvm/llvm_program.h @@ -99,13 +99,13 @@ class LlvmProgramImpl : public ProgramImpl { void print_memory_profiler_info( std::vector> &snode_trees_, - uint64 *result_buffer); + uint64 *result_buffer) override; void synchronize() override; - void check_runtime_error(uint64 *result_buffer); + void check_runtime_error(uint64 *result_buffer) override; - void finalize(); + void finalize() override; DeviceAllocation allocate_memory_ndarray(std::size_t alloc_size, uint64 *result_buffer) override; @@ -151,7 +151,7 @@ class LlvmProgramImpl : public ProgramImpl { std::unique_ptr compile_snode_tree_types_impl( SNodeTree *tree); - uint64 fetch_result_uint64(int i, uint64 *result_buffer); + uint64 fetch_result_uint64(int i, uint64 *result_buffer) override; template T runtime_query(const std::string &key, uint64 *result_buffer, Args... args) { diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp index d97052000d531..c2f6ecae0a373 100644 --- a/taichi/program/program.cpp +++ b/taichi/program/program.cpp @@ -216,13 +216,7 @@ SNode *Program::get_snode_root(int tree_id) { } void Program::check_runtime_error() { -#ifdef TI_WITH_LLVM - TI_ASSERT(arch_uses_llvm(config.arch)); - static_cast(program_impl_.get()) - ->check_runtime_error(result_buffer); -#else - TI_ERROR("Llvm disabled"); -#endif + program_impl_->check_runtime_error(result_buffer); } void Program::synchronize() { @@ -451,15 +445,7 @@ Kernel &Program::get_ndarray_writer(Ndarray *ndarray) { } uint64 Program::fetch_result_uint64(int i) { - if (arch_uses_llvm(config.arch)) { -#ifdef TI_WITH_LLVM - return static_cast(program_impl_.get()) - ->fetch_result(i, result_buffer); -#else - TI_NOT_IMPLEMENTED -#endif - } - return result_buffer[i]; + return program_impl_->fetch_result_uint64(i, result_buffer); } void Program::finalize() { @@ -506,13 +492,7 @@ void Program::finalize() { synchronize(); memory_pool_->terminate(); - if (arch_uses_llvm(config.arch)) { -#if TI_WITH_LLVM - static_cast(program_impl_.get())->finalize(); -#else - TI_NOT_IMPLEMENTED -#endif - } + program_impl_->finalize(); Stmt::reset_counter(); TaskLaunchRecord::reset_counter(); @@ -532,13 +512,7 @@ int Program::default_block_dim(const CompileConfig &config) { } void Program::print_memory_profiler_info() { -#ifdef TI_WITH_LLVM - TI_ASSERT(arch_uses_llvm(config.arch)); - static_cast(program_impl_.get()) - ->print_memory_profiler_info(snode_trees_, result_buffer); -#else - TI_ERROR("Llvm disabled"); -#endif + program_impl_->print_memory_profiler_info(snode_trees_, result_buffer); } std::size_t Program::get_snode_num_dynamically_allocated(SNode *snode) { diff --git a/taichi/program/program_impl.h b/taichi/program/program_impl.h index 714b128114a17..def3f65447330 100644 --- a/taichi/program/program_impl.h +++ b/taichi/program/program_impl.h @@ -101,21 +101,42 @@ class ProgramImpl { virtual ~ProgramImpl() { } + // TODO: Move to Runtime Object virtual uint64_t *get_ndarray_alloc_info_ptr(const DeviceAllocation &alloc) { TI_ERROR( "get_ndarray_alloc_info_ptr() not implemented on the current backend"); return nullptr; } + // TODO: Move to Runtime Object virtual void fill_ndarray(const DeviceAllocation &alloc, std::size_t size, uint32_t data) { TI_ERROR("fill_ndarray() not implemented on the current backend"); } + // TODO: Move to Runtime Object virtual void prepare_runtime_context(RuntimeContext *ctx) { } + virtual void print_memory_profiler_info( + std::vector> &snode_trees_, + uint64 *result_buffer) { + TI_ERROR( + "print_memory_profiler_info() not implemented on the current backend"); + } + + virtual void check_runtime_error(uint64 *result_buffer) { + TI_ERROR("check_runtime_error() not implemented on the current backend"); + } + + virtual void finalize() { + } + + virtual uint64 fetch_result_uint64(int i, uint64 *result_buffer) { + return result_buffer[i]; + } + private: }; From ef2760267383786a126b5d3e8ad860cc2841b13d Mon Sep 17 00:00:00 2001 From: jim19930609 Date: Thu, 16 Jun 2022 16:33:00 +0800 Subject: [PATCH 4/7] [llvm] [refactor] LLVMProgramImpl code clean up: part-4 --- taichi/llvm/llvm_program.h | 2 +- taichi/{system => llvm}/snode_tree_buffer_manager.cpp | 6 ------ taichi/{system => llvm}/snode_tree_buffer_manager.h | 0 3 files changed, 1 insertion(+), 7 deletions(-) rename taichi/{system => llvm}/snode_tree_buffer_manager.cpp (96%) rename taichi/{system => llvm}/snode_tree_buffer_manager.h (100%) diff --git a/taichi/llvm/llvm_program.h b/taichi/llvm/llvm_program.h index 69d306983bcd9..3ecdddee78f17 100644 --- a/taichi/llvm/llvm_program.h +++ b/taichi/llvm/llvm_program.h @@ -5,7 +5,7 @@ #include "taichi/llvm/llvm_device.h" #include "taichi/llvm/llvm_offline_cache.h" -#include "taichi/system/snode_tree_buffer_manager.h" +#include "taichi/llvm/snode_tree_buffer_manager.h" #include "taichi/inc/constants.h" #include "taichi/program/compile_config.h" #include "taichi/common/logging.h" diff --git a/taichi/system/snode_tree_buffer_manager.cpp b/taichi/llvm/snode_tree_buffer_manager.cpp similarity index 96% rename from taichi/system/snode_tree_buffer_manager.cpp rename to taichi/llvm/snode_tree_buffer_manager.cpp index 7a0d2accd90f6..6f3a79ea3e256 100644 --- a/taichi/system/snode_tree_buffer_manager.cpp +++ b/taichi/llvm/snode_tree_buffer_manager.cpp @@ -1,8 +1,6 @@ #include "snode_tree_buffer_manager.h" #include "taichi/program/program.h" -#ifdef TI_WITH_LLVM #include "taichi/llvm/llvm_program.h" -#endif TLANG_NAMESPACE_BEGIN @@ -40,7 +38,6 @@ Ptr SNodeTreeBufferManager::allocate(JITModule *runtime_jit, std::size_t alignment, const int snode_tree_id, uint64 *result_buffer) { -#ifdef TI_WITH_LLVM TI_TRACE("allocating memory for SNode Tree {}", snode_tree_id); TI_ASSERT_INFO(snode_tree_id < kMaxNumSnodeTreesLlvm, "LLVM backend supports up to {} snode trees", @@ -68,9 +65,6 @@ Ptr SNodeTreeBufferManager::allocate(JITModule *runtime_jit, sizes_[snode_tree_id] = size; return x.second; } -#else - TI_ERROR("Llvm disabled"); -#endif } void SNodeTreeBufferManager::destroy(SNodeTree *snode_tree) { diff --git a/taichi/system/snode_tree_buffer_manager.h b/taichi/llvm/snode_tree_buffer_manager.h similarity index 100% rename from taichi/system/snode_tree_buffer_manager.h rename to taichi/llvm/snode_tree_buffer_manager.h From 645c9fa8bd7306460ee5d9000adcb5cd6188c178 Mon Sep 17 00:00:00 2001 From: jim19930609 Date: Thu, 16 Jun 2022 16:38:18 +0800 Subject: [PATCH 5/7] Minor fix --- taichi/program/program.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp index c2f6ecae0a373..9b42633e5b036 100644 --- a/taichi/program/program.cpp +++ b/taichi/program/program.cpp @@ -491,8 +491,9 @@ void Program::finalize() { synchronize(); memory_pool_->terminate(); - - program_impl_->finalize(); + if (arch_uses_llvm(config.arch)) { + program_impl_->finalize(); + } Stmt::reset_counter(); TaskLaunchRecord::reset_counter(); From 8f2288ce53355ddaed2cb9c4408c867dd445de99 Mon Sep 17 00:00:00 2001 From: jim19930609 Date: Fri, 17 Jun 2022 15:20:45 +0800 Subject: [PATCH 6/7] [llvm] [refactor] LLVMProgramImpl code clean up: part-5 --- taichi/codegen/codegen_llvm.cpp | 2 +- taichi/codegen/codegen_llvm_quant.cpp | 2 +- taichi/llvm/llvm_program.h | 2 +- taichi/{struct => llvm}/struct_llvm.cpp | 5 +---- taichi/{struct => llvm}/struct_llvm.h | 3 --- taichi/program/program.cpp | 2 +- tests/cpp/codegen/refine_coordinates_test.cpp | 2 +- 7 files changed, 6 insertions(+), 12 deletions(-) rename taichi/{struct => llvm}/struct_llvm.cpp (99%) rename taichi/{struct => llvm}/struct_llvm.h (97%) diff --git a/taichi/codegen/codegen_llvm.cpp b/taichi/codegen/codegen_llvm.cpp index 659d131249fd8..049af887975b0 100644 --- a/taichi/codegen/codegen_llvm.cpp +++ b/taichi/codegen/codegen_llvm.cpp @@ -12,7 +12,7 @@ #include "taichi/llvm/launch_arg_info.h" #include "taichi/llvm/llvm_offline_cache.h" #include "taichi/llvm/llvm_program.h" -#include "taichi/struct/struct_llvm.h" +#include "taichi/llvm/struct_llvm.h" #include "taichi/util/file_sequence_writer.h" TLANG_NAMESPACE_BEGIN diff --git a/taichi/codegen/codegen_llvm_quant.cpp b/taichi/codegen/codegen_llvm_quant.cpp index c92a0720c5ec8..083f6242f6dad 100644 --- a/taichi/codegen/codegen_llvm_quant.cpp +++ b/taichi/codegen/codegen_llvm_quant.cpp @@ -2,7 +2,7 @@ #include "taichi/codegen/codegen_llvm.h" #include "taichi/ir/statements.h" -#include "taichi/struct/struct_llvm.h" +#include "taichi/llvm/struct_llvm.h" TLANG_NAMESPACE_BEGIN diff --git a/taichi/llvm/llvm_program.h b/taichi/llvm/llvm_program.h index 3ecdddee78f17..f50692094d6cd 100644 --- a/taichi/llvm/llvm_program.h +++ b/taichi/llvm/llvm_program.h @@ -14,7 +14,7 @@ #include "taichi/runtime/runtime.h" #include "taichi/system/threading.h" #include "taichi/struct/struct.h" -#include "taichi/struct/struct_llvm.h" +#include "taichi/llvm/struct_llvm.h" #include "taichi/program/snode_expr_utils.h" #include "taichi/system/memory_pool.h" #include "taichi/program/program_impl.h" diff --git a/taichi/struct/struct_llvm.cpp b/taichi/llvm/struct_llvm.cpp similarity index 99% rename from taichi/struct/struct_llvm.cpp rename to taichi/llvm/struct_llvm.cpp index adf6a9be08a9b..61917013fc802 100644 --- a/taichi/struct/struct_llvm.cpp +++ b/taichi/llvm/struct_llvm.cpp @@ -1,5 +1,4 @@ -#ifdef TI_WITH_LLVM -#include "taichi/struct/struct_llvm.h" +#include "taichi/llvm/struct_llvm.h" #include "llvm/IR/Verifier.h" #include "llvm/IR/IRBuilder.h" @@ -363,5 +362,3 @@ llvm::Function *StructCompilerLLVM::create_function(llvm::FunctionType *ft, } // namespace lang } // namespace taichi - -#endif //#ifdef TI_WITH_LLVM diff --git a/taichi/struct/struct_llvm.h b/taichi/llvm/struct_llvm.h similarity index 97% rename from taichi/struct/struct_llvm.h rename to taichi/llvm/struct_llvm.h index 5a725381ae7e5..3b887d7ee9c29 100644 --- a/taichi/struct/struct_llvm.h +++ b/taichi/llvm/struct_llvm.h @@ -1,6 +1,5 @@ #pragma once -#ifdef TI_WITH_LLVM // Codegen for the hierarchical data structure (LLVM) #include "taichi/llvm/llvm_program.h" #include "taichi/llvm/llvm_codegen_utils.h" @@ -56,5 +55,3 @@ class StructCompilerLLVM : public StructCompiler, public LLVMModuleBuilder { } // namespace lang } // namespace taichi - -#endif //#ifdef TI_WITH_LLVM diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp index c2f6ecae0a373..83bc53fd8298d 100644 --- a/taichi/program/program.cpp +++ b/taichi/program/program.cpp @@ -6,7 +6,7 @@ #include "taichi/program/extension.h" #include "taichi/backends/cpu/codegen_cpu.h" #include "taichi/struct/struct.h" -#include "taichi/struct/struct_llvm.h" +#include "taichi/llvm/struct_llvm.h" #include "taichi/backends/metal/api.h" #include "taichi/backends/wasm/aot_module_builder_impl.h" #include "taichi/backends/opengl/opengl_program.h" diff --git a/tests/cpp/codegen/refine_coordinates_test.cpp b/tests/cpp/codegen/refine_coordinates_test.cpp index cee7a13091f47..9f3b0b6faf4f1 100644 --- a/tests/cpp/codegen/refine_coordinates_test.cpp +++ b/tests/cpp/codegen/refine_coordinates_test.cpp @@ -13,7 +13,7 @@ #include "taichi/program/compile_config.h" #include "taichi/program/program.h" #include "taichi/llvm/llvm_program.h" -#include "taichi/struct/struct_llvm.h" +#include "taichi/llvm/struct_llvm.h" namespace taichi { From 89af1f366a385ae153c0384799fd69b35c065f6b Mon Sep 17 00:00:00 2001 From: jim19930609 Date: Fri, 17 Jun 2022 16:02:31 +0800 Subject: [PATCH 7/7] Minor bug fix --- taichi/program/program.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp index 83bc53fd8298d..493c732869e97 100644 --- a/taichi/program/program.cpp +++ b/taichi/program/program.cpp @@ -6,7 +6,6 @@ #include "taichi/program/extension.h" #include "taichi/backends/cpu/codegen_cpu.h" #include "taichi/struct/struct.h" -#include "taichi/llvm/struct_llvm.h" #include "taichi/backends/metal/api.h" #include "taichi/backends/wasm/aot_module_builder_impl.h" #include "taichi/backends/opengl/opengl_program.h"