stash

taichi-dev · Sep 19, 2022 · 7773702 · 7773702
1 parent 64450a6
commit 7773702
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/taichi/transforms/cache_loop_invariant_global_vars.cpp b/taichi/transforms/cache_loop_invariant_global_vars.cpp
@@ -174,7 +174,6 @@ class CacheLoopInvariantGlobalVars : public LoopInvariantDetector {
 
     while (true) {
       CacheLoopInvariantGlobalVars eliminator(config);
-      irpass::print(node);
       node->accept(&eliminator);
       if (eliminator.modifier.modify_ir())
         modified = true;

diff --git a/taichi/transforms/loop_invariant_code_motion.cpp b/taichi/transforms/loop_invariant_code_motion.cpp
@@ -52,6 +52,26 @@ class LoopInvariantCodeMotion : public LoopInvariantDetector {
     }
   }
 
+  void visit(ExternalPtrStmt *stmt) override {
+    if (is_loop_invariant(stmt, stmt->parent)) {
+      auto replacement = stmt->clone();
+      stmt->replace_usages_with(replacement.get());
+
+      modifier.insert_before(current_loop_stmt(), std::move(replacement));
+      modifier.erase(stmt);
+    }
+  }
+
+  void visit(ArgLoadStmt *stmt) override {
+    if (is_loop_invariant(stmt, stmt->parent)) {
+      auto replacement = stmt->clone();
+      stmt->replace_usages_with(replacement.get());
+
+      modifier.insert_before(current_loop_stmt(), std::move(replacement));
+      modifier.erase(stmt);
+    }
+  }
+
   static bool run(IRNode *node, const CompileConfig &config) {
     bool modified = false;