From 18b07472603388209c93cfe564f6f36b6c20b87b Mon Sep 17 00:00:00 2001
From: lin-hitonami <linjiang@taichi.graphics>
Date: Tue, 20 Sep 2022 15:03:21 +0800
Subject: [PATCH] add compile config

---
 taichi/program/compile_config.h                  | 1 +
 taichi/python/export_lang.cpp                    | 1 +
 taichi/transforms/compile_to_offloads.cpp        | 7 +++++--
 taichi/transforms/loop_invariant_code_motion.cpp | 6 +++---
 4 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/taichi/program/compile_config.h b/taichi/program/compile_config.h
index 7fd671f09c8cc9..a86a3f3f2bc998 100644
--- 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
index c3d58b362e52ed..f43b711b3269d2 100644
--- 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
index 9c73c6d4f03683..068aaf29b33f9d 100644
--- 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
index 0f8ec5fa7b8ab4..3a50421f19d996 100644
--- 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());