From 43f3b2d962e9ba0976cff98277b513dd893fec18 Mon Sep 17 00:00:00 2001
From: yulangz <1301481108@qq.com>
Date: Sun, 28 Apr 2024 09:06:43 +0000
Subject: [PATCH 1/4] Add DistributeFpnProposalsOpInferSymbolicShape; Add
 RoiAlignOpInferSymbolicShape; fix GatherOpInferSymbolicShape when axis input
 is tensor; fix Expr set_type implementation of derived class;

---
 paddle/cinn/ir/ir.cc                          | 57 ++++++++++++++
 paddle/cinn/ir/ir.h                           | 26 +++++--
 paddle/cinn/ir/ir_base.cc                     |  7 +-
 paddle/cinn/ir/ir_base.h                      |  2 +-
 .../infer_symbolic_shape/binary_infer_sym.cc  | 13 +++-
 .../multiary_infer_sym.cc                     | 19 +++++
 .../infer_symbolic_shape/multiary_infer_sym.h |  1 +
 .../infer_symbolic_shape/unary_infer_sym.cc   | 75 +++++++++++++++++++
 .../infer_symbolic_shape/unary_infer_sym.h    |  1 +
 paddle/fluid/pir/dialect/operator/ir/ops.yaml |  1 +
 paddle/phi/api/yaml/ops.yaml                  |  1 +
 11 files changed, 188 insertions(+), 15 deletions(-)

diff --git a/paddle/cinn/ir/ir.cc b/paddle/cinn/ir/ir.cc
index a121806e6f3bf8..8ff16c37faa45c 100644
--- a/paddle/cinn/ir/ir.cc
+++ b/paddle/cinn/ir/ir.cc
@@ -209,6 +209,11 @@ void Let::Verify() const {
 
 Type Let::type() const { return symbol.type(); }
 
+void Let::set_type(Type t) {
+  IrNode::set_type(t);
+  symbol->set_type(t);
+}
+
 Expr _Var_::Make(const std::string &name, const Type &type) {
   auto node = new _Var_(name, type);
   return Expr(node);
@@ -405,6 +410,11 @@ const std::string &Store::name() const {
 }
 
 Type Store::type() const { return value.type(); }
+
+void Store::set_type(Type t) {
+  IrNode::set_type(t);
+  value->set_type(t);
+}
 std::vector<Expr *> Store::expr_fields() {
   std::vector<Expr *> exprs({&tensor, &value});
   for (auto &idx : indices) exprs.push_back(&idx);
@@ -610,6 +620,12 @@ Type Load::type() const {
   return type;
 }
 
+void Load::set_type(Type t) {
+  CHECK(tensor.defined());
+  IrNode::set_type(t);
+  tensor->set_type(t);
+}
+
 std::vector<Expr *> Load::expr_fields() {
   std::vector<Expr *> exprs({&tensor});
   for (auto &idx : indices) exprs.push_back(&idx);
@@ -696,6 +712,11 @@ Type Broadcast::type() const {
   return value.type().ElementOf().with_lanes(lanes);
 }
 
+void Broadcast::set_type(Type t) {
+  IrNode::set_type(t);
+  value->set_type(t);
+}
+
 Expr Sum::Make(const std::vector<Expr> &vs) {
   CHECK(!vs.empty());
   if (vs.size() == 1) return vs.front();
@@ -772,6 +793,14 @@ Expr Reduce::Make(Reduce::ReduceType reduce_type,
   n->set_type(body.type());
   return Expr(n);
 }
+
+Type Reduce::type() const { return body.type().ElementOf(); }
+
+void Reduce::set_type(Type t) {
+  IrNode::set_type(t);
+  body->set_type(t);
+}
+
 std::vector<Expr *> Reduce::expr_fields() {
   std::vector<Expr *> res;
   if (init.defined()) {
@@ -798,6 +827,16 @@ void Reduce::Verify() const {
   CHECK_EQ(init.type(), body.type());
 }
 
+Type Select::type() const {
+  CHECK_EQ(true_value.type(), false_value.type());
+  return true_value.type();
+}
+
+void Select::set_type(Type t) {
+  IrNode::set_type(t);
+  true_value->set_type(t);
+}
+
 void Select::Verify() const {
   CHECK(condition.defined());
   CHECK(true_value.defined());
@@ -858,12 +897,30 @@ void MultiOperandVerify(llvm::ArrayRef<Expr> operands) {
   }
 }
 
+Type Product::type() const { return operands().front().type(); }
+
+void Product::set_type(Type t) {
+  IrNode::set_type(t);
+  for (auto &operand : operands()) {
+    operand->set_type(t);
+  }
+}
+
 void Product::Verify() const {
   CHECK_GT(operands().size(), 1UL)
       << "Product node should have more than 1 operands";
   MultiOperandVerify(operands());
 }
 
+Type Sum::type() const { return operands().front().type(); }
+
+void Sum::set_type(Type t) {
+  IrNode::set_type(t);
+  for (auto &operand : operands()) {
+    operand->set_type(t);
+  }
+}
+
 void Sum::Verify() const {
   CHECK_GT(operands().size(), 1UL)
       << "Sum node should have more than 1 operands";
diff --git a/paddle/cinn/ir/ir.h b/paddle/cinn/ir/ir.h
index d711e93ce61abb..a3e42a5a4d9a5d 100644
--- a/paddle/cinn/ir/ir.h
+++ b/paddle/cinn/ir/ir.h
@@ -305,6 +305,8 @@ struct Let : public ExprNode<Let> {
 
   Type type() const override;
 
+  void set_type(Type t) override;
+
   void Verify() const override;
 
   static const IrNodeTy _node_type_ = IrNodeTy::Let;
@@ -476,7 +478,9 @@ struct Reduce : public ExprNode<Reduce> {
                    Expr body,
                    const std::vector<Var>& reduce_axis);
 
-  Type type() const override { return body.type().ElementOf(); }
+  Type type() const override;
+
+  void set_type(Type t) override;
 
   std::vector<Expr*> expr_fields() override;
   std::vector<const Expr*> expr_fields() const override;
@@ -509,10 +513,9 @@ struct Select : public ExprNode<Select> {
     return Expr(node);
   }
 
-  Type type() const override {
-    CHECK_EQ(true_value.type(), false_value.type());
-    return true_value.type();
-  }
+  Type type() const override;
+
+  void set_type(Type t) override;
 
   void Verify() const override;
 
@@ -553,6 +556,8 @@ struct Load : public ExprNode<Load>, public LoadStoreAddrMnger {
 
   Type type() const override;
 
+  void set_type(Type t) override;
+
   static const IrNodeTy _node_type_ = IrNodeTy::Load;
 };
 
@@ -573,6 +578,7 @@ struct Store : public ExprNode<Store>, public LoadStoreAddrMnger {
   const std::string& name() const;
 
   Type type() const override;
+  void set_type(Type type) override;
   Expr index() const;
 
   static const IrNodeTy _node_type_ = IrNodeTy::Store;
@@ -899,6 +905,8 @@ struct Broadcast : public ExprNode<Broadcast> {
 
   Type type() const override;
 
+  void set_type(Type type) override;
+
   void Verify() const override;
 
   std::vector<Expr*> expr_fields() override { return {&value}; }
@@ -932,7 +940,9 @@ struct Product : public ExprNode<Product> {
 
   using ExprNode<Product>::operand;
 
-  Type type() const override { return operands().front().type(); }
+  Type type() const override;
+
+  void set_type(Type t) override;
 
   void Verify() const override;
 
@@ -944,7 +954,9 @@ struct Sum : public ExprNode<Sum> {
 
   using ExprNode<Sum>::operand;
 
-  Type type() const override { return operands().front().type(); }
+  Type type() const override;
+
+  void set_type(Type t) override;
 
   void Verify() const override;
 
diff --git a/paddle/cinn/ir/ir_base.cc b/paddle/cinn/ir/ir_base.cc
index c1b0580d16562e..da67b4e6e14c65 100644
--- a/paddle/cinn/ir/ir_base.cc
+++ b/paddle/cinn/ir/ir_base.cc
@@ -239,10 +239,11 @@ const Expr &IrNode::operand(int i) {
 void IrNode::set_type(Type type) { type_ = type; }
 
 void IrNode::convert_int32_to_int64() {
-  CHECK(type_ == Int(64) || type_ == Int(32) || type_.is_unk())
+  common::Type node_type = type();
+  CHECK(node_type == Int(64) || node_type == Int(32) || node_type.is_unk())
       << "Current only support convert int32_t to int64_t, but get type is "
-      << type_;
-  type_ = Int(64);
+      << node_type;
+  set_type(Int(64));
   for (Expr &operand : operands) {
     operand->convert_int32_to_int64();
   }
diff --git a/paddle/cinn/ir/ir_base.h b/paddle/cinn/ir/ir_base.h
index 236e8afb67fe86..e691038d3c83dd 100644
--- a/paddle/cinn/ir/ir_base.h
+++ b/paddle/cinn/ir/ir_base.h
@@ -162,7 +162,7 @@ class IrNode : public cinn::common::Object {
 
   virtual IrNodeTy node_type() const { return IrNodeTy::kUnk; }
   virtual Type type() const { return type_; }
-  void set_type(Type type);
+  virtual void set_type(Type t);
   //! Elevate int32 to int64 if needed
   void convert_int32_to_int64();
 
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
index 48fdb5b101d5b2..b94bde459c110c 100644
--- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
+++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -235,8 +235,15 @@ bool GatherOpInferSymbolicShape(
     return numel;
   }();
 
-  const auto &axis_shape_or_data =
-      shape_analysis->GetShapeOrDataForValue(op->operand_source(2));
+  int32_t axis;
+  if (op->num_operands() > 2) {
+    const auto &axis_shape_or_data =
+        shape_analysis->GetShapeOrDataForValue(op->operand_source(2));
+    axis =
+        static_cast<int>(axis_shape_or_data.data().value()[0].Get<int64_t>());
+  } else {
+    axis = op->attribute<pir::Int32Attribute>("axis").data();
+  }
 
   const std::vector<symbol::DimExpr> &input_sym_shape =
       input_shape_or_data.data().has_value()
@@ -248,8 +255,6 @@ bool GatherOpInferSymbolicShape(
           ? index_shape_or_data.data().value()
           : index_shape_or_data.shape();
 
-  int axis =
-      static_cast<int>(axis_shape_or_data.data().value()[0].Get<int64_t>());
   if (axis < 0) axis += input_sym_shape.size();
 
   const auto &out_sym_shape = [&] {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
index 4769a19293008e..945c70cd3089f3 100644
--- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
+++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -399,6 +399,25 @@ bool NearestInterpOpInferSymbolicShape(
   return BicubicInterpOpInferSymbolicShape(op, shape_analysis);
 }
 
+bool RoiAlignOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  const auto &x = op->operand_source(0);
+  const auto &boxes = op->operand_source(1);
+
+  const auto &num_boxes =
+      shape_analysis->GetShapeOrDataForValue(boxes).shape()[0];
+  symbol::DimExpr channel_num =
+      shape_analysis->GetShapeOrDataForValue(x).shape()[1];
+
+  int32_t out_h = op->attribute<pir::Int32Attribute>("pooled_height").data();
+  int32_t out_w = op->attribute<pir::Int32Attribute>("pooled_width").data();
+
+  std::vector<symbol::DimExpr> out_dim = {num_boxes, channel_num, out_h, out_w};
+  shape_analysis->SetShapeOrDataForValue(
+      op->result(0), symbol::TensorShapeOrDataDimExprs(out_dim));
+  return true;
+}
+
 bool MeshgridOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   const symbol::TensorListShapeOrDataDimExprs &shape_data_list =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
index 2e5f506b373ac4..3e852872c84213 100644
--- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
+++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
@@ -29,6 +29,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(LinearInterp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Meshgrid)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(NearestInterp)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(RoiAlign)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Stack)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TrilinearInterp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Where)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
index 1adc5767862474..4604f262669340 100644
--- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
+++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -247,6 +247,81 @@ bool DiagonalOpInferSymbolicShape(
   return true;
 }
 
+bool DistributeFpnProposalsOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  const auto &attributes = op->attributes();
+  int32_t min_level =
+      attributes.at("min_level").dyn_cast<pir::Int32Attribute>().data();
+  int32_t max_level =
+      attributes.at("max_level").dyn_cast<pir::Int32Attribute>().data();
+  int32_t num_levels = max_level - min_level + 1;
+
+  symbol::DimExpr num_rois = [&]() {
+    pir::Value rois_num = op->operand_source(1);
+    const auto &rois_num_shape_or_data =
+        shape_analysis->GetShapeOrDataForValue(rois_num);
+
+    PADDLE_ENFORCE_EQ(
+        rois_num_shape_or_data.shape()[0],
+        1,
+        phi::errors::InvalidArgument("DistributeFpnProposalsOp in pir model "
+                                     "only support batch_size=1 now."));
+
+    PADDLE_ENFORCE_EQ(rois_num_shape_or_data.data().has_value(),
+                      true,
+                      phi::errors::InvalidArgument(
+                          "InferSymbolicShape of DistributeFpnProposalsdOp "
+                          "only support input with rois_num."));
+
+    const auto &rois_num_value = rois_num_shape_or_data.data().value()[0];
+    CHECK(rois_num_value.isa<int64_t>() || rois_num_value.isa<std::string>())
+        << "rois_num must be int64 or SymName.";
+    if (rois_num_value.isa<int64_t>()) {
+      return symbol::DimExpr(rois_num_value.Get<int64_t>());
+    } else {
+      return symbol::DimExpr(rois_num_value.Get<std::string>());
+    }
+  }();
+
+  const auto &multi_rois_out_shape = [&]() {
+    symbol::TensorListShapeOrDataDimExprs multi_rois_out_shape;
+    if (num_levels == 1) {
+      multi_rois_out_shape.emplace_back(
+          symbol::TensorShapeOrDataDimExprs({num_rois, 4}));
+    } else {
+      symbol::DimExpr last_dim = num_rois;
+      for (int i = 0; i < num_levels - 1; i++) {
+        const auto &next_sym_name = shape_analysis->GetNextSymName();
+        std::vector<symbol::DimExpr> level_dim = {next_sym_name, 4};
+        multi_rois_out_shape.emplace_back(
+            symbol::TensorShapeOrDataDimExprs(level_dim));
+        last_dim = last_dim - level_dim[0];
+      }
+      multi_rois_out_shape.emplace_back(
+          symbol::TensorShapeOrDataDimExprs({last_dim, 4}));
+    }
+
+    return multi_rois_out_shape;
+  }();
+
+  const auto &rois_num_per_level_out_shape = [&]() {
+    symbol::TensorListShapeOrDataDimExprs rois_num_per_level_out_shape;
+    rois_num_per_level_out_shape.resize(num_levels,
+                                        symbol::TensorShapeOrDataDimExprs({1}));
+    return rois_num_per_level_out_shape;
+  }();
+
+  const auto &restore_ind = [&]() {
+    return symbol::TensorShapeOrDataDimExprs({num_levels, 1});
+  }();
+
+  shape_analysis->SetShapeOrDataForValue(op->result(0), multi_rois_out_shape);
+  shape_analysis->SetShapeOrDataForValue(op->result(1),
+                                         rois_num_per_level_out_shape);
+  shape_analysis->SetShapeOrDataForValue(op->result(2), restore_ind);
+  return true;
+}
+
 bool EinsumOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   PADDLE_THROW(phi::errors::Unimplemented(
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
index 8155b4f0ac5052..e260add8ddda0c 100644
--- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
+++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -29,6 +29,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumsum)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumsum_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DiagEmbed)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Diagonal)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(DistributeFpnProposals)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Einsum)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kthvalue)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logcumsumexp)
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
index 7b24d6c4fec9a0..7481e7bf31ef82 100644
--- a/paddle/fluid/pir/dialect/operator/ir/ops.yaml
+++ b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -525,6 +525,7 @@
     func : distribute_fpn_proposals
     data_type : fpn_rois
   optional : rois_num, multi_level_rois_num
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : distributed_fused_lamb
   args : (Tensor[] param, Tensor[] grad, Tensor fp32_fused_param, Tensor fp32_fused_grad, Tensor fp16_fused_param, Tensor fp16_fused_grad, Tensor moment1, Tensor moment2, Tensor beta1pow, Tensor beta2pow, Tensor fused_param_offsets, Tensor fp32_shard_fused_param_offsets, Tensor fp16_shard_fused_param_offsets, Tensor param_info, Tensor param_order, Tensor learning_rate, Tensor global_scale, float beta1, float beta2, float epsilon, float max_global_grad_norm, float weight_decay, bool clip_after_allreduce, int[] ring_ids= {}, int acc_steps = 1, bool use_master_param_norm = true, bool use_master_acc_grad = true, bool is_grad_scaled_by_nranks = true, int64_t nranks = 1, bool use_hierarchical_allreduce = false)
diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
index 1093c09c904ef1..ae7e30c41cceb7 100755
--- a/paddle/phi/api/yaml/ops.yaml
+++ b/paddle/phi/api/yaml/ops.yaml
@@ -2433,6 +2433,7 @@
     data_type : x
   optional : boxes_num
   backward : roi_align_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : roi_pool
   args : (Tensor x, Tensor boxes, Tensor boxes_num, int pooled_height=1, int pooled_width=1, float spatial_scale=1.0)

From b17b23d65e23e0e11d7d160051c8aba6e805a811 Mon Sep 17 00:00:00 2001
From: yulangz <1301481108@qq.com>
Date: Tue, 30 Apr 2024 15:34:15 +0000
Subject: [PATCH 2/4] fix

---
 paddle/cinn/ir/ir.cc                          | 44 -------------------
 paddle/cinn/ir/ir.h                           | 16 +------
 paddle/cinn/ir/ir_base.cc                     |  7 ++-
 paddle/cinn/ir/ir_base.h                      |  2 +-
 .../infer_symbolic_shape/unary_infer_sym.cc   |  2 +-
 5 files changed, 6 insertions(+), 65 deletions(-)

diff --git a/paddle/cinn/ir/ir.cc b/paddle/cinn/ir/ir.cc
index 87d61d11a2dbfa..268538cbaf6803 100644
--- a/paddle/cinn/ir/ir.cc
+++ b/paddle/cinn/ir/ir.cc
@@ -210,11 +210,6 @@ void Let::Verify() const {
 
 Type Let::type() const { return symbol.type(); }
 
-void Let::set_type(Type t) {
-  IrNode::set_type(t);
-  symbol->set_type(t);
-}
-
 Expr _Var_::Make(const std::string &name, const Type &type) {
   auto node = new _Var_(name, type);
   return Expr(node);
@@ -413,10 +408,6 @@ const std::string &Store::name() const {
 
 Type Store::type() const { return value.type(); }
 
-void Store::set_type(Type t) {
-  IrNode::set_type(t);
-  value->set_type(t);
-}
 std::vector<Expr *> Store::expr_fields() {
   std::vector<Expr *> exprs({&tensor, &value});
   for (auto &idx : indices) exprs.push_back(&idx);
@@ -624,12 +615,6 @@ Type Load::type() const {
   return type;
 }
 
-void Load::set_type(Type t) {
-  CHECK(tensor.defined());
-  IrNode::set_type(t);
-  tensor->set_type(t);
-}
-
 std::vector<Expr *> Load::expr_fields() {
   std::vector<Expr *> exprs({&tensor});
   for (auto &idx : indices) exprs.push_back(&idx);
@@ -716,11 +701,6 @@ Type Broadcast::type() const {
   return value.type().ElementOf().with_lanes(lanes);
 }
 
-void Broadcast::set_type(Type t) {
-  IrNode::set_type(t);
-  value->set_type(t);
-}
-
 Expr Sum::Make(const std::vector<Expr> &vs) {
   CHECK(!vs.empty());
   if (vs.size() == 1) return vs.front();
@@ -800,11 +780,6 @@ Expr Reduce::Make(Reduce::ReduceType reduce_type,
 
 Type Reduce::type() const { return body.type().ElementOf(); }
 
-void Reduce::set_type(Type t) {
-  IrNode::set_type(t);
-  body->set_type(t);
-}
-
 std::vector<Expr *> Reduce::expr_fields() {
   std::vector<Expr *> res;
   if (init.defined()) {
@@ -836,11 +811,6 @@ Type Select::type() const {
   return true_value.type();
 }
 
-void Select::set_type(Type t) {
-  IrNode::set_type(t);
-  true_value->set_type(t);
-}
-
 void Select::Verify() const {
   CHECK(condition.defined());
   CHECK(true_value.defined());
@@ -903,13 +873,6 @@ void MultiOperandVerify(llvm::ArrayRef<Expr> operands) {
 
 Type Product::type() const { return operands().front().type(); }
 
-void Product::set_type(Type t) {
-  IrNode::set_type(t);
-  for (auto &operand : operands()) {
-    operand->set_type(t);
-  }
-}
-
 void Product::Verify() const {
   CHECK_GT(operands().size(), 1UL)
       << "Product node should have more than 1 operands";
@@ -918,13 +881,6 @@ void Product::Verify() const {
 
 Type Sum::type() const { return operands().front().type(); }
 
-void Sum::set_type(Type t) {
-  IrNode::set_type(t);
-  for (auto &operand : operands()) {
-    operand->set_type(t);
-  }
-}
-
 void Sum::Verify() const {
   CHECK_GT(operands().size(), 1UL)
       << "Sum node should have more than 1 operands";
diff --git a/paddle/cinn/ir/ir.h b/paddle/cinn/ir/ir.h
index a3e42a5a4d9a5d..72429f6f0556a1 100644
--- a/paddle/cinn/ir/ir.h
+++ b/paddle/cinn/ir/ir.h
@@ -305,8 +305,6 @@ struct Let : public ExprNode<Let> {
 
   Type type() const override;
 
-  void set_type(Type t) override;
-
   void Verify() const override;
 
   static const IrNodeTy _node_type_ = IrNodeTy::Let;
@@ -480,8 +478,6 @@ struct Reduce : public ExprNode<Reduce> {
 
   Type type() const override;
 
-  void set_type(Type t) override;
-
   std::vector<Expr*> expr_fields() override;
   std::vector<const Expr*> expr_fields() const override;
 
@@ -515,8 +511,6 @@ struct Select : public ExprNode<Select> {
 
   Type type() const override;
 
-  void set_type(Type t) override;
-
   void Verify() const override;
 
   std::vector<Expr*> expr_fields() override {
@@ -556,8 +550,6 @@ struct Load : public ExprNode<Load>, public LoadStoreAddrMnger {
 
   Type type() const override;
 
-  void set_type(Type t) override;
-
   static const IrNodeTy _node_type_ = IrNodeTy::Load;
 };
 
@@ -578,7 +570,7 @@ struct Store : public ExprNode<Store>, public LoadStoreAddrMnger {
   const std::string& name() const;
 
   Type type() const override;
-  void set_type(Type type) override;
+
   Expr index() const;
 
   static const IrNodeTy _node_type_ = IrNodeTy::Store;
@@ -905,8 +897,6 @@ struct Broadcast : public ExprNode<Broadcast> {
 
   Type type() const override;
 
-  void set_type(Type type) override;
-
   void Verify() const override;
 
   std::vector<Expr*> expr_fields() override { return {&value}; }
@@ -942,8 +932,6 @@ struct Product : public ExprNode<Product> {
 
   Type type() const override;
 
-  void set_type(Type t) override;
-
   void Verify() const override;
 
   static const IrNodeTy _node_type_ = IrNodeTy::Product;
@@ -956,8 +944,6 @@ struct Sum : public ExprNode<Sum> {
 
   Type type() const override;
 
-  void set_type(Type t) override;
-
   void Verify() const override;
 
   static const IrNodeTy _node_type_ = IrNodeTy::Sum;
diff --git a/paddle/cinn/ir/ir_base.cc b/paddle/cinn/ir/ir_base.cc
index da67b4e6e14c65..c1b0580d16562e 100644
--- a/paddle/cinn/ir/ir_base.cc
+++ b/paddle/cinn/ir/ir_base.cc
@@ -239,11 +239,10 @@ const Expr &IrNode::operand(int i) {
 void IrNode::set_type(Type type) { type_ = type; }
 
 void IrNode::convert_int32_to_int64() {
-  common::Type node_type = type();
-  CHECK(node_type == Int(64) || node_type == Int(32) || node_type.is_unk())
+  CHECK(type_ == Int(64) || type_ == Int(32) || type_.is_unk())
       << "Current only support convert int32_t to int64_t, but get type is "
-      << node_type;
-  set_type(Int(64));
+      << type_;
+  type_ = Int(64);
   for (Expr &operand : operands) {
     operand->convert_int32_to_int64();
   }
diff --git a/paddle/cinn/ir/ir_base.h b/paddle/cinn/ir/ir_base.h
index e691038d3c83dd..236e8afb67fe86 100644
--- a/paddle/cinn/ir/ir_base.h
+++ b/paddle/cinn/ir/ir_base.h
@@ -162,7 +162,7 @@ class IrNode : public cinn::common::Object {
 
   virtual IrNodeTy node_type() const { return IrNodeTy::kUnk; }
   virtual Type type() const { return type_; }
-  virtual void set_type(Type t);
+  void set_type(Type type);
   //! Elevate int32 to int64 if needed
   void convert_int32_to_int64();
 
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
index 4604f262669340..d30cbdb735c396 100644
--- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
+++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -312,7 +312,7 @@ bool DistributeFpnProposalsOpInferSymbolicShape(
   }();
 
   const auto &restore_ind = [&]() {
-    return symbol::TensorShapeOrDataDimExprs({num_levels, 1});
+    return symbol::TensorShapeOrDataDimExprs({num_rois, 1});
   }();
 
   shape_analysis->SetShapeOrDataForValue(op->result(0), multi_rois_out_shape);

From 1f5dfc9d7501c0225e021182c2bea98b48f18b02 Mon Sep 17 00:00:00 2001
From: yulangz <1301481108@qq.com>
Date: Wed, 8 May 2024 06:55:16 +0000
Subject: [PATCH 3/4] change shape_analysis to infer_context

---
 .../infer_symbolic_shape/multiary_infer_sym.cc |  8 ++++----
 .../infer_symbolic_shape/unary_infer_sym.cc    | 18 +++++++++---------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
index 170ad3d15f00ac..4a8b0c1000cc0b 100644
--- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
+++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -505,20 +505,20 @@ bool MemoryEfficientAttentionOpInferSymbolicShape(
 }
 
 bool RoiAlignOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &x = op->operand_source(0);
   const auto &boxes = op->operand_source(1);
 
   const auto &num_boxes =
-      shape_analysis->GetShapeOrDataForValue(boxes).shape()[0];
+      infer_context->GetShapeOrDataForValue(boxes).shape()[0];
   symbol::DimExpr channel_num =
-      shape_analysis->GetShapeOrDataForValue(x).shape()[1];
+      infer_context->GetShapeOrDataForValue(x).shape()[1];
 
   int32_t out_h = op->attribute<pir::Int32Attribute>("pooled_height").data();
   int32_t out_w = op->attribute<pir::Int32Attribute>("pooled_width").data();
 
   std::vector<symbol::DimExpr> out_dim = {num_boxes, channel_num, out_h, out_w};
-  shape_analysis->SetShapeOrDataForValue(
+  infer_context->SetShapeOrDataForValue(
       op->result(0), symbol::TensorShapeOrDataDimExprs(out_dim));
   return true;
 }
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
index e5da9be8add0a8..90f46d0c6940ac 100644
--- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
+++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -247,7 +247,7 @@ bool DiagonalOpInferSymbolicShape(
 }
 
 bool DistributeFpnProposalsOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &attributes = op->attributes();
   int32_t min_level =
       attributes.at("min_level").dyn_cast<pir::Int32Attribute>().data();
@@ -258,7 +258,7 @@ bool DistributeFpnProposalsOpInferSymbolicShape(
   symbol::DimExpr num_rois = [&]() {
     pir::Value rois_num = op->operand_source(1);
     const auto &rois_num_shape_or_data =
-        shape_analysis->GetShapeOrDataForValue(rois_num);
+        infer_context->GetShapeOrDataForValue(rois_num);
 
     PADDLE_ENFORCE_EQ(
         rois_num_shape_or_data.shape()[0],
@@ -290,7 +290,7 @@ bool DistributeFpnProposalsOpInferSymbolicShape(
     } else {
       symbol::DimExpr last_dim = num_rois;
       for (int i = 0; i < num_levels - 1; i++) {
-        const auto &next_sym_name = shape_analysis->GetNextSymName();
+        const auto &next_sym_name = infer_context->GetNextSymName();
         std::vector<symbol::DimExpr> level_dim = {next_sym_name, 4};
         multi_rois_out_shape.emplace_back(
             symbol::TensorShapeOrDataDimExprs(level_dim));
@@ -314,15 +314,15 @@ bool DistributeFpnProposalsOpInferSymbolicShape(
     return symbol::TensorShapeOrDataDimExprs({num_rois, 1});
   }();
 
-  shape_analysis->SetShapeOrDataForValue(op->result(0), multi_rois_out_shape);
-  shape_analysis->SetShapeOrDataForValue(op->result(1),
-                                         rois_num_per_level_out_shape);
-  shape_analysis->SetShapeOrDataForValue(op->result(2), restore_ind);
+  infer_context->SetShapeOrDataForValue(op->result(0), multi_rois_out_shape);
+  infer_context->SetShapeOrDataForValue(op->result(1),
+                                        rois_num_per_level_out_shape);
+  infer_context->SetShapeOrDataForValue(op->result(2), restore_ind);
   return true;
 }
 
-bool EinsumOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+bool EinsumOpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
   PADDLE_THROW(phi::errors::Unimplemented(
       op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
   return true;

From f1d2dee98d764c36f1a3b6846b71e2866bbae782 Mon Sep 17 00:00:00 2001
From: yulangz <1301481108@qq.com>
Date: Wed, 8 May 2024 13:15:55 +0000
Subject: [PATCH 4/4] fix

---
 paddle/cinn/ir/ir.cc                          |  3 +-
 .../infer_symbolic_shape/unary_infer_sym.cc   | 35 ++++++++++---------
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/paddle/cinn/ir/ir.cc b/paddle/cinn/ir/ir.cc
index 268538cbaf6803..7157fcf5c15ce3 100644
--- a/paddle/cinn/ir/ir.cc
+++ b/paddle/cinn/ir/ir.cc
@@ -807,7 +807,8 @@ void Reduce::Verify() const {
 }
 
 Type Select::type() const {
-  CHECK_EQ(true_value.type(), false_value.type());
+  PADDLE_ENFORCE_EQ(
+      true_value.type(), false_value.type(), "Type of Select must be same");
   return true_value.type();
 }
 
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
index 90f46d0c6940ac..f84ca545188bfa 100644
--- a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
+++ b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -254,32 +254,35 @@ bool DistributeFpnProposalsOpInferSymbolicShape(
   int32_t max_level =
       attributes.at("max_level").dyn_cast<pir::Int32Attribute>().data();
   int32_t num_levels = max_level - min_level + 1;
+  int64_t batch_size = 1;
 
   symbol::DimExpr num_rois = [&]() {
     pir::Value rois_num = op->operand_source(1);
     const auto &rois_num_shape_or_data =
         infer_context->GetShapeOrDataForValue(rois_num);
 
-    PADDLE_ENFORCE_EQ(
-        rois_num_shape_or_data.shape()[0],
-        1,
-        phi::errors::InvalidArgument("DistributeFpnProposalsOp in pir model "
-                                     "only support batch_size=1 now."));
-
+    batch_size = rois_num_shape_or_data.shape()[0].Get<int64_t>();
     PADDLE_ENFORCE_EQ(rois_num_shape_or_data.data().has_value(),
                       true,
-                      phi::errors::InvalidArgument(
+                      ::common::errors::InvalidArgument(
                           "InferSymbolicShape of DistributeFpnProposalsdOp "
                           "only support input with rois_num."));
 
-    const auto &rois_num_value = rois_num_shape_or_data.data().value()[0];
-    CHECK(rois_num_value.isa<int64_t>() || rois_num_value.isa<std::string>())
-        << "rois_num must be int64 or SymName.";
-    if (rois_num_value.isa<int64_t>()) {
-      return symbol::DimExpr(rois_num_value.Get<int64_t>());
-    } else {
-      return symbol::DimExpr(rois_num_value.Get<std::string>());
+    symbol::DimExpr rois_total_num = 0;
+    for (int i = 0; i < batch_size; i++) {
+      const auto &rois_num_value = rois_num_shape_or_data.data().value()[i];
+
+      CHECK(rois_num_value.isa<int64_t>() || rois_num_value.isa<std::string>())
+          << "rois_num must be int64 or SymName.";
+      if (rois_num_value.isa<int64_t>()) {
+        return symbol::DimExpr(rois_num_value.Get<int64_t>());
+      } else {
+        return symbol::DimExpr(rois_num_value.Get<std::string>());
+      }
+      rois_total_num = rois_total_num + rois_num_value;
     }
+
+    return rois_total_num;
   }();
 
   const auto &multi_rois_out_shape = [&]() {
@@ -305,8 +308,8 @@ bool DistributeFpnProposalsOpInferSymbolicShape(
 
   const auto &rois_num_per_level_out_shape = [&]() {
     symbol::TensorListShapeOrDataDimExprs rois_num_per_level_out_shape;
-    rois_num_per_level_out_shape.resize(num_levels,
-                                        symbol::TensorShapeOrDataDimExprs({1}));
+    rois_num_per_level_out_shape.resize(
+        num_levels, symbol::TensorShapeOrDataDimExprs({batch_size}));
     return rois_num_per_level_out_shape;
   }();