[aot] [vulkan] Add AotKernel and its Vulkan impl (#4387)

* [aot] [vulkan] Add AotKernel and its Vulkan impl

* rename to launch

* Auto Format

* fix win

* fix win

* im damned

* weird...

* fix

* fix

* maybe fix

* Auto Format

* horrible c++

* the final fix

Co-authored-by: Taichi Gardener <[email protected]>

Author: k-ye

taichi/aot/module_loader.cpp

@@ -0,0 +1,18 @@
+#include "taichi/aot/module_loader.h"
+
+namespace taichi {
+namespace lang {
+
+AotKernel *AotModuleLoader::get_kernel(const std::string &name) {
+  auto itr = loaded_kernels_.find(name);
+  if (itr != loaded_kernels_.end()) {
+    return itr->second.get();
+  }
+  auto k = make_new_kernel(name);
+  auto *kptr = k.get();
+  loaded_kernels_[name] = std::move(k);
+  return kptr;
+}
+
+}  // namespace lang
+}  // namespace taichi

taichi/aot/module_loader.h

@@ -1,6 +1,8 @@
 #pragma once
 
+#include <memory>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include "taichi/aot/module_data.h"
@@ -11,17 +13,62 @@
 namespace taichi {
 namespace lang {
 
-class AotModuleLoader {
+class RuntimeContext;
+
+// TODO: Instead of prefixing all these classes with "Aot", just put them into
+// the `aot` namespace.
+class TI_DLL_EXPORT AotKernel {
+ public:
+  // Rule of 5 to make MSVC happy
+  AotKernel() = default;
+  virtual ~AotKernel() = default;
+  AotKernel(const AotKernel &) = delete;
+  AotKernel &operator=(const AotKernel &) = delete;
+  AotKernel(AotKernel &&) = default;
+  AotKernel &operator=(AotKernel &&) = default;
+
+  /**
+   * @brief Launches the kernel to the device
+   *
+   * This does not manage the device to host synchronization.
+   *
+   * @param ctx Host context
+   */
+  virtual void launch(RuntimeContext *ctx) = 0;
+};
+
+class TI_DLL_EXPORT AotModuleLoader {
  public:
+  // Rule of 5 to make MSVC happy
+  AotModuleLoader() = default;
   virtual ~AotModuleLoader() = default;
+  AotModuleLoader(const AotModuleLoader &) = delete;
+  AotModuleLoader &operator=(const AotModuleLoader &) = delete;
+  AotModuleLoader(AotModuleLoader &&) = default;
+  AotModuleLoader &operator=(AotModuleLoader &&) = default;
 
   // TODO: Add method get_kernel(...) once the kernel field data will be
   // generic/common across all backends.
 
   virtual bool get_field(const std::string &name,
                          aot::CompiledFieldData &field) = 0;
 
+  /**
+   * @brief Get the kernel object
+   *
+   * @param name Name of the kernel
+   * @return AotKernel*
+   */
+  AotKernel *get_kernel(const std::string &name);
+
   virtual size_t get_root_size() const = 0;
+
+ protected:
+  virtual std::unique_ptr<AotKernel> make_new_kernel(
+      const std::string &name) = 0;
+
+ private:
+  std::unordered_map<std::string, std::unique_ptr<AotKernel>> loaded_kernels_;
 };
 
 // Only responsible for reporting device capabilities

taichi/backends/vulkan/aot_module_loader_impl.cpp

@@ -3,9 +3,30 @@
 #include <fstream>
 #include <type_traits>
 
+#include "taichi/backends/vulkan/runtime.h"
+
 namespace taichi {
 namespace lang {
 namespace vulkan {
+namespace {
+
+using KernelHandle = VkRuntime::KernelHandle;
+
+class KernelImpl : public AotKernel {
+ public:
+  explicit KernelImpl(VkRuntime *runtime, KernelHandle handle)
+      : runtime_(runtime), handle_(handle) {
+  }
+
+  void launch(RuntimeContext *ctx) override {
+    runtime_->launch_kernel(handle_, ctx);
+  }
+
+ private:
+  VkRuntime *const runtime_;
+  const KernelHandle handle_;
+};
+}  // namespace
 
 AotModuleLoaderImpl::AotModuleLoaderImpl(const std::string &output_dir) {
   const std::string bin_path = fmt::format("{}/metadata.tcb", output_dir);
@@ -55,6 +76,17 @@ bool AotModuleLoaderImpl::get_kernel(const std::string &name,
   return false;
 }
 
+std::unique_ptr<AotKernel> AotModuleLoaderImpl::make_new_kernel(
+    const std::string &name) {
+  VkRuntime::RegisterParams kparams;
+  if (!get_kernel(name, kparams)) {
+    TI_DEBUG("Failed to load kernel {}", name);
+    return nullptr;
+  }
+  auto handle = runtime_->register_taichi_kernel(kparams);
+  return std::make_unique<KernelImpl>(runtime_, handle);
+}
+
 bool AotModuleLoaderImpl::get_field(const std::string &name,
                                     aot::CompiledFieldData &field) {
   TI_ERROR("AOT: get_field for Vulkan not implemented yet");

taichi/backends/vulkan/aot_module_loader_impl.h

@@ -13,6 +13,8 @@ namespace taichi {
 namespace lang {
 namespace vulkan {
 
+class VkRuntime;
+
 class TI_DLL_EXPORT AotModuleLoaderImpl : public AotModuleLoader {
  public:
   explicit AotModuleLoaderImpl(const std::string &output_dir);
@@ -25,10 +27,12 @@ class TI_DLL_EXPORT AotModuleLoaderImpl : public AotModuleLoader {
   size_t get_root_size() const override;
 
  private:
+  std::unique_ptr<AotKernel> make_new_kernel(const std::string &name) override;
   std::vector<uint32_t> read_spv_file(const std::string &output_dir,
                                       const TaskAttributes &k);
 
   TaichiAotData ti_aot_data_;
+  VkRuntime *runtime_{nullptr};
 };
 
 }  // namespace vulkan