[CINN]apply new backend ir on some ir utils (#70952)

Hongqing-work · web-flow · commit 216c7f7be361 · 2025-02-10T20:45:14.000+08:00
* [CINN]apply new backend ir on some ir utils

* del ir_collector for later update
diff --git a/paddle/cinn/ir/utils/ir_copy.cc b/paddle/cinn/ir/utils/ir_copy.cc
@@ -25,12 +25,17 @@
 #include "paddle/cinn/ir/ir_printer.h"
 #include "paddle/cinn/ir/module.h"
 #include "paddle/cinn/ir/schedule/ir_schedule.h"
+#include "paddle/cinn/ir/stmt_visitors.h"
+
+using cinn::ir::stmt::BlockRef;
+using cinn::ir::stmt::StmtRef;
 
 namespace cinn {
 namespace ir {
 namespace ir_utils {
 namespace {
-struct IRCopyVisitor : public ir::IRVisitorRequireReImpl<Expr> {
+struct IRCopyVisitor : public ir::IRVisitorRequireReImpl<Expr>,
+                       public stmt::StmtVisitor<StmtRef, BlockRef> {
  public:
   explicit IRCopyVisitor(bool copy_buffer_node)
       : copy_buffer_node(copy_buffer_node) {}
@@ -128,6 +133,11 @@ struct IRCopyVisitor : public ir::IRVisitorRequireReImpl<Expr> {
     return LoweredFunc(func);
   }
 
+  StmtRef VisitStmt(const StmtRef& stmt) {
+    return StmtVisitor::VisitStmt(stmt);
+  }
+  BlockRef VisitBlock(const BlockRef& block) override;
+
  protected:
   // The methods of ir nodes follows the order defined in node.h
   Expr Visit(const ir::IntImm* op) override {
@@ -521,6 +531,10 @@ struct IRCopyVisitor : public ir::IRVisitorRequireReImpl<Expr> {
   NODETY_UNARY_OP_FOR_EACH(OP_UNARY_HANDLE)
 #undef OP_UNARY_HANDLE
 
+#define __(stmt__) StmtRef VisitStmt(const stmt::stmt__& stmt) override;
+  NODETY_FORALL_STMT(__)
+#undef __
+
   std::vector<Expr> Visit(const std::vector<Expr>& vs) {
     std::vector<Expr> copied;
     for (auto& e : vs) {
@@ -557,6 +571,75 @@ Expr IRCopyVisitor::Visit(const ir::intrinsics::BuiltinIntrin* op) {
   return intrinsics::BuiltinIntrin::Make(
       op->name, op->args, op->id, op->arg_nums, op->type());
 }
+
+// copy for stmt
+StmtRef IRCopyVisitor::VisitStmt(const stmt::Let& stmt) {
+  return stmt::Let(Visit(&stmt->symbol()), Visit(&stmt->body()));
+}
+StmtRef IRCopyVisitor::VisitStmt(const stmt::Store& stmt) {
+  return stmt::Store(
+      Visit(&stmt->tensor()), Visit(&stmt->value()), Visit(stmt->indices()));
+}
+StmtRef IRCopyVisitor::VisitStmt(const stmt::Alloc& stmt) {
+  Expr condition;
+  Expr body;
+  if (stmt->condition().defined()) condition = Visit(&stmt->condition());
+  if (stmt->body().defined()) body = Visit(&stmt->body());
+  return stmt::Alloc(Visit(&stmt->destination()),
+                     stmt->type(),
+                     Visit(stmt->extents()),
+                     condition,
+                     body);
+}
+StmtRef IRCopyVisitor::VisitStmt(const stmt::Free& stmt) {
+  return stmt::Free(Visit(&stmt->destination()));
+}
+StmtRef IRCopyVisitor::VisitStmt(const stmt::IfThenElse& stmt) {
+  return stmt::IfThenElse(Visit(&stmt->condition()),
+                          VisitBlock(stmt->true_case()),
+                          VisitBlock(stmt->false_case()));
+}
+StmtRef IRCopyVisitor::VisitStmt(const stmt::For& stmt) {
+  return stmt::For(stmt->loop_var(),
+                   Visit(&stmt->min()),
+                   Visit(&stmt->extent()),
+                   stmt->for_type(),
+                   stmt->device_api(),
+                   VisitBlock(stmt->body()),
+                   stmt->vectorize_info(),
+                   stmt->bind_info());
+}
+StmtRef IRCopyVisitor::VisitStmt(const stmt::Evaluate& stmt) {
+  return stmt::Evaluate(Visit(&stmt->value()));
+}
+StmtRef IRCopyVisitor::VisitStmt(const stmt::Schedule& stmt) {
+  std::vector<Var> iter_vars;
+  for (auto iter_var : stmt->iter_vars()) {
+    auto* var = iter_var.As<_Var_>();
+    PADDLE_ENFORCE_NE(var,
+                      nullptr,
+                      ::common::errors::InvalidArgument(
+                          "Schedule iter_var is not a valid _Var_ type."));
+    iter_vars.emplace_back(Visit(var));
+  }
+  return stmt::Schedule(iter_vars,
+                        Visit(stmt->iter_values()),
+                        Visit(stmt->read_buffers()),
+                        Visit(stmt->write_buffers()),
+                        stmt->name(),
+                        VisitBlock(stmt->body()),
+                        stmt->attrs(),
+                        stmt->reduce_method());
+}
+// copy for block
+BlockRef IRCopyVisitor::VisitBlock(const stmt::BlockRef& block) {
+  std::vector<StmtRef> new_stmts;
+  for (const auto& stmt : block->stmts()) {
+    new_stmts.emplace_back(VisitStmt(stmt));
+  }
+  return stmt::BlockRef(new_stmts);
+}
+
 }  // namespace
 Expr IRCopy(const Expr& x, bool copy_buffer_node) {
   IRCopyVisitor visitor(copy_buffer_node);
@@ -572,6 +655,11 @@ std::vector<Expr> IRCopy(const std::vector<Expr>& x, bool copy_buffer_node) {
   return res;
 }
 
+BlockRef IRCopy(const BlockRef& x, bool copy_buffer_node) {
+  IRCopyVisitor visitor(copy_buffer_node);
+  return visitor.VisitBlock(x);
+}
+
 ir::ModuleExpr IRCopy(const ir::ModuleExpr& x, bool copy_buffer_node) {
   return ir::ModuleExpr(IRCopy(x.GetExprs(), copy_buffer_node));
 }
diff --git a/paddle/cinn/ir/utils/ir_copy.h b/paddle/cinn/ir/utils/ir_copy.h
@@ -19,6 +19,7 @@
 
 #include "paddle/cinn/ir/ir.h"
 #include "paddle/cinn/ir/lowered_func.h"
+#include "paddle/cinn/ir/stmt.h"
 
 namespace cinn {
 
@@ -33,6 +34,9 @@ Expr IRCopy(const Expr& x, bool copy_buffer_node = true);
 std::vector<Expr> IRCopy(const std::vector<Expr>& x,
                          bool copy_buffer_node = true);
 
+std::vector<stmt::StmtRef> IRCopy(const std::vector<stmt::StmtRef>& x,
+                                  bool copy_buffer_node = true);
+stmt::BlockRef IRCopy(const stmt::BlockRef& x, bool copy_buffer_node = true);
 ir::ModuleExpr IRCopy(const ir::ModuleExpr& x, bool copy_buffer_node = true);
 
 ir::Module IRCopy(const Module& m, bool copy_buffer_node = true);
