[llvm] (Decomp of #5251 5/n) Add the parallel compilation worker to L…

…lvmProgramImpl (#5364) * [llvm] (Decomp of #5251 5/n) Add the parallel compilation worker to LlvmProgramImpl * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update export_lang.cpp * num_compile_threads * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
taichi-dev · Jul 8, 2022 · 5f222b9 · 5f222b9
1 parent 67d99d1
commit 5f222b9
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 3 deletions.
diff --git a/taichi/program/compile_config.cpp b/taichi/program/compile_config.cpp
@@ -58,6 +58,7 @@ CompileConfig::CompileConfig() {
   print_kernel_llvm_ir = false;
   print_kernel_nvptx = false;
   print_kernel_llvm_ir_optimized = false;
+  num_compile_threads = 0;
 
   // CUDA backend options:
   device_memory_GB = 1;  // by default, preallocate 1 GB GPU memory

diff --git a/taichi/program/compile_config.h b/taichi/program/compile_config.h
@@ -115,6 +115,8 @@ struct CompileConfig {
   int offline_cache_max_size_of_files{1024 * 1024};  // bytes
   double offline_cache_cleaning_factor{0.25};        // [0.f, 1.f]
 
+  int num_compile_threads{0};
+
   CompileConfig();
 };
 

diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp
@@ -244,7 +244,9 @@ void export_lang(py::module &m) {
       .def_readwrite("offline_cache_max_size_of_files",
                      &CompileConfig::offline_cache_max_size_of_files)
       .def_readwrite("offline_cache_cleaning_factor",
-                     &CompileConfig::offline_cache_cleaning_factor);
+                     &CompileConfig::offline_cache_cleaning_factor)
+      .def_readwrite("num_compile_threads",
+                     &CompileConfig::num_compile_threads);
 
   m.def("reset_default_compile_config",
         [&]() { default_compile_config = CompileConfig(); });

diff --git a/taichi/runtime/program_impls/llvm/llvm_program.cpp b/taichi/runtime/program_impls/llvm/llvm_program.cpp
@@ -19,7 +19,8 @@ namespace lang {
 
 LlvmProgramImpl::LlvmProgramImpl(CompileConfig &config_,
                                  KernelProfilerBase *profiler)
-    : ProgramImpl(config_) {
+    : ProgramImpl(config_),
+      compilation_workers("compile", config_.num_compile_threads) {
   runtime_exec_ = std::make_unique<LlvmRuntimeExecutor>(config_, profiler);
   cache_data_ = std::make_unique<LlvmOfflineCache>();
 }

diff --git a/taichi/runtime/program_impls/llvm/llvm_program.h b/taichi/runtime/program_impls/llvm/llvm_program.h
@@ -8,7 +8,7 @@
 #include "taichi/runtime/llvm/llvm_runtime_executor.h"
 #include "taichi/system/memory_pool.h"
 #include "taichi/program/program_impl.h"
-
+#include "taichi/program/parallel_executor.h"
 #define TI_RUNTIME_HOST
 #include "taichi/program/context.h"
 #undef TI_RUNTIME_HOST
@@ -268,6 +268,7 @@ class LlvmProgramImpl : public ProgramImpl {
     // 2. Destructs runtime_exec_
     runtime_exec_.reset();
   }
+  ParallelExecutor compilation_workers;  // parallel compilation
 
  private:
   std::size_t num_snode_trees_processed_{0};