[llvm] [aot] Add LLVM to CAPI part 1: Implemented capi::LlvmRuntime c…

…lass (#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

@@ -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<Runtime *>(new VulkanRuntimeOwned));
 #endif  // TI_WITH_VULKAN
+#ifdef TI_WITH_LLVM
+    case TI_ARCH_X64:
+      return (TiRuntime)(static_cast<Runtime *>(
+          new capi::LlvmRuntime(taichi::Arch::x64)));
+    case TI_ARCH_ARM64:
+      return (TiRuntime)(static_cast<Runtime *>(
+          new capi::LlvmRuntime(taichi::Arch::arm64)));
+    case TI_ARCH_CUDA:
+      return (TiRuntime)(static_cast<Runtime *>(
+          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;

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);
 }
 

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<taichi::lang::LlvmRuntimeExecutor>(cfg, nullptr);
+
+  taichi::lang::Device *compute_device = executor_->get_compute_device();
+  memory_pool_ =
+      taichi::arch_is_cpu(arch)
+          ? std::make_unique<taichi::lang::MemoryPool>(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

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<taichi::lang::LlvmRuntimeExecutor> executor_{nullptr};
+  std::unique_ptr<taichi::lang::MemoryPool> memory_pool_{nullptr};
+};
+
+}  // namespace capi
+
+#endif  // TI_WITH_LLVM

taichi/codegen/llvm/codegen_llvm.h

@@ -9,26 +9,14 @@
 #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 {
 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<OffloadedTask> tasks;
-  std::unique_ptr<llvm::Module> module{nullptr};
-  LLVMCompiledData() = default;
-  LLVMCompiledData(LLVMCompiledData &&) = default;
-  LLVMCompiledData(std::vector<OffloadedTask> tasks,
-                   std::unique_ptr<llvm::Module> 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;

taichi/codegen/llvm/llvm_compiled_data.h

@@ -0,0 +1,37 @@
+#pragma once
+
+#include <memory>
+
+#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<OffloadedTask> tasks;
+  std::unique_ptr<llvm::Module> module{nullptr};
+  LLVMCompiledData() = default;
+  LLVMCompiledData(LLVMCompiledData &&) = default;
+  LLVMCompiledData(std::vector<OffloadedTask> tasks,
+                   std::unique_ptr<llvm::Module> module)
+      : tasks(std::move(tasks)), module(std::move(module)) {
+  }
+  LLVMCompiledData clone() const;
+  TI_IO_DEF(tasks);
+};
+
+}  // namespace lang
+}  // namespace taichi

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 {