From a4cf4d4fb46b5e743f664645144d4785d1384939 Mon Sep 17 00:00:00 2001 From: Ailing Zhang Date: Wed, 20 Oct 2021 05:06:21 -0400 Subject: [PATCH] [refactor] Get rid of unnecessary snodes map. --- taichi/backends/cc/cc_program.cpp | 1 - taichi/backends/cc/cc_program.h | 1 - taichi/backends/metal/metal_program.cpp | 1 - taichi/backends/metal/metal_program.h | 1 - taichi/backends/opengl/opengl_program.cpp | 1 - taichi/backends/opengl/opengl_program.h | 1 - taichi/backends/vulkan/vulkan_program.cpp | 1 - taichi/backends/vulkan/vulkan_program.h | 1 - taichi/llvm/llvm_program.cpp | 18 +++++++----------- taichi/llvm/llvm_program.h | 1 - taichi/program/async_engine.cpp | 3 +-- taichi/program/async_engine.h | 1 - taichi/program/program.cpp | 4 ++-- taichi/program/program.h | 3 --- taichi/program/program_impl.h | 1 - taichi/program/state_flow_graph.cpp | 7 +------ taichi/program/state_flow_graph.h | 3 +-- 17 files changed, 12 insertions(+), 37 deletions(-) diff --git a/taichi/backends/cc/cc_program.cpp b/taichi/backends/cc/cc_program.cpp index 9c704c19c9cb7..6d31daddcbbe4 100644 --- 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::unordered_map &, 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 index e6400ad6e3c06..2a6248df5b7b3 100644 --- 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::unordered_map &, uint64 *result_buffer) override; void synchronize() override { diff --git a/taichi/backends/metal/metal_program.cpp b/taichi/backends/metal/metal_program.cpp index 21623f0c5cfb8..7a9a48ece4613 100644 --- 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::unordered_map &, 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 index 7d9ca9da5beef..af8149b030cac 100644 --- 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> &snode_trees_, - std::unordered_map &snodes, uint64 *result_buffer) override; void synchronize() override { diff --git a/taichi/backends/opengl/opengl_program.cpp b/taichi/backends/opengl/opengl_program.cpp index c29ba1175ced7..ffc6a56b2bfc7 100644 --- 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::unordered_map &, 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 index e6467a5fab9b3..5ea6c3ffadc07 100644 --- 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> &snode_trees_, - std::unordered_map &snodes, uint64 *result_buffer) override; void synchronize() override { diff --git a/taichi/backends/vulkan/vulkan_program.cpp b/taichi/backends/vulkan/vulkan_program.cpp index 19fc16f7b673f..a04d516be0067 100644 --- 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::unordered_map &, 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 index 79052380912d5..0c12c26c93677 100644 --- 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::unordered_map &, uint64 *result_buffer) override; void synchronize() override { diff --git a/taichi/llvm/llvm_program.cpp b/taichi/llvm/llvm_program.cpp index 173d5b194fd39..b39c8fcc62e1d 100644 --- 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> &snode_trees_, - std::unordered_map &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 scomp = std::make_unique( - 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 scomp = + std::make_unique(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 index 3b7aefc15218e..3c54da549bf46 100644 --- 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> &snode_trees_, - std::unordered_map &snodes, uint64 *result_buffer) override; template diff --git a/taichi/program/async_engine.cpp b/taichi/program/async_engine.cpp index 730451f719dd8..49c6bdba18b6a 100644 --- 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 &snodes, const BackendExecCompilationFunc &compile_to_backend) : queue(&ir_bank_, compile_to_backend), config_(config), - sfg(std::make_unique(this, &ir_bank_, config, snodes)) { + sfg(std::make_unique(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 index a3f7b08cc26a3..d6578a9adebeb 100644 --- a/taichi/program/async_engine.h +++ b/taichi/program/async_engine.h @@ -138,7 +138,6 @@ class AsyncEngine { std::unique_ptr sfg; explicit AsyncEngine(const CompileConfig *const config, - const std::unordered_map &snodes, const BackendExecCompilationFunc &compile_to_backend); void clear_cache() { diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp index 3cc4dde4cfe63..88de6a5d15647 100644 --- 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( - &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 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 index 620e7da40712b..ead7c6d65d7b9 100644 --- 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 - snodes; // TODO: seems LLVM specific but used by state_flow_graph.cpp. - std::unique_ptr async_engine{nullptr}; std::vector> kernels; diff --git a/taichi/program/program_impl.h b/taichi/program/program_impl.h index ccefab029790a..9f033d0867061 100644 --- 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> &snode_trees_, - std::unordered_map &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 index 5b02f9c52a453..9eb4d443b32c9 100644 --- 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 &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::get_pending_tasks() const { diff --git a/taichi/program/state_flow_graph.h b/taichi/program/state_flow_graph.h index c14785ef964eb..df89a8ec5579e 100644 --- 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 &snodes); + const CompileConfig *const config); std::vector get_pending_tasks() const;