[opt] Avoid storing constants across offloaded tasks

taichi-dev · Apr 17, 2020 · 7d3bd6c · 7d3bd6c
1 parent b79d93d
commit 7d3bd6c
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/taichi/ir/ir.cpp b/taichi/ir/ir.cpp
@@ -815,6 +815,10 @@ void SNodeOpExpression::flatten(VecStatement &ret) {
   stmt = ret.back().get();
 }
 
+std::unique_ptr<ConstStmt> ConstStmt::copy() {
+  return std::make_unique<ConstStmt>(val);
+}
+
 For::For(const Expr &s, const Expr &e, const std::function<void(Expr)> &func) {
   auto i = Expr(std::make_shared<IdExpression>());
   auto stmt_unique = std::make_unique<FrontendForStmt>(i, s, e);

diff --git a/taichi/ir/ir.h b/taichi/ir/ir.h
@@ -1622,6 +1622,8 @@ class ConstStmt : public Stmt {
     return false;
   }
 
+  std::unique_ptr<ConstStmt> copy();
+
   TI_STMT_DEF_FIELDS(ret_type, val);
   DEFINE_ACCEPT
 };

diff --git a/taichi/transforms/offload.cpp b/taichi/transforms/offload.cpp
@@ -234,6 +234,10 @@ class IdentifyValuesUsedInOtherOffloads : public BasicStmtVisitor {
       return;
     if (stmt_to_offloaded[stmt] == current_offloaded)
       return;
+    if (stmt->is<ConstStmt>()) {
+      // Directly insert copies of ConstStmts later
+      return;
+    }
     if (local_to_global.find(stmt) == local_to_global.end()) {
       // Not yet allocated
       local_to_global[stmt] = allocate_global(stmt->ret_type);
@@ -435,11 +439,19 @@ class FixCrossOffloadReferences : public BasicStmtVisitor {
       auto op = stmt->operand(i);
       if (op == nullptr)
         continue;
-      if (local_to_global_offset.find(op) == local_to_global_offset.end())
-        continue;
       if (stmt_to_offloaded[stmt] ==
           stmt_to_offloaded[op])  // same OffloadedStmt
         continue;
+      if (op->is<ConstStmt>()) {
+        auto copy = op->as<ConstStmt>()->copy();
+        stmt_to_offloaded[copy.get()] = stmt_to_offloaded[stmt];
+        stmt->set_operand(i, copy.get());
+        stmt->insert_before_me(std::move(copy));
+        modified = true;
+        continue;
+      }
+      if (local_to_global_offset.find(op) == local_to_global_offset.end())
+        continue;
 
       auto global = Stmt::make<GlobalTemporaryStmt>(local_to_global_offset[op],
                                                     op->ret_type);