PaddlePaddle
diff --git a/‎paddle/cinn/common/object.h
Copy file name to clipboard
+37Lines changed: 37 additions & 0 deletions b/‎paddle/cinn/common/object.h
Copy file name to clipboard
+37Lines changed: 37 additions & 0 deletions
diff --git a/‎paddle/cinn/ir/CMakeLists.txt
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/cinn/ir/CMakeLists.txt
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/cinn/ir/ir_base.cc
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions b/‎paddle/cinn/ir/ir_base.cc
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions
diff --git a/‎paddle/cinn/ir/ir_base.h
Copy file name to clipboard
+19Lines changed: 19 additions & 0 deletions b/‎paddle/cinn/ir/ir_base.h
Copy file name to clipboard
+19Lines changed: 19 additions & 0 deletions
diff --git a/‎paddle/cinn/ir/ir_printer.cc
Copy file name to clipboard
+242-1Lines changed: 242 additions & 1 deletion b/‎paddle/cinn/ir/ir_printer.cc
Copy file name to clipboard
+242-1Lines changed: 242 additions & 1 deletion
@@ -75,5 +75,42 @@ struct Object {
 using object_ptr = Object*;
 using shared_object = Shared<Object>;
 
+/*
+ * \brief Delete the default copy/move constructor and assign operator
+ * \param TypeName The class typename.
+ */
+#define CINN_DELETE_COPY_MOVE_AND_ASSIGN_FOR_OBJECT_NODE(TypeName) \
+  TypeName(const TypeName& other) = delete;                        \
+  TypeName(TypeName&& other) = delete;                             \
+  TypeName& operator=(const TypeName& other) = delete;             \
+  TypeName& operator=(TypeName&& other) = delete;
+
+/*
+ * \brief Define the default copy/move constructor and assign operator
+ * \param TypeName The class typename.
+ */
+#define CINN_DEFINE_COPY_MOVE_AND_ASSIGN_FOR_OBJECT_REF(TypeName) \
+  TypeName(const TypeName& other) = default;                      \
+  TypeName(TypeName&& other) = default;                           \
+  TypeName& operator=(const TypeName& other) = default;           \
+  TypeName& operator=(TypeName&& other) = default;
+
+/*
+ * \brief Define object reference methods
+ * \param TypeName The object type name
+ * \param ParentType The parent type of the objectref
+ * \param ObjectName The type name of the object
+ */
+#define CINN_DEFINE_OBJECT_REF_METHODS(TypeName, ParentType, ObjectName) \
+  TypeName() = default;                                                  \
+  explicit TypeName(ObjectName* n) : ParentType(n) {}                    \
+  explicit TypeName(const cinn::common::Shared<Object>& ref)             \
+      : ParentType(ref) {}                                               \
+  CINN_DEFINE_COPY_MOVE_AND_ASSIGN_FOR_OBJECT_REF(TypeName);             \
+  const ObjectName* operator->() const {                                 \
+    return static_cast<const ObjectName*>(this->p_);                     \
+  }                                                                      \
+  ObjectName* operator->() { return static_cast<ObjectName*>(this->p_); }
+
 }  // namespace common
 }  // namespace cinn
@@ -19,6 +19,8 @@ gather_srcs(
   intrinsic_ops.cc
   layout.cc
   schedule_block_graph.cc
+  stmt.cc
+  stmt_visitors.cc
   dim.cc)
 
 add_subdirectory(ir_analyzer)
 
@@ -66,6 +66,21 @@ std::ostream &operator<<(std::ostream &os, IrNodeTy type) {
   return os;
 }
 
+std::ostream &operator<<(std::ostream &os, StmtNodeTy type) {
+  switch (type) {
+#define __m(t__)                         \
+  case StmtNodeTy::t__:                  \
+    os << "<stmt node: " << #t__ << ">"; \
+    break;
+
+    NODETY_FORALL_STMT(__m)
+#undef __m
+    default:
+      PADDLE_THROW(
+          ::common::errors::InvalidArgument("unknown StmtNodeTy found"));
+  }
+}
+
 Expr Zero(const Type &type) {
   if (type.is_bfloat16()) return Expr(bfloat16(0.f));
   if (type.is_float16()) return Expr(float16(0.f));
 
@@ -116,6 +116,7 @@ class Dim;
 #define NODETY_CONTROL_OP_FOR_INTRINSIC(macro__) \
   macro__(IntrinsicOp)                      \
 
+// TODO(Hongqing-work): change NODETY_FORALL to NODETY_FORALL_EXPR
 #define NODETY_FORALL(__m)              \
   NODETY_PRIMITIVE_TYPE_FOR_EACH(__m)   \
   NODETY_OP_FOR_EACH(__m)               \
@@ -126,6 +127,16 @@ class Dim;
   NODETY_PRIMITIVE_TYPE_FOR_EACH(__m)                    \
   NODETY_OP_FOR_EACH(__m)                                \
   NODETY_CONTROL_OP_FOR_EACH(__m)
+
+#define NODETY_FORALL_STMT(macro__) \
+  macro__(Let)                      \
+  macro__(Store)                    \
+  macro__(Alloc)                    \
+  macro__(Free)                     \
+  macro__(IfThenElse)               \
+  macro__(For)                      \
+  macro__(Schedule)                 \
+  macro__(Evaluate)
 // clang-format on
 
 //! Define IrNodeTy
@@ -143,6 +154,13 @@ enum class IrNodeTy {
 #undef __m
 // @}
 
+//! Define StmtNodeTy
+// @{
+#define __m(x__) x__,
+enum class StmtNodeTy { kUnk = -1, NODETY_FORALL_STMT(__m) };
+#undef __m
+// @}
+
 //! String representations for IrNodeTy.
 // @{
 #define __m(x__) #x__,
@@ -152,6 +170,7 @@ const std::vector<std::string> kIrNodeTyReprs(
 // @}
 
 std::ostream& operator<<(std::ostream& os, IrNodeTy type);
+std::ostream& operator<<(std::ostream& os, StmtNodeTy type);
 
 struct Expr;
 
 
@@ -33,7 +33,17 @@ using cinn::common::bfloat16;
 using cinn::common::float16;
 
 void IrPrinter::Print(const Expr &e) {
-  IRVisitorRequireReImpl::Visit(&e);
+  Visit(e);
+  os_ << str_;
+  str_ = "";
+}
+void IrPrinter::Print(const stmt::StmtRef &stmt) {
+  VisitStmt(stmt);
+  os_ << str_;
+  str_ = "";
+}
+void IrPrinter::Print(const stmt::BlockRef &block) {
+  VisitBlock(block);
   os_ << str_;
   str_ = "";
 }
@@ -803,6 +813,219 @@ void IrPrinter::Visit(const intrinsics::BuiltinIntrin *x) {
   str_ += ")";
 }
 
+void IrPrinter::VisitStmt(const stmt::Let &stmt) {
+  PADDLE_ENFORCE_EQ(
+      stmt->type().valid(),
+      true,
+      ::common::errors::InvalidArgument(
+          "The type of `f` is not valid. "
+          "Please ensure that `f->type()` returns a valid type."));
+
+  str_ += stmt->type().to_string();
+  str_ += " ";
+  Visit(stmt->symbol());
+  if (stmt->body().defined()) {
+    str_ += " = ";
+    Visit(stmt->body());
+  }
+}
+
+void IrPrinter::VisitStmt(const stmt::Store &stmt) {
+  if (stmt->is_addr_tensor()) {
+    auto *tensor_node = stmt->tensor().As<ir::_Tensor_>();
+    PADDLE_ENFORCE_NOT_NULL(tensor_node,
+                            ::common::errors::InvalidArgument(
+                                "The tensor node is not valid. "
+                                "Please ensure that `stmt->tensor` is "
+                                "properly assigned to a tensor node."));
+
+    str_ += tensor_node->name;
+  } else if (stmt->is_addr_scalar()) {
+    Visit(stmt->tensor());
+  } else {
+    CINN_NOT_IMPLEMENTED
+  }
+
+  str_ += "[";
+  for (std::size_t i = 0; i + 1 < stmt->indices().size(); i++) {
+    Visit(stmt->indices()[i]);
+    str_ += ", ";
+  }
+  if (!stmt->indices().empty()) Visit(stmt->indices().back());
+  str_ += "] = ";
+  Visit(stmt->value());
+}
+
+void IrPrinter::VisitStmt(const stmt::Alloc &stmt) {
+  auto *buffer = stmt->destination().As<ir::_Buffer_>();
+  PADDLE_ENFORCE_NOT_NULL(buffer,
+                          ::common::errors::InvalidArgument(
+                              "The destination is not a valid buffer. "
+                              "Please ensure that `stmt->destination` is "
+                              "properly assigned to a buffer."));
+  str_ += "alloc(";
+  str_ += buffer->name;
+  str_ += ", ";
+  Visit(stmt->extents());
+  str_ += ")";
+}
+
+void IrPrinter::VisitStmt(const stmt::Free &stmt) {
+  auto *buffer = stmt->destination().As<ir::_Buffer_>();
+  PADDLE_ENFORCE_NOT_NULL(buffer,
+                          ::common::errors::InvalidArgument(
+                              "The destination is not a valid buffer. "
+                              "Please ensure that `stmt->destination` is "
+                              "properly assigned to a buffer."));
+
+  str_ += "free(";
+  str_ += buffer->name;
+  str_ += ")";
+}
+
+void IrPrinter::VisitStmt(const stmt::IfThenElse &stmt) {
+  str_ += "if (";
+  Visit(stmt->condition());
+  str_ += ") ";
+  VisitBlock(stmt->true_case());
+
+  if (!stmt->false_case()->stmts().empty()) {
+    str_ += " else ";
+    VisitBlock(stmt->false_case());
+  }
+}
+
+void IrPrinter::VisitStmt(const stmt::For &stmt) {
+  if (stmt->is_parallel()) {
+    str_ += "parallel for (";
+  } else if (stmt->is_unrolled()) {
+    str_ += "unroll for (";
+  } else if (stmt->is_vectorized()) {
+    int factor = stmt->vectorize_info().factor;
+    str_ += "vectorize[";
+    str_ += std::to_string(factor);
+    str_ += "] for (";
+  } else if (stmt->is_binded()) {
+    auto &bind_info = stmt->bind_info();
+    if (bind_info.valid()) {
+      char axis_name = 'x' + bind_info.offset;
+      auto for_type = bind_info.for_type;
+      std::string prefix =
+          for_type == ForType::GPUBlock ? "blockIdx." : "threadIdx.";
+      str_ += "thread_bind[";
+      str_ += prefix;
+      str_ += axis_name;
+      str_ += "] for (";
+    } else {
+      str_ += "thread_bind[invalid info] for (";
+    }
+  } else if (stmt->is_serial()) {
+    str_ += "serial for (";
+  } else if (stmt->is_default()) {
+    str_ += "default for (";
+  } else {
+    str_ += "for (";
+  }
+  Visit(stmt->loop_var());
+  str_ += ", ";
+  Visit(stmt->min());
+  str_ += ", ";
+  Visit(stmt->extent());
+  str_ += ") ";
+  VisitBlock(stmt->body());
+}
+
+void IrPrinter::VisitStmt(const stmt::Schedule &stmt) {
+  str_ += "Schedule (";
+  str_ += stmt->name();
+  str_ += ") {\n";
+  // print block vars and bindings
+  auto iter_vars = stmt->iter_vars();
+  auto iter_values = stmt->iter_values();
+  PADDLE_ENFORCE_EQ(iter_vars.size(),
+                    iter_values.size(),
+                    ::common::errors::InvalidArgument(
+                        "iter_vars.size() != iter_values.size()"));
+  IncIndent();
+  if (!iter_vars.empty()) DoIndent();
+  for (std::size_t i = 0; i < iter_vars.size(); i++) {
+    if (i) str_ += ", ";
+    str_ += iter_vars[i]->name;
+  }
+  if (!iter_vars.empty()) str_ += " = axis.bind(";
+  for (std::size_t i = 0; i < iter_values.size(); i++) {
+    if (i) str_ += ", ";
+    Visit(iter_values[i]);
+  }
+  if (!iter_vars.empty()) str_ += ")\n";
+
+  // print block body
+  if (!stmt->read_buffers().empty()) {
+    DoIndent();
+    str_ += "read_buffers(";
+    auto &read_buffers = stmt->read_buffers();
+    for (std::size_t i = 0; i < read_buffers.size(); i++) {
+      if (i) str_ += ", ";
+      Visit(read_buffers[i]);
+    }
+    str_ += ")\n";
+  }
+  if (!stmt->write_buffers().empty()) {
+    DoIndent();
+    str_ += "write_buffers(";
+    auto &write_buffers = stmt->write_buffers();
+    for (std::size_t i = 0; i < write_buffers.size(); i++) {
+      if (i) str_ += ", ";
+      Visit(write_buffers[i]);
+    }
+    str_ += ")\n";
+  }
+  if (!stmt->attrs().empty()) {
+    DoIndent();
+    str_ += "attrs(";
+    bool comma = false;
+    for (auto &&kv : stmt->attrs()) {
+      if (comma) str_ += ", ";
+      str_ += kv.first;
+      str_ += ":";
+      absl::visit(
+          [this](auto &&arg) {
+            std::ostringstream ss;
+            ss << arg;
+            this->str_ += ss.str();
+          },
+          kv.second);
+      comma = true;
+    }
+    str_ += ")\n";
+  }
+
+  for (const auto &inner_stmt : stmt->body()->stmts()) {
+    DoIndent();
+    VisitStmt(inner_stmt);
+    str_ += "\n";
+  }
+  DecIndent();
+  DoIndent();
+  str_ += "}";
+}
+
+void IrPrinter::VisitStmt(const stmt::Evaluate &stmt) { Visit(stmt->value()); }
+
+void IrPrinter::VisitBlock(const stmt::BlockRef &block) {
+  str_ += "{\n";
+
+  IncIndent();
+  for (const auto &stmt : block->stmts()) {
+    DoIndent();
+    VisitStmt(stmt);
+    str_ += "\n";
+  }
+  DecIndent();
+  DoIndent();
+  str_ += "}";
+}
+
 std::ostream &operator<<(std::ostream &os, Expr a) {
   std::stringstream ss;
   IrPrinter printer(ss);
@@ -844,5 +1067,23 @@ std::ostream &operator<<(std::ostream &os, const std::vector<Dim> &a) {
   return os;
 }
 
+namespace stmt {
+std::ostream &operator<<(std::ostream &os, const stmt::BlockRef &block) {
+  std::stringstream ss;
+  IrPrinter printer(ss);
+  printer.Print(block);
+  os << ss.str();
+  return os;
+}
+
+std::ostream &operator<<(std::ostream &os, const stmt::StmtRef &stmt) {
+  std::stringstream ss;
+  IrPrinter printer(ss);
+  printer.Print(stmt);
+  os << ss.str();
+  return os;
+}
+}  // namespace stmt
+
 }  // namespace ir
 }  // namespace cinn