[metal] Rearrange how KernelManager is initialized (taichi-dev#3109)

firedtoad · Oct 8, 2021 · ed5d899 · ed5d899
1 parent 4908d5b
commit ed5d899
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 15 deletions.
diff --git a/taichi/backends/metal/kernel_manager.cpp b/taichi/backends/metal/kernel_manager.cpp
@@ -1122,6 +1122,7 @@ class KernelManager::Impl {
   nsobj_unique_ptr<MTLBuffer> global_tmps_buffer_{nullptr};
   std::unique_ptr<BufferMemoryView> runtime_mem_{nullptr};
   nsobj_unique_ptr<MTLBuffer> runtime_buffer_{nullptr};
+  int last_snode_id_used_in_runtime_{-1};
   // TODO: Rename these to 'print_assert_{mem|buffer}_'
   std::unique_ptr<BufferMemoryView> print_mem_{nullptr};
   nsobj_unique_ptr<MTLBuffer> print_buffer_{nullptr};

diff --git a/taichi/backends/metal/kernel_manager.h b/taichi/backends/metal/kernel_manager.h
@@ -27,12 +27,10 @@ class KernelManager {
  public:
   struct Params {
     CompiledRuntimeModule compiled_runtime_module;
-    CompiledStructs compiled_structs;
     CompileConfig *config;
-    MemoryPool *mem_pool;
     uint64_t *host_result_buffer;
+    MemoryPool *mem_pool;
     KernelProfilerBase *profiler;
-    int root_id;
   };
 
   explicit KernelManager(Params params);

diff --git a/taichi/backends/metal/metal_program.cpp b/taichi/backends/metal/metal_program.cpp
@@ -29,11 +29,18 @@ void MetalProgramImpl::materialize_runtime(MemoryPool *memory_pool,
                                            KernelProfilerBase *profiler,
                                            uint64 **result_buffer_ptr) {
   TI_ASSERT(*result_buffer_ptr == nullptr);
+  TI_ASSERT(metal_kernel_mgr_ == nullptr);
   *result_buffer_ptr = (uint64 *)memory_pool->allocate(
       sizeof(uint64) * taichi_result_buffer_entries, 8);
-  params_.mem_pool = memory_pool;
-  params_.profiler = profiler;
   compiled_runtime_module_ = metal::compile_runtime_module();
+
+  metal::KernelManager::Params params;
+  params.compiled_runtime_module = compiled_runtime_module_.value();
+  params.config = config;
+  params.host_result_buffer = *result_buffer_ptr;
+  params.mem_pool = memory_pool;
+  params.profiler = profiler;
+  metal_kernel_mgr_ = std::make_unique<metal::KernelManager>(std::move(params));
 }
 
 void MetalProgramImpl::materialize_snode_tree(
@@ -47,15 +54,6 @@ void MetalProgramImpl::materialize_snode_tree(
   auto *const root = tree->root();
 
   metal_compiled_structs_ = metal::compile_structs(*root);
-  if (metal_kernel_mgr_ == nullptr) {
-    params_.compiled_structs = metal_compiled_structs_.value();
-    params_.compiled_runtime_module = compiled_runtime_module_.value();
-    params_.config = config;
-    params_.host_result_buffer = result_buffer;
-    params_.root_id = root->id;
-    metal_kernel_mgr_ =
-        std::make_unique<metal::KernelManager>(std::move(params_));
-  }
   metal_kernel_mgr_->add_compiled_snode_tree(metal_compiled_structs_.value());
 }
 

diff --git a/taichi/backends/metal/metal_program.h b/taichi/backends/metal/metal_program.h
@@ -1,4 +1,7 @@
 #pragma once
+
+#include <vector>
+
 #include "taichi/backends/metal/kernel_manager.h"
 #include "taichi/backends/metal/struct_metal.h"
 #include "taichi/system/memory_pool.h"
@@ -47,8 +50,8 @@ class MetalProgramImpl : public ProgramImpl {
   std::optional<metal::CompiledRuntimeModule> compiled_runtime_module_{
       std::nullopt};
   std::optional<metal::CompiledStructs> metal_compiled_structs_{std::nullopt};
+  std::vector<metal::CompiledStructs> compiled_snode_trees_;
   std::unique_ptr<metal::KernelManager> metal_kernel_mgr_{nullptr};
-  metal::KernelManager::Params params_;
 };
 
 }  // namespace lang

diff --git a/taichi/program/program_impl.h b/taichi/program/program_impl.h
@@ -9,6 +9,7 @@
 
 namespace taichi {
 namespace lang {
+
 class ProgramImpl {
  public:
   // TODO: Make it safer, we exposed it for now as it's directly accessed