add compile config

taichi-dev · Sep 20, 2022 · 18b0747 · 18b0747
1 parent 347a373
commit 18b0747
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 5 deletions.
diff --git a/taichi/program/compile_config.h b/taichi/program/compile_config.h
@@ -28,6 +28,7 @@ struct CompileConfig {
   bool lower_access;
   bool simplify_after_lower_access;
   bool move_loop_invariant_outside_if;
+  bool cache_loop_invariant_global_vars{false};
   bool demote_dense_struct_fors;
   bool advanced_optimization;
   bool constant_folding;

diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp
@@ -166,6 +166,7 @@ void export_lang(py::module &m) {
       .def_readwrite("lower_access", &CompileConfig::lower_access)
       .def_readwrite("move_loop_invariant_outside_if",
                      &CompileConfig::move_loop_invariant_outside_if)
+      .def_readwrite("cache_loop_invariant_global_vars", &CompileConfig::cache_loop_invariant_global_vars)
       .def_readwrite("default_cpu_block_dim",
                      &CompileConfig::default_cpu_block_dim)
       .def_readwrite("cpu_block_dim_adaptive",

diff --git a/taichi/transforms/compile_to_offloads.cpp b/taichi/transforms/compile_to_offloads.cpp
@@ -183,8 +183,11 @@ void offload_to_executable(IRNode *ir,
   irpass::demote_atomics(ir, config);
   print("Atomics demoted I");
   irpass::analysis::verify(ir);
-  irpass::cache_loop_invariant_global_vars(ir, config);
-  print("Cache loop-invariant global vars");
+  if (config.cache_loop_invariant_global_vars) {
+    irpass::cache_loop_invariant_global_vars(ir, config);
+    print("Cache loop-invariant global vars");
+  }
+
   if (config.demote_dense_struct_fors) {
     irpass::demote_dense_struct_fors(ir, config.packed);
     irpass::type_check(ir, config);

diff --git a/taichi/transforms/loop_invariant_code_motion.cpp b/taichi/transforms/loop_invariant_code_motion.cpp
@@ -33,7 +33,7 @@ class LoopInvariantCodeMotion : public LoopInvariantDetector {
   }
 
   void visit(GlobalPtrStmt *stmt) override {
-    if (is_loop_invariant(stmt, stmt->parent)) {
+    if (config.cache_loop_invariant_global_vars && is_loop_invariant(stmt, stmt->parent)) {
       auto replacement = stmt->clone();
       stmt->replace_usages_with(replacement.get());
 
@@ -43,7 +43,7 @@ class LoopInvariantCodeMotion : public LoopInvariantDetector {
   }
 
   void visit(ExternalPtrStmt *stmt) override {
-    if (is_loop_invariant(stmt, stmt->parent)) {
+    if (config.cache_loop_invariant_global_vars && is_loop_invariant(stmt, stmt->parent)) {
       auto replacement = stmt->clone();
       stmt->replace_usages_with(replacement.get());
 
@@ -53,7 +53,7 @@ class LoopInvariantCodeMotion : public LoopInvariantDetector {
   }
 
   void visit(ArgLoadStmt *stmt) override {
-    if (is_loop_invariant(stmt, stmt->parent)) {
+    if (config.cache_loop_invariant_global_vars && is_loop_invariant(stmt, stmt->parent)) {
       auto replacement = stmt->clone();
       stmt->replace_usages_with(replacement.get());