[CINN][backend]add parent for stmt and block (#71042)

Hongqing-work · web-flow · commit a965a4fe7f4c · 2025-02-08T18:31:33.000+08:00
* [CINN][backend]add parent for stmt and block

* fix test
diff --git a/paddle/cinn/ir/stmt.cc b/paddle/cinn/ir/stmt.cc
@@ -23,6 +23,15 @@ namespace stmt {
 
 using cinn::common::make_shared;
 
+const BlockRef _Stmt_::GetParentBlockRef() const { return BlockRef{parent_}; }
+
+void _Stmt_::set_block_fields(const std::vector<BlockRef> &blocks) {
+  blocks_ = blocks;
+  for (auto &block : blocks_) {
+    block->set_parent(this);
+  }
+}
+
 BlockRef _Block_::Make(const std::vector<StmtRef> &stmts) {
   BlockRef ref(new _Block_());
   ref->set_stmts(stmts);
@@ -215,8 +224,7 @@ IfThenElse _IfThenElse_::Make(Expr condition,
           "The true_case is not defined. "
           "A valid true_case expression is required for _IfThenElse_."));
   ref->set_condition(condition);
-  ref->set_true_case(true_case);
-  ref->set_false_case(false_case);
+  ref->set_block_fields({true_case, false_case});
   return ref;
 }
 
diff --git a/paddle/cinn/ir/stmt.h b/paddle/cinn/ir/stmt.h
@@ -33,6 +33,8 @@ class For;
 class Evaluate;
 class Schedule;
 
+class _Block_;
+
 /*
  * \brief Define statement type
  * \param Type The statement type
@@ -62,19 +64,21 @@ class _Stmt_ : public cinn::common::Object {
 
   virtual StmtNodeTy stmt_type() const { return StmtNodeTy::kUnk; }
 
+  const BlockRef GetParentBlockRef() const;
   const std::vector<BlockRef>& block_fields() const { return blocks_; }
-  virtual void set_block_fields(const std::vector<BlockRef>& blocks) {
-    blocks_ = blocks;
-  }
+  virtual void set_block_fields(const std::vector<BlockRef>& blocks);
   virtual Type type() const { return type_; }
   void set_type(Type type) { type_ = type; }
 
   virtual void Verify() const { CINN_NOT_IMPLEMENTED }
 
  protected:
+  _Block_* parent_;
   std::vector<BlockRef> blocks_;
 
  private:
+  friend class _Block_;
+  void set_parent(_Block_* parent) { parent_ = parent; }
   // Forbidden cast these methods of Object.
   // TODO(Hongqing-work): restruct the Object class.
   static constexpr char* __type_info__ = "StmtNode";
@@ -139,14 +143,23 @@ class StmtRef : public cinn::common::Shared<Object> {
 
 class _Block_ : public cinn::common::Object {
  public:
+  const StmtRef GetParentStmtRef() const { return StmtRef{parent_}; }
   const std::vector<StmtRef>& stmts() const { return stmts_; }
-  void set_stmts(const std::vector<StmtRef>& stmts) { stmts_ = stmts; }
+  void set_stmts(const std::vector<StmtRef>& stmts) {
+    stmts_ = stmts;
+    for (auto& stmt : stmts_) stmt->set_parent(this);
+  }
 
   static BlockRef Make(const std::vector<StmtRef>& stmts);
 
  private:
   _Block_() = default;
   std::vector<StmtRef> stmts_;
+
+  _Stmt_* parent_;
+  friend class _Stmt_;
+  void set_parent(_Stmt_* parent) { parent_ = parent; }
+
   // Forbidden cast these methods of Object.
   static constexpr char* __type_info__ = "BlockNode";
   const char* type_info() const override { return __type_info__; }
@@ -324,24 +337,22 @@ class _IfThenElse_ : public _Stmt_ {
   const BlockRef& true_case() const { return blocks_[0]; }
   const BlockRef& false_case() const { return blocks_[1]; }
   void set_condition(Expr condition) { condition_ = condition; }
-  void set_true_case(BlockRef true_case) { blocks_[0] = true_case; }
-  void set_false_case(BlockRef false_case) { blocks_[1] = false_case; }
+  void set_true_case(BlockRef true_case) {
+    set_block_fields({true_case, blocks_[1]});
+  }
+  void set_false_case(BlockRef false_case) {
+    set_block_fields({blocks_[0], false_case});
+  }
 
   void Verify() const override;
 
   virtual void set_block_fields(const std::vector<BlockRef>& blocks) {
     PADDLE_ENFORCE_EQ(
-        blocks.size() == 1U || blocks.size() == 2U,
+        blocks.size() == 2U,
         true,
         ::common::errors::InvalidArgument(
-            "blocks size for IfThenElse must be 1 or 2, but got %d",
-            blocks.size()));
-    if (blocks.size() == 2U) {
-      blocks_ = blocks;
-    } else {
-      blocks_[0] = blocks[0];
-      blocks_[1] = BlockRef(std::vector<StmtRef>());
-    }
+            "blocks size for IfThenElse must be 2, but got %d", blocks.size()));
+    _Stmt_::set_block_fields(blocks);
   }
 
   CINN_DELETE_COPY_MOVE_AND_ASSIGN_FOR_OBJECT_NODE(_IfThenElse_)
@@ -380,7 +391,7 @@ class _For_ : public _Stmt_, public ForBase {
   void set_loop_var(Var loop_var) { loop_var_ = loop_var; }
   void set_min(Expr min) { min_ = min; }
   void set_extent(Expr extent) { extent_ = extent; }
-  void set_body(BlockRef body) { blocks_[0] = body; }
+  void set_body(BlockRef body) { set_block_fields({body}); }
   void set_device_api(DeviceAPI device_api) { device_api_ = device_api; }
   void set_metadata(LLVMForLoopMeta metadata) { metadata_ = metadata; }
 
@@ -392,7 +403,7 @@ class _For_ : public _Stmt_, public ForBase {
         true,
         ::common::errors::InvalidArgument(
             "blocks size for For must be 1, but got %d", blocks.size()));
-    blocks_ = blocks;
+    _Stmt_::set_block_fields(blocks);
   }
 
   CINN_DELETE_COPY_MOVE_AND_ASSIGN_FOR_OBJECT_NODE(_For_)
@@ -487,7 +498,7 @@ class _Schedule_ : public _Stmt_ {
   void set_write_buffers(const std::vector<Expr>& write_buffers) {
     write_buffers_ = write_buffers;
   }
-  void set_body(const BlockRef& body) { blocks_[0] = body; }
+  void set_body(const BlockRef& body) { set_block_fields({body}); }
   void set_attrs(const std::map<std::string, attr_t>& attrs) { attrs_ = attrs; }
   void set_name(const std::string& name) { name_ = name; }
   void set_reduce_method(const ReduceMethod& reduce_method) {
@@ -502,7 +513,7 @@ class _Schedule_ : public _Stmt_ {
         true,
         ::common::errors::InvalidArgument(
             "blocks size for For must be 1, but got %d", blocks.size()));
-    blocks_ = blocks;
+    _Stmt_::set_block_fields(blocks);
   }
 
   CINN_DELETE_COPY_MOVE_AND_ASSIGN_FOR_OBJECT_NODE(_Schedule_)
diff --git a/test/cpp/pir/cinn/adt/merge_block_utils_test.cc b/test/cpp/pir/cinn/adt/merge_block_utils_test.cc
@@ -49,12 +49,13 @@ bool IsBlockForAllEqual(const ForTreeNode& first, const ForTreeNode& second) {
 ir::stmt::For MakeForLoops(const std::vector<int> extents, int index) {
   ir::stmt::StmtRef body_stmt;
   if (index == extents.size() - 1) {
-    body_stmt = ir::stmt::Schedule(std::vector<Var>(),
-                                   std::vector<Expr>(),
-                                   std::vector<Expr>(),
-                                   std::vector<Expr>(),
-                                   "block",
-                                   ir::stmt::BlockRef(0));
+    body_stmt = ir::stmt::Schedule(
+        std::vector<Var>(),
+        std::vector<Expr>(),
+        std::vector<Expr>(),
+        std::vector<Expr>(),
+        "block",
+        ir::stmt::BlockRef(std::vector<ir::stmt::StmtRef>()));
   } else {
     body_stmt = MakeForLoops(extents, index + 1);
   }
