From c40e19a678f32e93f9ef8ab534da364cac24b8f6 Mon Sep 17 00:00:00 2001 From: Zhanlue Yang Date: Wed, 13 Jul 2022 09:37:28 +0800 Subject: [PATCH] [llvm] [aot] Add LLVM to CAPI part 1: Implemented capi::LlvmRuntime class (#5393) * [llvm] [aot] Add LLVM to CAPI part 1: Implement capi::LlvmRuntime class * Added LlvmRuntime initialization to ti_create_runtime() * Fixed issue with use of memory pool * Fixed minor issue --- c_api/src/taichi_core_impl.cpp | 12 +++++ c_api/src/taichi_core_impl.h | 9 ++-- c_api/src/taichi_llvm_impl.cpp | 57 ++++++++++++++++++++++++ c_api/src/taichi_llvm_impl.h | 41 +++++++++++++++++ taichi/codegen/llvm/codegen_llvm.h | 27 +---------- taichi/codegen/llvm/llvm_compiled_data.h | 37 +++++++++++++++ taichi/runtime/llvm/llvm_offline_cache.h | 3 +- 7 files changed, 156 insertions(+), 30 deletions(-) create mode 100644 c_api/src/taichi_llvm_impl.cpp create mode 100644 c_api/src/taichi_llvm_impl.h create mode 100644 taichi/codegen/llvm/llvm_compiled_data.h diff --git a/c_api/src/taichi_core_impl.cpp b/c_api/src/taichi_core_impl.cpp index bbcd69f928998..ce5ee136971ce 100644 --- a/c_api/src/taichi_core_impl.cpp +++ b/c_api/src/taichi_core_impl.cpp @@ -1,5 +1,6 @@ #include "taichi_core_impl.h" #include "taichi_vulkan_impl.h" +#include "taichi_llvm_impl.h" #include "taichi/program/ndarray.h" Runtime::Runtime(taichi::Arch arch) : arch(arch) { @@ -47,6 +48,17 @@ TiRuntime ti_create_runtime(TiArch arch) { case TI_ARCH_VULKAN: return (TiRuntime)(static_cast(new VulkanRuntimeOwned)); #endif // TI_WITH_VULKAN +#ifdef TI_WITH_LLVM + case TI_ARCH_X64: + return (TiRuntime)(static_cast( + new capi::LlvmRuntime(taichi::Arch::x64))); + case TI_ARCH_ARM64: + return (TiRuntime)(static_cast( + new capi::LlvmRuntime(taichi::Arch::arm64))); + case TI_ARCH_CUDA: + return (TiRuntime)(static_cast( + new capi::LlvmRuntime(taichi::Arch::cuda))); +#endif // TI_WITH_LLVM default: TI_WARN("ignored attempt to create runtime on unknown arch"); return TI_NULL_HANDLE; diff --git a/c_api/src/taichi_core_impl.h b/c_api/src/taichi_core_impl.h index 94a3fc7238111..3c6766b209977 100644 --- a/c_api/src/taichi_core_impl.h +++ b/c_api/src/taichi_core_impl.h @@ -54,12 +54,15 @@ class AotModule { namespace { -taichi::lang::DeviceAllocation devmem2devalloc(Runtime &runtime, - TiMemory devmem) { +[[maybe_unused]] taichi::lang::DeviceAllocation devmem2devalloc( + Runtime &runtime, + TiMemory devmem) { return taichi::lang::DeviceAllocation{ &runtime.get(), (taichi::lang::DeviceAllocationId)((size_t)devmem - 1)}; } -TiMemory devalloc2devmem(const taichi::lang::DeviceAllocation &devalloc) { + +[[maybe_unused]] TiMemory devalloc2devmem( + const taichi::lang::DeviceAllocation &devalloc) { return (TiMemory)((size_t)devalloc.alloc_id + 1); } diff --git a/c_api/src/taichi_llvm_impl.cpp b/c_api/src/taichi_llvm_impl.cpp new file mode 100644 index 0000000000000..cefba2ce1e658 --- /dev/null +++ b/c_api/src/taichi_llvm_impl.cpp @@ -0,0 +1,57 @@ +#include "taichi_core_impl.h" +#include "taichi_llvm_impl.h" + +#ifdef TI_WITH_LLVM + +#include "taichi/runtime/llvm/llvm_runtime_executor.h" + +namespace capi { + +LlvmRuntime::LlvmRuntime(taichi::Arch arch) : Runtime(arch) { + taichi::lang::CompileConfig cfg; + cfg.arch = arch; + executor_ = std::make_unique(cfg, nullptr); + + taichi::lang::Device *compute_device = executor_->get_compute_device(); + memory_pool_ = + taichi::arch_is_cpu(arch) + ? std::make_unique(arch, compute_device) + : nullptr; + + // materialize_runtime() takes in a uint64_t** (pointer object's address) and + // modifies the address it points to. + // + // Therefore we can't use host_result_buffer_.data() here, + // since it returns a temporary copy of the internal data pointer, + // thus we won't be able to modify the address where the std::array's data + // pointer is pointing to. + executor_->materialize_runtime(memory_pool_.get(), nullptr /*kNoProfiler*/, + &result_buffer); +} + +taichi::lang::Device &LlvmRuntime::get() { + taichi::lang::Device *device = executor_->get_compute_device(); + return *device; +} + +TiAotModule LlvmRuntime::load_aot_module(const char *module_path) { + TI_NOT_IMPLEMENTED; +} + +void LlvmRuntime::buffer_copy(const taichi::lang::DevicePtr &dst, + const taichi::lang::DevicePtr &src, + size_t size) { + TI_NOT_IMPLEMENTED; +} + +void LlvmRuntime::submit() { + TI_NOT_IMPLEMENTED; +} + +void LlvmRuntime::wait() { + TI_NOT_IMPLEMENTED; +} + +} // namespace capi + +#endif // TI_WITH_LLVM diff --git a/c_api/src/taichi_llvm_impl.h b/c_api/src/taichi_llvm_impl.h new file mode 100644 index 0000000000000..6298ca8bc2ad4 --- /dev/null +++ b/c_api/src/taichi_llvm_impl.h @@ -0,0 +1,41 @@ +#pragma once + +#ifdef TI_WITH_LLVM +#include "taichi_core_impl.h" + +namespace taichi { +namespace lang { +class LlvmRuntimeExecutor; +class MemoryPool; +} // namespace lang +} // namespace taichi + +namespace capi { + +class LlvmRuntime : public Runtime { + public: + LlvmRuntime(taichi::Arch arch); + + private: + /* Internally used interfaces */ + taichi::lang::Device &get() override; + + TiAotModule load_aot_module(const char *module_path) override; + + void buffer_copy(const taichi::lang::DevicePtr &dst, + const taichi::lang::DevicePtr &src, + size_t size) override; + + void submit() override; + + void wait() override; + + private: + taichi::uint64 *result_buffer{nullptr}; + std::unique_ptr executor_{nullptr}; + std::unique_ptr memory_pool_{nullptr}; +}; + +} // namespace capi + +#endif // TI_WITH_LLVM diff --git a/taichi/codegen/llvm/codegen_llvm.h b/taichi/codegen/llvm/codegen_llvm.h index 8d41f3ea1d2fd..9cb97fd056f49 100644 --- a/taichi/codegen/llvm/codegen_llvm.h +++ b/taichi/codegen/llvm/codegen_llvm.h @@ -9,6 +9,7 @@ #include "taichi/ir/ir.h" #include "taichi/runtime/llvm/launch_arg_info.h" #include "taichi/codegen/llvm/llvm_codegen_utils.h" +#include "taichi/codegen/llvm/llvm_compiled_data.h" #include "taichi/program/program.h" namespace taichi { @@ -16,19 +17,6 @@ namespace lang { class CodeGenLLVM; -class OffloadedTask { - public: - std::string name; - int block_dim{0}; - int grid_dim{0}; - - OffloadedTask(const std::string &name = "", - int block_dim = 0, - int grid_dim = 0) - : name(name), block_dim(block_dim), grid_dim(grid_dim){}; - TI_IO_DEF(name, block_dim, grid_dim); -}; - class FunctionCreationGuard { public: CodeGenLLVM *mb; @@ -42,19 +30,6 @@ class FunctionCreationGuard { ~FunctionCreationGuard(); }; -struct LLVMCompiledData { - std::vector tasks; - std::unique_ptr module{nullptr}; - LLVMCompiledData() = default; - LLVMCompiledData(LLVMCompiledData &&) = default; - LLVMCompiledData(std::vector tasks, - std::unique_ptr module) - : tasks(std::move(tasks)), module(std::move(module)) { - } - LLVMCompiledData clone() const; - TI_IO_DEF(tasks); -}; - class CodeGenLLVM : public IRVisitor, public LLVMModuleBuilder { public: Kernel *kernel; diff --git a/taichi/codegen/llvm/llvm_compiled_data.h b/taichi/codegen/llvm/llvm_compiled_data.h new file mode 100644 index 0000000000000..17bde818b3135 --- /dev/null +++ b/taichi/codegen/llvm/llvm_compiled_data.h @@ -0,0 +1,37 @@ +#pragma once + +#include + +#include "llvm/IR/Module.h" + +namespace taichi { +namespace lang { + +class OffloadedTask { + public: + std::string name; + int block_dim{0}; + int grid_dim{0}; + + OffloadedTask(const std::string &name = "", + int block_dim = 0, + int grid_dim = 0) + : name(name), block_dim(block_dim), grid_dim(grid_dim){}; + TI_IO_DEF(name, block_dim, grid_dim); +}; + +struct LLVMCompiledData { + std::vector tasks; + std::unique_ptr module{nullptr}; + LLVMCompiledData() = default; + LLVMCompiledData(LLVMCompiledData &&) = default; + LLVMCompiledData(std::vector tasks, + std::unique_ptr module) + : tasks(std::move(tasks)), module(std::move(module)) { + } + LLVMCompiledData clone() const; + TI_IO_DEF(tasks); +}; + +} // namespace lang +} // namespace taichi diff --git a/taichi/runtime/llvm/llvm_offline_cache.h b/taichi/runtime/llvm/llvm_offline_cache.h index 8535fbb3abd73..4da6e9acadf52 100644 --- a/taichi/runtime/llvm/llvm_offline_cache.h +++ b/taichi/runtime/llvm/llvm_offline_cache.h @@ -8,7 +8,8 @@ #include "taichi/runtime/llvm/launch_arg_info.h" #include "taichi/program/kernel.h" #include "taichi/util/io.h" -#include "taichi/codegen/llvm/codegen_llvm.h" +#include "taichi/codegen/llvm/llvm_compiled_data.h" + namespace taichi { namespace lang {