[refactor] Get rid of unnecessary snodes map.

taichi-dev · Oct 20, 2021 · a4cf4d4 · a4cf4d4
1 parent 8c5c147
commit a4cf4d4
Show file tree

Hide file tree

Showing 17 changed files with 12 additions and 37 deletions.
diff --git a/taichi/backends/cc/cc_program.cpp b/taichi/backends/cc/cc_program.cpp
@@ -35,7 +35,6 @@ void CCProgramImpl::materialize_runtime(MemoryPool *memory_pool,
 void CCProgramImpl::materialize_snode_tree(
     SNodeTree *tree,
     std::vector<std::unique_ptr<SNodeTree>> &,
-    std::unordered_map<int, SNode *> &,
     uint64 *result_buffer) {
   auto *const root = tree->root();
   CCLayoutGen gen(this, root);

diff --git a/taichi/backends/cc/cc_program.h b/taichi/backends/cc/cc_program.h
@@ -43,7 +43,6 @@ class CCProgramImpl : public ProgramImpl {
 
   void materialize_snode_tree(SNodeTree *tree,
                               std::vector<std::unique_ptr<SNodeTree>> &,
-                              std::unordered_map<int, SNode *> &,
                               uint64 *result_buffer) override;
 
   void synchronize() override {

diff --git a/taichi/backends/metal/metal_program.cpp b/taichi/backends/metal/metal_program.cpp
@@ -46,7 +46,6 @@ void MetalProgramImpl::materialize_runtime(MemoryPool *memory_pool,
 void MetalProgramImpl::materialize_snode_tree(
     SNodeTree *tree,
     std::vector<std::unique_ptr<SNodeTree>> &,
-    std::unordered_map<int, SNode *> &,
     uint64 *result_buffer) {
   // TODO: support materializing multiple snode trees
   TI_ASSERT_INFO(config->use_llvm,

diff --git a/taichi/backends/metal/metal_program.h b/taichi/backends/metal/metal_program.h
@@ -33,7 +33,6 @@ class MetalProgramImpl : public ProgramImpl {
   void materialize_snode_tree(
       SNodeTree *tree,
       std::vector<std::unique_ptr<SNodeTree>> &snode_trees_,
-      std::unordered_map<int, SNode *> &snodes,
       uint64 *result_buffer) override;
 
   void synchronize() override {

diff --git a/taichi/backends/opengl/opengl_program.cpp b/taichi/backends/opengl/opengl_program.cpp
@@ -41,7 +41,6 @@ void OpenglProgramImpl::compile_snode_tree_types(SNodeTree *tree) {
 void OpenglProgramImpl::materialize_snode_tree(
     SNodeTree *tree,
     std::vector<std::unique_ptr<SNodeTree>> &,
-    std::unordered_map<int, SNode *> &,
     uint64 *result_buffer) {
 #ifdef TI_WITH_OPENGL
   compile_snode_tree_types(tree);

diff --git a/taichi/backends/opengl/opengl_program.h b/taichi/backends/opengl/opengl_program.h
@@ -38,7 +38,6 @@ class OpenglProgramImpl : public ProgramImpl {
   void materialize_snode_tree(
       SNodeTree *tree,
       std::vector<std::unique_ptr<SNodeTree>> &snode_trees_,
-      std::unordered_map<int, SNode *> &snodes,
       uint64 *result_buffer) override;
 
   void synchronize() override {

diff --git a/taichi/backends/vulkan/vulkan_program.cpp b/taichi/backends/vulkan/vulkan_program.cpp
@@ -97,7 +97,6 @@ void VulkanProgramImpl::materialize_runtime(MemoryPool *memory_pool,
 void VulkanProgramImpl::materialize_snode_tree(
     SNodeTree *tree,
     std::vector<std::unique_ptr<SNodeTree>> &,
-    std::unordered_map<int, SNode *> &,
     uint64 *result_buffer) {
   vulkan_runtime_->materialize_snode_tree(tree);
 }

diff --git a/taichi/backends/vulkan/vulkan_program.h b/taichi/backends/vulkan/vulkan_program.h
@@ -42,7 +42,6 @@ class VulkanProgramImpl : public ProgramImpl {
 
   void materialize_snode_tree(SNodeTree *tree,
                               std::vector<std::unique_ptr<SNodeTree>> &,
-                              std::unordered_map<int, SNode *> &,
                               uint64 *result_buffer) override;
 
   void synchronize() override {

diff --git a/taichi/llvm/llvm_program.cpp b/taichi/llvm/llvm_program.cpp
@@ -218,20 +218,16 @@ void LlvmProgramImpl::initialize_llvm_runtime_snodes(const SNodeTree *tree,
 void LlvmProgramImpl::materialize_snode_tree(
     SNodeTree *tree,
     std::vector<std::unique_ptr<SNodeTree>> &snode_trees_,
-    std::unordered_map<int, SNode *> &snodes,
     uint64 *result_buffer) {
   auto *const root = tree->root();
-  auto host_module = clone_struct_compiler_initial_context(
-      snode_trees_, llvm_context_host.get());
-  std::unique_ptr<StructCompiler> scomp = std::make_unique<StructCompilerLLVM>(
-      host_arch(), this, std::move(host_module));
-  scomp->run(*root);
-
-  for (auto snode : scomp->snodes) {
-    snodes[snode->id] = snode;
-  }
-
   if (arch_is_cpu(config->arch)) {
+    auto host_module = clone_struct_compiler_initial_context(
+        snode_trees_, llvm_context_host.get());
+    std::unique_ptr<StructCompiler> scomp =
+        std::make_unique<StructCompilerLLVM>(host_arch(), this,
+                                             std::move(host_module));
+    scomp->run(*root);
+
     initialize_llvm_runtime_snodes(tree, scomp.get(), result_buffer);
   } else if (config->arch == Arch::cuda) {
     auto device_module = clone_struct_compiler_initial_context(

diff --git a/taichi/llvm/llvm_program.h b/taichi/llvm/llvm_program.h
@@ -60,7 +60,6 @@ class LlvmProgramImpl : public ProgramImpl {
   void materialize_snode_tree(
       SNodeTree *tree,
       std::vector<std::unique_ptr<SNodeTree>> &snode_trees_,
-      std::unordered_map<int, SNode *> &snodes,
       uint64 *result_buffer) override;
 
   template <typename T>

diff --git a/taichi/program/async_engine.cpp b/taichi/program/async_engine.cpp
@@ -188,11 +188,10 @@ ExecutionQueue::ExecutionQueue(
 }
 
 AsyncEngine::AsyncEngine(const CompileConfig *const config,
-                         const std::unordered_map<int, SNode *> &snodes,
                          const BackendExecCompilationFunc &compile_to_backend)
     : queue(&ir_bank_, compile_to_backend),
       config_(config),
-      sfg(std::make_unique<StateFlowGraph>(this, &ir_bank_, config, snodes)) {
+      sfg(std::make_unique<StateFlowGraph>(this, &ir_bank_, config)) {
   Timeline::get_this_thread_instance().set_name("host");
   ir_bank_.set_sfg(sfg.get());
 }

diff --git a/taichi/program/async_engine.h b/taichi/program/async_engine.h
@@ -138,7 +138,6 @@ class AsyncEngine {
   std::unique_ptr<StateFlowGraph> sfg;
 
   explicit AsyncEngine(const CompileConfig *const config,
-                       const std::unordered_map<int, SNode *> &snodes,
                        const BackendExecCompilationFunc &compile_to_backend);
 
   void clear_cache() {

diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp
@@ -138,7 +138,7 @@ Program::Program(Arch desired_arch)
     TI_WARN("Running in async mode. This is experimental.");
     TI_ASSERT(is_extension_supported(config.arch, Extension::async_mode));
     async_engine = std::make_unique<AsyncEngine>(
-        &config, snodes, [this](Kernel &kernel, OffloadedStmt *offloaded) {
+        &config, [this](Kernel &kernel, OffloadedStmt *offloaded) {
           return this->compile(kernel, offloaded);
         });
   }
@@ -201,7 +201,7 @@ SNodeTree *Program::add_snode_tree(std::unique_ptr<SNode> root,
   if (compile_only) {
     program_impl_->compile_snode_tree_types(tree.get());
   } else {
-    program_impl_->materialize_snode_tree(tree.get(), snode_trees_, snodes,
+    program_impl_->materialize_snode_tree(tree.get(), snode_trees_,
                                           result_buffer);
   }
   snode_trees_.push_back(std::move(tree));

diff --git a/taichi/program/program.h b/taichi/program/program.h
@@ -105,9 +105,6 @@ class Program {
 
   uint64 *result_buffer{nullptr};  // Note result_buffer is used by all backends
 
-  std::unordered_map<int, SNode *>
-      snodes;  // TODO: seems LLVM specific but used by state_flow_graph.cpp.
-
   std::unique_ptr<AsyncEngine> async_engine{nullptr};
 
   std::vector<std::unique_ptr<Kernel>> kernels;

diff --git a/taichi/program/program_impl.h b/taichi/program/program_impl.h
@@ -43,7 +43,6 @@ class ProgramImpl {
   virtual void materialize_snode_tree(
       SNodeTree *tree,
       std::vector<std::unique_ptr<SNodeTree>> &snode_trees_,
-      std::unordered_map<int, SNode *> &snodes,
       uint64 *result_buffer_ptr) = 0;
 
   virtual void destroy_snode_tree(SNodeTree *snode_tree) = 0;

diff --git a/taichi/program/state_flow_graph.cpp b/taichi/program/state_flow_graph.cpp
@@ -240,8 +240,7 @@ void StateFlowGraph::Node::disconnect_with(StateFlowGraph::Node *other) {
 
 StateFlowGraph::StateFlowGraph(AsyncEngine *engine,
                                IRBank *ir_bank,
-                               const CompileConfig *const config,
-                               const std::unordered_map<int, SNode *> &snodes)
+                               const CompileConfig *const config)
     : first_pending_task_index_(1 /*after initial node*/),
       ir_bank_(ir_bank),
       engine_(engine),
@@ -255,10 +254,6 @@ StateFlowGraph::StateFlowGraph(AsyncEngine *engine,
   initial_node_->input_edges.node_id = 0;
   initial_node_->output_edges.node_id = 0;
   initial_node_->mark_executed();
-
-  for (const auto snode : snodes) {
-    list_up_to_date_[snode.second] = false;
-  }
 }
 
 std::vector<StateFlowGraph::Node *> StateFlowGraph::get_pending_tasks() const {

diff --git a/taichi/program/state_flow_graph.h b/taichi/program/state_flow_graph.h
@@ -227,8 +227,7 @@ class StateFlowGraph {
 
   StateFlowGraph(AsyncEngine *engine,
                  IRBank *ir_bank,
-                 const CompileConfig *const config,
-                 const std::unordered_map<int, SNode *> &snodes);
+                 const CompileConfig *const config);
 
   std::vector<Node *> get_pending_tasks() const;