From 365827df08fd2aa35ac56ae7afae8c436aaa74f4 Mon Sep 17 00:00:00 2001 From: PGZXB <420254146@qq.com> Date: Wed, 9 Mar 2022 19:40:16 +0800 Subject: [PATCH 01/10] Impl offline-cache --- .gitmodules | 3 + cmake/TaichiCore.cmake | 1 + external/PicoSHA2 | 1 + taichi/backends/cpu/codegen_cpu.cpp | 5 +- taichi/backends/cpu/codegen_cpu.h | 6 +- taichi/codegen/codegen.cpp | 5 +- taichi/codegen/codegen.h | 3 +- taichi/codegen/codegen_llvm.cpp | 14 ++- taichi/codegen/codegen_llvm.h | 6 +- taichi/llvm/llvm_offline_cache.h | 137 ++++++++++++++++++++++++++++ taichi/llvm/llvm_program.cpp | 89 +++++++++++++++++- taichi/llvm/llvm_program.h | 17 ++++ taichi/program/compile_config.h | 4 + taichi/program/kernel.cpp | 7 +- taichi/program/kernel.h | 14 +++ taichi/python/export_lang.cpp | 5 +- 16 files changed, 303 insertions(+), 14 deletions(-) create mode 160000 external/PicoSHA2 create mode 100644 taichi/llvm/llvm_offline_cache.h diff --git a/.gitmodules b/.gitmodules index be8c04fa168b0..c81394ee31a41 100644 --- a/.gitmodules +++ b/.gitmodules @@ -56,3 +56,6 @@ [submodule "external/FP16"] path = external/FP16 url = https://github.com/Maratyszcza/FP16 +[submodule "external/PicoSHA2"] + path = external/PicoSHA2 + url = https://github.com/okdshin/PicoSHA2.git diff --git a/cmake/TaichiCore.cmake b/cmake/TaichiCore.cmake index 3545dc68b0fb2..18a3ea60b4bc8 100644 --- a/cmake/TaichiCore.cmake +++ b/cmake/TaichiCore.cmake @@ -235,6 +235,7 @@ endif() include_directories(${CMAKE_SOURCE_DIR}) include_directories(external/include) include_directories(external/spdlog/include) +include_directories(external/PicoSHA2) if (TI_WITH_OPENGL) target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/glad/include) endif() diff --git a/external/PicoSHA2 b/external/PicoSHA2 new file mode 160000 index 0000000000000..1677374f23352 --- /dev/null +++ b/external/PicoSHA2 @@ -0,0 +1 @@ +Subproject commit 1677374f23352716fc52183255a40c1b8e1d53eb diff --git a/taichi/backends/cpu/codegen_cpu.cpp b/taichi/backends/cpu/codegen_cpu.cpp index 0b1809e4d75fe..bedeaea907885 100644 --- a/taichi/backends/cpu/codegen_cpu.cpp +++ b/taichi/backends/cpu/codegen_cpu.cpp @@ -16,7 +16,8 @@ class CodeGenLLVMCPU : public CodeGenLLVM { public: using IRVisitor::visit; - CodeGenLLVMCPU(Kernel *kernel, IRNode *ir) : CodeGenLLVM(kernel, ir) { + CodeGenLLVMCPU(Kernel *kernel, IRNode *ir, bool needs_cache) + : CodeGenLLVM(kernel, ir, nullptr, needs_cache) { TI_AUTO_PROF } @@ -194,7 +195,7 @@ class CodeGenLLVMCPU : public CodeGenLLVM { FunctionType CodeGenCPU::codegen() { TI_AUTO_PROF - return CodeGenLLVMCPU(kernel, ir).gen(); + return CodeGenLLVMCPU(kernel, ir, needs_cache_).gen(); } TLANG_NAMESPACE_END diff --git a/taichi/backends/cpu/codegen_cpu.h b/taichi/backends/cpu/codegen_cpu.h index c3d723c75eff2..a00458a73165a 100644 --- a/taichi/backends/cpu/codegen_cpu.h +++ b/taichi/backends/cpu/codegen_cpu.h @@ -8,10 +8,14 @@ TLANG_NAMESPACE_BEGIN class CodeGenCPU : public KernelCodeGen { public: - CodeGenCPU(Kernel *kernel, IRNode *ir = nullptr) : KernelCodeGen(kernel, ir) { + CodeGenCPU(Kernel *kernel, IRNode *ir = nullptr, bool needs_cache = false) + : KernelCodeGen(kernel, ir), needs_cache_(needs_cache) { } FunctionType codegen() override; + + private: + bool needs_cache_{false}; }; TLANG_NAMESPACE_END diff --git a/taichi/codegen/codegen.cpp b/taichi/codegen/codegen.cpp index 263f15d039d4d..b7bb02dd835d8 100644 --- a/taichi/codegen/codegen.cpp +++ b/taichi/codegen/codegen.cpp @@ -32,10 +32,11 @@ KernelCodeGen::KernelCodeGen(Kernel *kernel, IRNode *ir) std::unique_ptr KernelCodeGen::create(Arch arch, Kernel *kernel, - Stmt *stmt) { + Stmt *stmt, + bool needs_cache) { #ifdef TI_WITH_LLVM if (arch_is_cpu(arch) && arch != Arch::wasm) { - return std::make_unique(kernel, stmt); + return std::make_unique(kernel, stmt, needs_cache); } else if (arch == Arch::wasm) { return std::make_unique(kernel, stmt); } else if (arch == Arch::cuda) { diff --git a/taichi/codegen/codegen.h b/taichi/codegen/codegen.h index e700fe0f349e2..42ab840a0b081 100644 --- a/taichi/codegen/codegen.h +++ b/taichi/codegen/codegen.h @@ -19,7 +19,8 @@ class KernelCodeGen { static std::unique_ptr create(Arch arch, Kernel *kernel, - Stmt *stmt = nullptr); + Stmt *stmt = nullptr, + bool needs_cache = false); virtual FunctionType codegen() = 0; }; diff --git a/taichi/codegen/codegen_llvm.cpp b/taichi/codegen/codegen_llvm.cpp index 241cb15161bb1..e6fd8d8e677fc 100644 --- a/taichi/codegen/codegen_llvm.cpp +++ b/taichi/codegen/codegen_llvm.cpp @@ -297,7 +297,8 @@ void CodeGenLLVM::emit_struct_meta_base(const std::string &name, CodeGenLLVM::CodeGenLLVM(Kernel *kernel, IRNode *ir, - std::unique_ptr &&module) + std::unique_ptr &&module, + bool needs_cache) // TODO: simplify LLVMModuleBuilder ctor input : LLVMModuleBuilder( module == nullptr ? kernel->program->get_llvm_program_impl() @@ -306,6 +307,7 @@ CodeGenLLVM::CodeGenLLVM(Kernel *kernel, : std::move(module), kernel->program->get_llvm_program_impl()->get_llvm_context( kernel->arch)), + needs_cache_(needs_cache), kernel(kernel), ir(ir), prog(kernel->program) { @@ -2240,6 +2242,16 @@ FunctionType CodeGenLLVM::compile_module_to_executable() { TI_AUTO_PROF eliminate_unused_functions(); + auto *llvm_prog = prog->get_llvm_program_impl(); + if (needs_cache_) { + std::vector offloaded_task_name_list; + for (auto &task : offloaded_tasks) { + offloaded_task_name_list.push_back(task.name); + } + llvm_prog->cache_kernel(this->kernel->get_key(), this->module.get(), + std::move(offloaded_task_name_list)); + } + tlctx->add_module(std::move(module)); for (auto &task : offloaded_tasks) { diff --git a/taichi/codegen/codegen_llvm.h b/taichi/codegen/codegen_llvm.h index e32ee3cdce298..8afda8b7ea0ef 100644 --- a/taichi/codegen/codegen_llvm.h +++ b/taichi/codegen/codegen_llvm.h @@ -48,6 +48,9 @@ class FunctionCreationGuard { }; class CodeGenLLVM : public IRVisitor, public LLVMModuleBuilder { + private: + bool needs_cache_{false}; + public: Kernel *kernel; IRNode *ir; @@ -82,7 +85,8 @@ class CodeGenLLVM : public IRVisitor, public LLVMModuleBuilder { CodeGenLLVM(Kernel *kernel, IRNode *ir = nullptr, - std::unique_ptr &&module = nullptr); + std::unique_ptr &&module = nullptr, + bool needs_cache = false); Arch current_arch() { return kernel->arch; diff --git a/taichi/llvm/llvm_offline_cache.h b/taichi/llvm/llvm_offline_cache.h new file mode 100644 index 0000000000000..1d2776f06915d --- /dev/null +++ b/taichi/llvm/llvm_offline_cache.h @@ -0,0 +1,137 @@ +#pragma once + +#include "llvm/AsmParser/Parser.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/raw_os_ostream.h" + +#include "taichi/common/core.h" +#include "taichi/common/logging.h" +#include "taichi/program/compile_config.h" +#include "taichi/ir/ir.h" + +TLANG_NAMESPACE_BEGIN +struct LlvmOfflineCache { + struct KernelCacheData { + std::string kernel_key; + std::unique_ptr owned_module{nullptr}; + llvm::Module *module{nullptr}; + std::vector offloaded_task_name_list; + + KernelCacheData() = default; + KernelCacheData(KernelCacheData &&) = default; + KernelCacheData &operator=(KernelCacheData &&) = default; + ~KernelCacheData() = default; + }; + + std::unordered_map kernels; +}; + +class LlvmOfflineCacheFileReader { + public: + LlvmOfflineCacheFileReader(const std::string &path) : path_(path) { + } + + bool get_kernel_cache(LlvmOfflineCache::KernelCacheData &res, + const std::string &key, + llvm::LLVMContext &llvm_ctx) { + res.kernel_key = key; + std::string filename_prefix = path_ + "/" + key; + { + std::string filename = filename_prefix + ".ll"; + llvm::SMDiagnostic err; + res.owned_module = llvm::parseAssemblyFile(filename, err, llvm_ctx); + res.module = res.owned_module.get(); + if (!res.module) + return false; + } + { + std::string filename = filename_prefix + "_otnl.txt"; + std::ifstream in(filename, std::ios::in | std::ios::binary); + if (!in.is_open()) + return false; + while (true) { + std::string line; + std::getline(in, line, '\n'); + if (line.empty()) + break; + res.offloaded_task_name_list.push_back(std::move(line)); + } + } + return true; + } + + private: + std::string path_; +}; + +class LlvmOfflineCacheFileWriter { + public: + LlvmOfflineCacheFileWriter(const std::string &path) : path_(path) { + } + + void set_data(LlvmOfflineCache &&data) { + this->mangled_ = false; + this->data_ = std::move(data); + } + + void add_kernel_cache(const std::string &key, + LlvmOfflineCache::KernelCacheData &&kernel_cache) { + data_.kernels[key] = std::move(kernel_cache); + } + + void dump() { + for (auto &[k, v] : data_.kernels) { + std::string filename_prefix = path_ + "/" + k; + { + std::string filename = filename_prefix + ".ll"; + std::ofstream os(filename, std::ios::out | std::ios::binary); + TI_ERROR_IF(!os.is_open(), "File {} open failed", filename); + llvm::SMDiagnostic err; + llvm::LLVMContext ctx; + llvm::raw_os_ostream llvm_os(os); + if (v.module) { + mangle_offloaded_task_name(k, v.module, v.offloaded_task_name_list); + v.module->print(llvm_os, nullptr); + } else if (v.owned_module) { + mangle_offloaded_task_name(k, v.owned_module.get(), + v.offloaded_task_name_list); + v.owned_module->print(llvm_os, nullptr); + } else + TI_ASSERT(false); + } + { + std::string filename = filename_prefix + "_otnl.txt"; + std::ofstream os(filename, std::ios::out | std::ios::binary); + TI_ERROR_IF(!os.is_open(), "File {} open failed", filename); + for (const auto &name : v.offloaded_task_name_list) { + os << name << '\n'; + } + } + } + } + + private: + void mangle_offloaded_task_name( + const std::string &kernel_key, + llvm::Module *module, + std::vector &offloaded_task_name_list) { + if (!mangled_) { + std::size_t cnt = 0; + for (auto &e : offloaded_task_name_list) { + std::string mangled_name = kernel_key + std::to_string(cnt++); + auto func = module->getFunction(e); + TI_ASSERT(func != nullptr); + func->setName(mangled_name); + e = mangled_name; + } + } + } + + std::string path_; + LlvmOfflineCache data_; + bool mangled_{false}; +}; + +TLANG_NAMESPACE_END diff --git a/taichi/llvm/llvm_program.cpp b/taichi/llvm/llvm_program.cpp index c07596707360f..07d718e4411e5 100644 --- a/taichi/llvm/llvm_program.cpp +++ b/taichi/llvm/llvm_program.cpp @@ -1,19 +1,21 @@ #include "llvm_program.h" #include "llvm/IR/Module.h" +#include "picosha2.h" + #include "taichi/backends/cuda/cuda_driver.h" #include "taichi/backends/arch.h" +#include "taichi/llvm/llvm_offline_cache.h" #include "taichi/platform/cuda/detect_cuda.h" #include "taichi/math/arithmetic.h" #include "taichi/runtime/llvm/mem_request.h" #include "taichi/util/str.h" #include "taichi/codegen/codegen.h" #include "taichi/ir/statements.h" +#include "taichi/ir/transforms.h" #include "taichi/backends/cpu/cpu_device.h" #include "taichi/backends/cuda/cuda_device.h" -#include "taichi/backends/cuda/cuda_device.h" - #if defined(TI_WITH_CUDA) #include "taichi/backends/cuda/cuda_driver.h" #include "taichi/backends/cuda/codegen_cuda.h" @@ -129,12 +131,38 @@ void LlvmProgramImpl::maybe_initialize_cuda_llvm_context() { } } +LlvmProgramImpl::~LlvmProgramImpl() { + if (config->offline_cache && this->supports_offline_cache()) { + LlvmOfflineCacheFileWriter writer(config->offline_cache_file_path); + writer.set_data(std::move(cache_data_)); + writer.dump(); + } +} + FunctionType LlvmProgramImpl::compile(Kernel *kernel, OffloadedStmt *offloaded) { + bool needs_cache = false; + if (config->offline_cache && this->supports_offline_cache() && + !kernel->is_evaluator) { + std::string kernel_key, hashed_kernel_key; + irpass::re_id(kernel->ir.get()); + irpass::print(kernel->ir.get(), &kernel_key); + picosha2::hash256_hex_string(kernel_key, hashed_kernel_key); + auto func = this->create_kernel_function_from_offline_cache( + hashed_kernel_key, kernel); + if (func) { + kernel->set_from_offline_cache(); + return func; + } else { + kernel->set_key(hashed_kernel_key); + needs_cache = true; + } + } if (!kernel->lowered()) { kernel->lower(); } - auto codegen = KernelCodeGen::create(kernel->arch, kernel, offloaded); + auto codegen = + KernelCodeGen::create(kernel->arch, kernel, offloaded, needs_cache); return codegen->codegen(); } @@ -642,5 +670,60 @@ void LlvmProgramImpl::fill_ndarray(const DeviceAllocation &alloc, std::fill((uint32_t *)ptr, (uint32_t *)ptr + size, data); } } + +FunctionType LlvmProgramImpl::create_kernel_function_from_offline_cache( + const std::string &kernel_key, + Kernel *kernel) { + TI_ASSERT(config->offline_cache); + using task_fp_type = int32 (*)(void *); + + LlvmOfflineCacheFileReader reader(config->offline_cache_file_path); + LlvmOfflineCache::KernelCacheData cache_data; + auto *tlctx = this->get_llvm_context(config->arch); + auto &llvm_ctx = *tlctx->get_this_thread_context(); + + if (!reader.get_kernel_cache(cache_data, kernel_key, llvm_ctx)) + return nullptr; + + std::vector func_list; + tlctx->add_module(std::move(cache_data.owned_module)); + for (const auto &func_name : cache_data.offloaded_task_name_list) { + void *kernel_symbol = tlctx->lookup_function_pointer(func_name); + TI_ASSERT(kernel_symbol); + func_list.push_back((task_fp_type)kernel_symbol); + } + + return [kernel, flist = std::move(func_list)](RuntimeContext &ctx) -> void { + auto args = kernel->args; + // For taichi ndarrays, context.args saves pointer to its + // |DeviceAllocation|, CPU backend actually want to use the raw ptr + // here. + for (std::size_t i = 0; i < args.size(); ++i) { + if (args[i].is_array && ctx.is_device_allocation[i] && args[i].size > 0) { + DeviceAllocation *ptr = + static_cast(ctx.get_arg(i)); + uint64 host_ptr = (uint64)kernel->program->get_llvm_program_impl() + ->get_ndarray_alloc_info_ptr(*ptr); + ctx.set_arg(i, host_ptr); + ctx.set_device_allocation(i, false); + } + } + for (auto func : flist) { + func(&ctx); + } + }; +} + +void LlvmProgramImpl::cache_kernel( + const std::string &kernel_key, + llvm::Module *module, + std::vector &&offloaded_task_name_list) { + TI_ASSERT(cache_data_.kernels.find(kernel_key) == cache_data_.kernels.end()); + auto &kernel_cache = cache_data_.kernels[kernel_key]; + kernel_cache.kernel_key = kernel_key; + kernel_cache.owned_module = llvm::CloneModule(*module); + kernel_cache.offloaded_task_name_list = offloaded_task_name_list; +} + } // namespace lang } // namespace taichi diff --git a/taichi/llvm/llvm_program.h b/taichi/llvm/llvm_program.h index 72ef3c111b42f..248e5acaeb6f3 100644 --- a/taichi/llvm/llvm_program.h +++ b/taichi/llvm/llvm_program.h @@ -1,5 +1,6 @@ #pragma once #include "taichi/llvm/llvm_device.h" +#include "taichi/llvm/llvm_offline_cache.h" #include "taichi/system/snode_tree_buffer_manager.h" #include "taichi/inc/constants.h" #include "taichi/program/compile_config.h" @@ -37,6 +38,7 @@ class CpuDevice; class LlvmProgramImpl : public ProgramImpl { public: LlvmProgramImpl(CompileConfig &config, KernelProfilerBase *profiler); + ~LlvmProgramImpl() override; void initialize_host(); @@ -114,6 +116,19 @@ class LlvmProgramImpl : public ProgramImpl { std::size_t size, uint32_t data); + bool supports_offline_cache() const { + auto arch = config->arch; + return arch_is_cpu(arch) && arch != Arch::wasm && !config->async_mode; + } + + FunctionType create_kernel_function_from_offline_cache( + const std::string &kernel_key, + Kernel *kernel); + + void cache_kernel(const std::string &kernel_key, + llvm::Module *module, + std::vector &&offloaded_task_name_list); + private: std::unique_ptr clone_struct_compiler_initial_context( const std::vector> &snode_trees_, @@ -172,6 +187,8 @@ class LlvmProgramImpl : public ProgramImpl { std::unordered_map snode_tree_allocs_; + LlvmOfflineCache cache_data_; + std::shared_ptr device_{nullptr}; cuda::CudaDevice *cuda_device(); cpu::CpuDevice *cpu_device(); diff --git a/taichi/program/compile_config.h b/taichi/program/compile_config.h index 9b329f6f1e4ae..0c80fef5d526b 100644 --- a/taichi/program/compile_config.h +++ b/taichi/program/compile_config.h @@ -105,6 +105,10 @@ struct CompileConfig { bool experimental_auto_mesh_local{false}; int auto_mesh_local_default_occupacy{4}; + // Offline cache options + bool offline_cache{false}; + std::string offline_cache_file_path{"."}; + CompileConfig(); }; diff --git a/taichi/program/kernel.cpp b/taichi/program/kernel.cpp index 582ee70c9eb70..863070aeec6af 100644 --- a/taichi/program/kernel.cpp +++ b/taichi/program/kernel.cpp @@ -2,6 +2,7 @@ #include "taichi/backends/cuda/cuda_driver.h" #include "taichi/codegen/codegen.h" +#include "taichi/common/logging.h" #include "taichi/common/task.h" #include "taichi/ir/statements.h" #include "taichi/ir/transforms.h" @@ -106,8 +107,10 @@ void Kernel::operator()(LaunchContextBuilder &ctx_builder) { compile(); } - for (auto &offloaded : ir->as()->statements) { - account_for_offloaded(offloaded->as()); + if (!this->from_offline_cache_) { + for (auto &offloaded : ir->as()->statements) { + account_for_offloaded(offloaded->as()); + } } compiled_(ctx_builder.get_context()); diff --git a/taichi/program/kernel.h b/taichi/program/kernel.h index c3a270ab098ed..19399e37078cc 100644 --- a/taichi/program/kernel.h +++ b/taichi/program/kernel.h @@ -117,6 +117,18 @@ class TI_DLL_EXPORT Kernel : public Callable { return task_counter_++; } + void set_key(const std::string &key) { + this->kernel_key_ = key; + } + + const std::string &get_key() const { + return this->kernel_key_; + } + + void set_from_offline_cache() { + this->from_offline_cache_ = true; + } + [[nodiscard]] std::string get_name() const override; /** * Whether the given |arch| is supported in the lower() method. @@ -141,6 +153,8 @@ class TI_DLL_EXPORT Kernel : public Callable { // OffloadedStmt for async execution bool lowered_{false}; std::atomic task_counter_{0}; + std::string kernel_key_; + bool from_offline_cache_{false}; }; TLANG_NAMESPACE_END diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp index a8cbfa1bafe33..e28bc7b9f4696 100644 --- a/taichi/python/export_lang.cpp +++ b/taichi/python/export_lang.cpp @@ -217,7 +217,10 @@ void export_lang(py::module &m) { .def_readwrite("experimental_auto_mesh_local", &CompileConfig::experimental_auto_mesh_local) .def_readwrite("auto_mesh_local_default_occupacy", - &CompileConfig::auto_mesh_local_default_occupacy); + &CompileConfig::auto_mesh_local_default_occupacy) + .def_readwrite("offline_cache", &CompileConfig::offline_cache) + .def_readwrite("offline_cache_file_path", + &CompileConfig::offline_cache_file_path); m.def("reset_default_compile_config", [&]() { default_compile_config = CompileConfig(); }); From 7951f69c7c26acbf604fe3fe195391f8bbc23ff6 Mon Sep 17 00:00:00 2001 From: PGZXB <420254146@qq.com> Date: Thu, 10 Mar 2022 09:51:05 +0800 Subject: [PATCH 02/10] Move implementation of some functions of LlvmOfflineCacheFileReader/Writer to llvm_offline_cache.cpp --- taichi/llvm/llvm_offline_cache.cpp | 87 ++++++++++++++++++++++++++++++ taichi/llvm/llvm_offline_cache.h | 81 ++-------------------------- 2 files changed, 91 insertions(+), 77 deletions(-) create mode 100644 taichi/llvm/llvm_offline_cache.cpp diff --git a/taichi/llvm/llvm_offline_cache.cpp b/taichi/llvm/llvm_offline_cache.cpp new file mode 100644 index 0000000000000..99a93d67c99e9 --- /dev/null +++ b/taichi/llvm/llvm_offline_cache.cpp @@ -0,0 +1,87 @@ +#include "llvm_offline_cache.h" + +#include "llvm/AsmParser/Parser.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/raw_os_ostream.h" +#include "llvm/IR/Module.h" + +TLANG_NAMESPACE_BEGIN + +bool LlvmOfflineCacheFileReader::get_kernel_cache( + LlvmOfflineCache::KernelCacheData &res, + const std::string &key, + llvm::LLVMContext &llvm_ctx) { + res.kernel_key = key; + std::string filename_prefix = path_ + "/" + key; + { + std::string filename = filename_prefix + ".ll"; + llvm::SMDiagnostic err; + res.owned_module = llvm::parseAssemblyFile(filename, err, llvm_ctx); + res.module = res.owned_module.get(); + if (!res.module) + return false; + } + { + std::string filename = filename_prefix + "_otnl.txt"; + std::ifstream in(filename, std::ios::in | std::ios::binary); + if (!in.is_open()) + return false; + while (true) { + std::string line; + std::getline(in, line, '\n'); + if (line.empty()) + break; + res.offloaded_task_name_list.push_back(std::move(line)); + } + } + return true; +} + +void LlvmOfflineCacheFileWriter::dump() { + for (auto &[k, v] : data_.kernels) { + std::string filename_prefix = path_ + "/" + k; + { + std::string filename = filename_prefix + ".ll"; + std::ofstream os(filename, std::ios::out | std::ios::binary); + TI_ERROR_IF(!os.is_open(), "File {} open failed", filename); + llvm::SMDiagnostic err; + llvm::LLVMContext ctx; + llvm::raw_os_ostream llvm_os(os); + if (v.module) { + mangle_offloaded_task_name(k, v.module, v.offloaded_task_name_list); + v.module->print(llvm_os, nullptr); + } else if (v.owned_module) { + mangle_offloaded_task_name(k, v.owned_module.get(), + v.offloaded_task_name_list); + v.owned_module->print(llvm_os, nullptr); + } else + TI_ASSERT(false); + } + { + std::string filename = filename_prefix + "_otnl.txt"; + std::ofstream os(filename, std::ios::out | std::ios::binary); + TI_ERROR_IF(!os.is_open(), "File {} open failed", filename); + for (const auto &name : v.offloaded_task_name_list) { + os << name << '\n'; + } + } + } +} + +void LlvmOfflineCacheFileWriter::mangle_offloaded_task_name( + const std::string &kernel_key, + llvm::Module *module, + std::vector &offloaded_task_name_list) { + if (!mangled_) { + std::size_t cnt = 0; + for (auto &e : offloaded_task_name_list) { + std::string mangled_name = kernel_key + std::to_string(cnt++); + auto func = module->getFunction(e); + TI_ASSERT(func != nullptr); + func->setName(mangled_name); + e = mangled_name; + } + } +} + +TLANG_NAMESPACE_END diff --git a/taichi/llvm/llvm_offline_cache.h b/taichi/llvm/llvm_offline_cache.h index 1d2776f06915d..d3752690a9e36 100644 --- a/taichi/llvm/llvm_offline_cache.h +++ b/taichi/llvm/llvm_offline_cache.h @@ -1,15 +1,7 @@ #pragma once -#include "llvm/AsmParser/Parser.h" -#include "llvm/IR/LLVMContext.h" -#include "llvm/IR/Module.h" -#include "llvm/Support/SourceMgr.h" -#include "llvm/Support/raw_os_ostream.h" - #include "taichi/common/core.h" -#include "taichi/common/logging.h" -#include "taichi/program/compile_config.h" -#include "taichi/ir/ir.h" +#include "taichi/llvm/llvm_fwd.h" TLANG_NAMESPACE_BEGIN struct LlvmOfflineCache { @@ -35,32 +27,7 @@ class LlvmOfflineCacheFileReader { bool get_kernel_cache(LlvmOfflineCache::KernelCacheData &res, const std::string &key, - llvm::LLVMContext &llvm_ctx) { - res.kernel_key = key; - std::string filename_prefix = path_ + "/" + key; - { - std::string filename = filename_prefix + ".ll"; - llvm::SMDiagnostic err; - res.owned_module = llvm::parseAssemblyFile(filename, err, llvm_ctx); - res.module = res.owned_module.get(); - if (!res.module) - return false; - } - { - std::string filename = filename_prefix + "_otnl.txt"; - std::ifstream in(filename, std::ios::in | std::ios::binary); - if (!in.is_open()) - return false; - while (true) { - std::string line; - std::getline(in, line, '\n'); - if (line.empty()) - break; - res.offloaded_task_name_list.push_back(std::move(line)); - } - } - return true; - } + llvm::LLVMContext &llvm_ctx); private: std::string path_; @@ -81,53 +48,13 @@ class LlvmOfflineCacheFileWriter { data_.kernels[key] = std::move(kernel_cache); } - void dump() { - for (auto &[k, v] : data_.kernels) { - std::string filename_prefix = path_ + "/" + k; - { - std::string filename = filename_prefix + ".ll"; - std::ofstream os(filename, std::ios::out | std::ios::binary); - TI_ERROR_IF(!os.is_open(), "File {} open failed", filename); - llvm::SMDiagnostic err; - llvm::LLVMContext ctx; - llvm::raw_os_ostream llvm_os(os); - if (v.module) { - mangle_offloaded_task_name(k, v.module, v.offloaded_task_name_list); - v.module->print(llvm_os, nullptr); - } else if (v.owned_module) { - mangle_offloaded_task_name(k, v.owned_module.get(), - v.offloaded_task_name_list); - v.owned_module->print(llvm_os, nullptr); - } else - TI_ASSERT(false); - } - { - std::string filename = filename_prefix + "_otnl.txt"; - std::ofstream os(filename, std::ios::out | std::ios::binary); - TI_ERROR_IF(!os.is_open(), "File {} open failed", filename); - for (const auto &name : v.offloaded_task_name_list) { - os << name << '\n'; - } - } - } - } + void dump(); private: void mangle_offloaded_task_name( const std::string &kernel_key, llvm::Module *module, - std::vector &offloaded_task_name_list) { - if (!mangled_) { - std::size_t cnt = 0; - for (auto &e : offloaded_task_name_list) { - std::string mangled_name = kernel_key + std::to_string(cnt++); - auto func = module->getFunction(e); - TI_ASSERT(func != nullptr); - func->setName(mangled_name); - e = mangled_name; - } - } - } + std::vector &offloaded_task_name_list); std::string path_; LlvmOfflineCache data_; From 5f04fdcc43ae433cdf01b52fd7bdd9dc3aa2a116 Mon Sep 17 00:00:00 2001 From: PGZXB <420254146@qq.com> Date: Thu, 10 Mar 2022 10:22:59 +0800 Subject: [PATCH 03/10] Reset default offline_cache_file_path to ./__ticache__ --- .gitignore | 1 + taichi/llvm/llvm_offline_cache.h | 2 ++ taichi/program/compile_config.h | 2 +- 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index fd39d08f9acea..07b519a6eea36 100644 --- a/.gitignore +++ b/.gitignore @@ -85,3 +85,4 @@ _build !docs/**/*.json imgui.ini /venv/ +__ticache__ diff --git a/taichi/llvm/llvm_offline_cache.h b/taichi/llvm/llvm_offline_cache.h index d3752690a9e36..b30e85210cdb7 100644 --- a/taichi/llvm/llvm_offline_cache.h +++ b/taichi/llvm/llvm_offline_cache.h @@ -2,6 +2,7 @@ #include "taichi/common/core.h" #include "taichi/llvm/llvm_fwd.h" +#include "taichi/util/io.h" TLANG_NAMESPACE_BEGIN struct LlvmOfflineCache { @@ -36,6 +37,7 @@ class LlvmOfflineCacheFileReader { class LlvmOfflineCacheFileWriter { public: LlvmOfflineCacheFileWriter(const std::string &path) : path_(path) { + taichi::create_directories(path); } void set_data(LlvmOfflineCache &&data) { diff --git a/taichi/program/compile_config.h b/taichi/program/compile_config.h index 0c80fef5d526b..18b62d2de430d 100644 --- a/taichi/program/compile_config.h +++ b/taichi/program/compile_config.h @@ -107,7 +107,7 @@ struct CompileConfig { // Offline cache options bool offline_cache{false}; - std::string offline_cache_file_path{"."}; + std::string offline_cache_file_path{"./__ticache__"}; CompileConfig(); }; From 75286ca0a07f8665afccf8322a22ba217e658410 Mon Sep 17 00:00:00 2001 From: PGZXB <420254146@qq.com> Date: Fri, 11 Mar 2022 18:25:42 +0800 Subject: [PATCH 04/10] Fix bug, add is-grad to hashed-key --- taichi/llvm/llvm_program.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/taichi/llvm/llvm_program.cpp b/taichi/llvm/llvm_program.cpp index 07d718e4411e5..c1956b8a5f1ff 100644 --- a/taichi/llvm/llvm_program.cpp +++ b/taichi/llvm/llvm_program.cpp @@ -148,6 +148,8 @@ FunctionType LlvmProgramImpl::compile(Kernel *kernel, irpass::re_id(kernel->ir.get()); irpass::print(kernel->ir.get(), &kernel_key); picosha2::hash256_hex_string(kernel_key, hashed_kernel_key); + if (kernel->grad) hashed_kernel_key.push_back('g'); + else hashed_kernel_key.push_back('n'); auto func = this->create_kernel_function_from_offline_cache( hashed_kernel_key, kernel); if (func) { From 24d879944c557e8f0b71b779b69dc6371f71addd Mon Sep 17 00:00:00 2001 From: Taichi Gardener Date: Fri, 11 Mar 2022 10:29:26 +0000 Subject: [PATCH 05/10] Auto Format --- taichi/llvm/llvm_program.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/taichi/llvm/llvm_program.cpp b/taichi/llvm/llvm_program.cpp index c1956b8a5f1ff..55e6f58b652bc 100644 --- a/taichi/llvm/llvm_program.cpp +++ b/taichi/llvm/llvm_program.cpp @@ -148,8 +148,10 @@ FunctionType LlvmProgramImpl::compile(Kernel *kernel, irpass::re_id(kernel->ir.get()); irpass::print(kernel->ir.get(), &kernel_key); picosha2::hash256_hex_string(kernel_key, hashed_kernel_key); - if (kernel->grad) hashed_kernel_key.push_back('g'); - else hashed_kernel_key.push_back('n'); + if (kernel->grad) + hashed_kernel_key.push_back('g'); + else + hashed_kernel_key.push_back('n'); auto func = this->create_kernel_function_from_offline_cache( hashed_kernel_key, kernel); if (func) { From f4e4136f5a66821b540cb1b1853bb8dcdf953e12 Mon Sep 17 00:00:00 2001 From: PGZXB <420254146@qq.com> Date: Mon, 14 Mar 2022 14:00:50 +0800 Subject: [PATCH 06/10] Fix default offline_cache_file_path to get_repo_dir()/ticache --- taichi/program/compile_config.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/taichi/program/compile_config.h b/taichi/program/compile_config.h index 18b62d2de430d..9548851c3847c 100644 --- a/taichi/program/compile_config.h +++ b/taichi/program/compile_config.h @@ -107,7 +107,7 @@ struct CompileConfig { // Offline cache options bool offline_cache{false}; - std::string offline_cache_file_path{"./__ticache__"}; + std::string offline_cache_file_path{get_repo_dir() + "/ticache"}; CompileConfig(); }; From 6aba9219c52891b61f5909b057cad388a86b38f0 Mon Sep 17 00:00:00 2001 From: PGZXB <420254146@qq.com> Date: Mon, 14 Mar 2022 14:02:24 +0800 Subject: [PATCH 07/10] Remove __ticache__ from .gitignore --- .gitignore | 1 - 1 file changed, 1 deletion(-) diff --git a/.gitignore b/.gitignore index 07b519a6eea36..fd39d08f9acea 100644 --- a/.gitignore +++ b/.gitignore @@ -85,4 +85,3 @@ _build !docs/**/*.json imgui.ini /venv/ -__ticache__ From de919ba7006622c8f5da3ef53a234cedb8a9d864 Mon Sep 17 00:00:00 2001 From: PGZXB <420254146@qq.com> Date: Tue, 15 Mar 2022 21:44:51 +0800 Subject: [PATCH 08/10] Fix bus --- taichi/program/compile_config.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/taichi/program/compile_config.h b/taichi/program/compile_config.h index 9548851c3847c..fc5c4d8ccdcf9 100644 --- a/taichi/program/compile_config.h +++ b/taichi/program/compile_config.h @@ -107,7 +107,7 @@ struct CompileConfig { // Offline cache options bool offline_cache{false}; - std::string offline_cache_file_path{get_repo_dir() + "/ticache"}; + std::string offline_cache_file_path{get_repo_dir() + "ticache"}; CompileConfig(); }; From f24ad380138551f4f3fa81c21a996a5f28d8a2bb Mon Sep 17 00:00:00 2001 From: PGZXB <420254146@qq.com> Date: Wed, 16 Mar 2022 11:09:05 +0800 Subject: [PATCH 09/10] Move dumping-cache-file from ~LlvmProgramImpl() to ProgramImpl::dump_cache_data_to_disk() called in Program::finalize() --- taichi/llvm/llvm_program.cpp | 16 ++++++++-------- taichi/llvm/llvm_program.h | 3 ++- taichi/program/program.cpp | 1 + taichi/program/program_impl.h | 6 ++++++ 4 files changed, 17 insertions(+), 9 deletions(-) diff --git a/taichi/llvm/llvm_program.cpp b/taichi/llvm/llvm_program.cpp index 55e6f58b652bc..027a7ebba5290 100644 --- a/taichi/llvm/llvm_program.cpp +++ b/taichi/llvm/llvm_program.cpp @@ -131,14 +131,6 @@ void LlvmProgramImpl::maybe_initialize_cuda_llvm_context() { } } -LlvmProgramImpl::~LlvmProgramImpl() { - if (config->offline_cache && this->supports_offline_cache()) { - LlvmOfflineCacheFileWriter writer(config->offline_cache_file_path); - writer.set_data(std::move(cache_data_)); - writer.dump(); - } -} - FunctionType LlvmProgramImpl::compile(Kernel *kernel, OffloadedStmt *offloaded) { bool needs_cache = false; @@ -729,5 +721,13 @@ void LlvmProgramImpl::cache_kernel( kernel_cache.offloaded_task_name_list = offloaded_task_name_list; } +void LlvmProgramImpl::dump_cache_data_to_disk() { + if (config->offline_cache && this->supports_offline_cache()) { + LlvmOfflineCacheFileWriter writer(config->offline_cache_file_path); + writer.set_data(std::move(cache_data_)); + writer.dump(); + } +} + } // namespace lang } // namespace taichi diff --git a/taichi/llvm/llvm_program.h b/taichi/llvm/llvm_program.h index 248e5acaeb6f3..2134b070956fc 100644 --- a/taichi/llvm/llvm_program.h +++ b/taichi/llvm/llvm_program.h @@ -38,7 +38,6 @@ class CpuDevice; class LlvmProgramImpl : public ProgramImpl { public: LlvmProgramImpl(CompileConfig &config, KernelProfilerBase *profiler); - ~LlvmProgramImpl() override; void initialize_host(); @@ -173,6 +172,8 @@ class LlvmProgramImpl : public ProgramImpl { DevicePtr get_snode_tree_device_ptr(int tree_id) override; + void dump_cache_data_to_disk() override; + private: std::unique_ptr llvm_context_host_{nullptr}; std::unique_ptr llvm_context_device_{nullptr}; diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp index 48fd81c676000..c0dfb3326752a 100644 --- a/taichi/program/program.cpp +++ b/taichi/program/program.cpp @@ -518,6 +518,7 @@ void Program::finalize() { finalized_ = true; num_instances_ -= 1; + program_impl_->dump_cache_data_to_disk(); TI_TRACE("Program ({}) finalized_.", fmt::ptr(this)); } diff --git a/taichi/program/program_impl.h b/taichi/program/program_impl.h index 3ba3645598a70..6ba07b56c850f 100644 --- a/taichi/program/program_impl.h +++ b/taichi/program/program_impl.h @@ -65,6 +65,12 @@ class ProgramImpl { */ virtual std::unique_ptr make_aot_module_builder() = 0; + /** + * Dump Offline-cache data to disk + */ + virtual void dump_cache_data_to_disk() { + } + virtual Device *get_compute_device() { return nullptr; } From 5c885b1a1dd746d9aeb64439737343ffd06bb27d Mon Sep 17 00:00:00 2001 From: PGZXB <420254146@qq.com> Date: Wed, 16 Mar 2022 18:23:25 +0800 Subject: [PATCH 10/10] Fix TLANG_NAMESPACE_BEGIN to namespace taichi { namespace lang { --- taichi/llvm/llvm_offline_cache.cpp | 6 ++++-- taichi/llvm/llvm_offline_cache.h | 7 +++++-- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/taichi/llvm/llvm_offline_cache.cpp b/taichi/llvm/llvm_offline_cache.cpp index 99a93d67c99e9..595544a599f57 100644 --- a/taichi/llvm/llvm_offline_cache.cpp +++ b/taichi/llvm/llvm_offline_cache.cpp @@ -5,7 +5,8 @@ #include "llvm/Support/raw_os_ostream.h" #include "llvm/IR/Module.h" -TLANG_NAMESPACE_BEGIN +namespace taichi { +namespace lang { bool LlvmOfflineCacheFileReader::get_kernel_cache( LlvmOfflineCache::KernelCacheData &res, @@ -84,4 +85,5 @@ void LlvmOfflineCacheFileWriter::mangle_offloaded_task_name( } } -TLANG_NAMESPACE_END +} // namespace lang +} // namespace taichi diff --git a/taichi/llvm/llvm_offline_cache.h b/taichi/llvm/llvm_offline_cache.h index b30e85210cdb7..fafd4d08bae5e 100644 --- a/taichi/llvm/llvm_offline_cache.h +++ b/taichi/llvm/llvm_offline_cache.h @@ -4,7 +4,9 @@ #include "taichi/llvm/llvm_fwd.h" #include "taichi/util/io.h" -TLANG_NAMESPACE_BEGIN +namespace taichi { +namespace lang { + struct LlvmOfflineCache { struct KernelCacheData { std::string kernel_key; @@ -63,4 +65,5 @@ class LlvmOfflineCacheFileWriter { bool mangled_{false}; }; -TLANG_NAMESPACE_END +} // namespace lang +} // namespace taichi