move collect_undefined_vars from optim/ to ir/utils/ (#57230)

Courtesy-Xs · web-flow · commit 114813a67b18 · 2023-09-14T10:43:38.000+08:00
diff --git a/paddle/cinn/backends/llvm/codegen_x86.cc b/paddle/cinn/backends/llvm/codegen_x86.cc
@@ -28,7 +28,7 @@
 #include "paddle/cinn/common/target.h"
 #include "paddle/cinn/ir/ir.h"
 #include "paddle/cinn/ir/op/ir_operators.h"
-#include "paddle/cinn/optim/collect_undefined_vars.h"
+#include "paddle/cinn/ir/utils/ir_nodes_collector.h"
 #include "paddle/cinn/runtime/intrinsic.h"
 
 namespace cinn::backends {
@@ -98,7 +98,7 @@ void CodeGenX86::CreateParallelLaunch(Expr body, int num_task) {
                                              llvm::Function::PrivateLinkage,
                                              "__parallel_lambda",
                                              m_);
-  std::vector<std::string> vars = optim::CollectUndefinedVars(&body);
+  std::vector<std::string> vars = ir::CollectUndefinedVars(&body);
   uint64_t nbytes;
   auto* data = PackVars(vars, &nbytes);
 
diff --git a/paddle/cinn/ir/utils/ir_nodes_collector.cc b/paddle/cinn/ir/utils/ir_nodes_collector.cc
@@ -207,6 +207,92 @@ std::set<Expr> CollectReferencedTensors(
   return ts0;
 }
 
+std::vector<std::string> CollectUndefinedVars(const Expr* e) {
+  struct Mutator : public ir::IRMutator<const Expr*> {
+    using ir::IRMutator<const Expr*>::Visit;
+    std::vector<std::string> undefined_vars;
+    std::set<std::string> defined_vars;
+    std::set<std::string> used_vars;
+
+    void CollectVarDef(const std::string& var) {
+      CHECK(!defined_vars.count(var))
+          << "var " << var << " has been defined, please check";
+      CHECK(!used_vars.count(var))
+          << "var " << var << " is wrongly used before definition";
+      defined_vars.insert(var);
+    }
+
+    void ClearVar(const std::string& var) {
+      defined_vars.erase(var);
+      used_vars.erase(var);
+    }
+
+    void CollectVarUse(const std::string& var) {
+      used_vars.insert(var);
+      if (defined_vars.count(var) == 0) {
+        undefined_vars.push_back(var);
+      }
+    }
+
+    void Visit(const ir::Let* op, const Expr* expr) override {
+      Expr symbol = op->symbol;
+      auto var = symbol.as_var_ref();
+      CHECK(var.defined());
+      CollectVarDef(var->name);
+      auto* node = expr->As<ir::Let>();
+      Visit(&node->body, &node->body);
+    }
+
+    void Visit(const ir::For* op, const Expr* expr) override {
+      CollectVarDef(op->loop_var->name);
+      auto* node = expr->As<ir::For>();
+      Visit(&node->min, &node->min);
+      Visit(&node->extent, &node->extent);
+      Visit(&node->body, &node->body);
+      ClearVar(op->loop_var->name);
+    }
+
+    void Visit(const ir::Load* op, const Expr* expr) override {
+      auto tensor = op->tensor.as_tensor_ref();
+      CollectVarUse(tensor->name);
+      auto* node = expr->As<ir::Load>();
+      for (auto& idx : node->indices) Visit(&idx, &idx);
+    }
+
+    void Visit(const ir::Store* op, const Expr* expr) override {
+      auto tensor = op->tensor.as_tensor_ref();
+      CollectVarUse(tensor->name);
+      auto* node = expr->As<ir::Store>();
+      for (auto& idx : node->indices) Visit(&idx, &idx);
+      Visit(&node->value, &node->value);
+    }
+
+    void Visit(const ir::_Var_* op, const Expr* expr) override {
+      CollectVarUse(op->name);
+      auto* node = expr->As<ir::_Var_>();
+      if (node->lower_bound.defined()) {
+        Visit(&node->lower_bound, &node->lower_bound);
+      }
+      if (node->upper_bound.defined()) {
+        Visit(&node->upper_bound, &node->upper_bound);
+      }
+    }
+
+    void Visit(const ir::Reduce* op, const Expr* expr) override {
+      for (auto& axis : op->reduce_axis) {
+        CollectVarDef(axis->name);
+      }
+      auto* node = expr->As<ir::Reduce>();
+      if (node->init.defined()) Visit(&node->init, &node->init);
+      Visit(&node->body, &node->body);
+    }
+  };
+
+  Mutator mutator;
+  mutator.Visit(e, e);
+  return mutator.undefined_vars;
+}
+
 std::set<std::string> CollectTensorNeedsWrite(const Expr* e) {
   std::set<std::string> tensor_written;
   IrNodesCollector::handler_t handler = [&](const Expr* x) {
diff --git a/paddle/cinn/ir/utils/ir_nodes_collector.h b/paddle/cinn/ir/utils/ir_nodes_collector.h
@@ -65,6 +65,23 @@ std::map<std::string, Expr> CollectTensorMap(
       return true;
     });
 
+/**
+ * Collect undefined vars in the scope.
+ *
+ * e.g.
+ *
+ * The expression:
+ * for i
+ *  for j
+ *    a[i, j] = b[i, j]
+ *
+ * here a, b are vars without definition
+ */
+std::vector<std::string> CollectUndefinedVars(const Expr* e);
+
+/**
+ * Collect the Tensor Nodes which will be Writed by Store or Call Nodes
+ */
 std::set<std::string> CollectTensorNeedsWrite(const Expr* e);
 
 }  // namespace ir
diff --git a/paddle/cinn/optim/CMakeLists.txt b/paddle/cinn/optim/CMakeLists.txt
@@ -23,7 +23,6 @@ gather_srcs(
   replace_const_param_to_integer.cc
   lower_intrin.cc
   cast_bool_to_int8.cc
-  collect_undefined_vars.cc
   var_mod_simplify.cc
   remove_schedule_block.cc)
 
diff --git a/paddle/cinn/optim/collect_undefined_vars.cc b/paddle/cinn/optim/collect_undefined_vars.cc
diff --git a/paddle/cinn/optim/collect_undefined_vars.h b/paddle/cinn/optim/collect_undefined_vars.h
