[Opt] Make irpass::replace_all_usages_with bottom-up

taichi/ir/ir.cpp

@@ -204,8 +204,7 @@ Stmt *Stmt::insert_after_me(std::unique_ptr<Stmt> &&new_stmt) {
 }
 
 void Stmt::replace_with(Stmt *new_stmt) {
-  auto root = get_ir_root();
-  irpass::replace_all_usages_with(root, this, new_stmt);
+  irpass::replace_all_usages_with(nullptr, this, new_stmt);
 }
 
 void Stmt::replace_with(VecStatement &&new_statements, bool replace_usages) {

taichi/transforms/ir_printer.cpp

@@ -63,6 +63,13 @@ class IRPrinter : public IRVisitor {
   }
 
   static void run(IRNode *node, std::string *output) {
+    if (node == nullptr) {
+      TI_WARN("IRPrinter: Printing nullptr.");
+      if (output) {
+        *output = std::string();
+      }
+      return;
+    }
     auto p = IRPrinter(output);
     p.print("kernel {{");
     node->accept(&p);

taichi/transforms/offload.cpp

@@ -50,6 +50,7 @@ class Offloader {
 
     for (int i = 0; i < (int)root_statements.size(); i++) {
       auto &stmt = root_statements[i];
+      // Note that stmt->parent is root_block, which doesn't contain stmt now.
       if (auto s = stmt->cast<RangeForStmt>(); s && !s->strictly_serialized) {
         assemble_serial_statements();
         auto offloaded =

taichi/transforms/statement_usage_replace.cpp

@@ -7,6 +7,8 @@ TLANG_NAMESPACE_BEGIN
 // Replace all usages statement A with a new statement B.
 // Note that the original statement A is NOT replaced.
 class StatementUsageReplace : public IRVisitor {
+  // The reason why we don't use BasicStmtVisitor is we don't want to go into
+  // FrontendForStmt.
  public:
   Stmt *old_stmt, *new_stmt;
 
@@ -16,17 +18,16 @@ class StatementUsageReplace : public IRVisitor {
     invoke_default_visitor = true;
   }
 
-  void visit(Stmt *stmt) {
-    // if (stmt->may_have_operand(old_stmt))
+  void visit(Stmt *stmt) override {
     stmt->replace_operand_with(old_stmt, new_stmt);
   }
 
-  void visit(WhileStmt *stmt) {
+  void visit(WhileStmt *stmt) override {
     stmt->replace_operand_with(old_stmt, new_stmt);
     stmt->body->accept(this);
   }
 
-  void visit(IfStmt *if_stmt) {
+  void visit(IfStmt *if_stmt) override {
     if_stmt->replace_operand_with(old_stmt, new_stmt);
     if (if_stmt->true_statements)
       if_stmt->true_statements->accept(this);
@@ -35,34 +36,52 @@ class StatementUsageReplace : public IRVisitor {
     }
   }
 
-  void visit(Block *stmt_list) {
+  void visit(Block *stmt_list) override {
     for (auto &stmt : stmt_list->statements) {
       stmt->accept(this);
     }
   }
 
-  void visit(RangeForStmt *stmt) {
+  void visit(RangeForStmt *stmt) override {
     stmt->replace_operand_with(old_stmt, new_stmt);
     stmt->body->accept(this);
   }
 
-  void visit(StructForStmt *stmt) {
+  void visit(StructForStmt *stmt) override {
     stmt->body->accept(this);
   }
 
-  void visit(OffloadedStmt *stmt) {
+  void visit(OffloadedStmt *stmt) override {
     stmt->all_blocks_accept(this);
   }
 
-  static void run(IRNode *node, Stmt *old_stmt, Stmt *new_stmt) {
-    StatementUsageReplace inst(old_stmt, new_stmt);
-    node->accept(&inst);
+  static void run(IRNode *root, Stmt *old_stmt, Stmt *new_stmt) {
+    StatementUsageReplace replacer(old_stmt, new_stmt);
+    if (root != nullptr) {
+      // If root is specified, simply traverse the root.
+      root->accept(&replacer);
+      return;
+    }
+
+    // statements inside old_stmt->parent
+    TI_ASSERT(old_stmt->parent != nullptr);
+    old_stmt->parent->accept(&replacer);
+    auto current_block = old_stmt->parent->parent;
+
+    // statements outside old_stmt->parent: bottom-up
+    while (current_block != nullptr) {
+      for (auto &stmt : current_block->statements) {
+        stmt->replace_operand_with(old_stmt, new_stmt);
+      }
+      current_block = current_block->parent;
+    }
   }
 };
 
 namespace irpass {
 
 void replace_all_usages_with(IRNode *root, Stmt *old_stmt, Stmt *new_stmt) {
+  TI_AUTO_PROF;
   StatementUsageReplace::run(root, old_stmt, new_stmt);
 }