[CINN] Refine implement of infer_symbolic_shape for while op (#62665)

zyfncg · web-flow · commit e65239e45ee5 · 2024-03-19T11:19:14.000+08:00
* refine impl of infer symbolic shape for while op

* optimize infer

* refine code

* refine code impl

* polish code

* revert log
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/substitute_dim_expr_based_on_constraints_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/substitute_dim_expr_based_on_constraints_pass.cc
@@ -28,9 +28,14 @@ namespace ir {
 namespace {
 
 template <typename DoEachT>
-void VisitEachOp(cinn::dialect::GroupOp op, const DoEachT& DoEach) {
-  for (pir::Operation* sub_op : op.GetOperators()) {
-    DoEach(sub_op);
+void VisitEachOp(pir::Operation* op, const DoEachT& DoEach) {
+  DoEach(op);
+  for (auto& region : *op) {
+    for (auto& block : region) {
+      for (auto& op_in_block : block) {
+        DoEach(&op_in_block);
+      }
+    }
   }
 }
 
@@ -113,7 +118,7 @@ int GetDimExprPriority(const symbol::DimExpr& dim_expr) {
 std::unordered_map<symbol::DimExpr, symbol::DimExpr> GetDimExprSubstitution(
     pir::ShapeConstraintIRAnalysis* shape_analysis) {
   const std::vector<symbol::DimExprConstraint>& dim_expr_constraints =
-      shape_analysis->CreateDimExprBuilder().constraints();
+      shape_analysis->DimExprBuilder().constraints();
   const cinn::common::UnionFindSet<symbol::DimExpr>& union_find_set = [&]() {
     cinn::common::UnionFindSet<symbol::DimExpr> union_find_set;
     for (const auto& constraint : dim_expr_constraints) {
@@ -147,15 +152,14 @@ std::unordered_map<symbol::DimExpr, symbol::DimExpr> GetDimExprSubstitution(
   return substitution_pattern;
 }
 
-void SubstituteDimExprBasedOnConstraints(pir::Operation* op) {
+void SubstituteDimExprBasedOnConstraints(pir::Operation* region_op) {
   VLOG(4) << "SubstituteDimExprBasedOnConstraints start";
-  auto group_op = op->dyn_cast<cinn::dialect::GroupOp>();
   pir::ShapeConstraintIRAnalysis* shape_analysis =
-      &pir::ShapeAnalysisManager::Instance().Get(group_op->GetParentProgram());
+      &pir::ShapeAnalysisManager::Instance().Get(region_op->GetParentProgram());
   const std::unordered_map<symbol::DimExpr, symbol::DimExpr>&
       substitution_pattern = GetDimExprSubstitution(shape_analysis);
 
-  VisitEachOp(group_op, [&](pir::Operation* op) {
+  VisitEachOp(region_op, [&](pir::Operation* op) {
     VisitEachValue(op, [&](pir::Value value) {
       if (!shape_analysis->HasShapeOrDataForValue(value)) {
         VLOG(4) << "Can not find ShapeOrData for value of op(" << op->name()
@@ -173,6 +177,9 @@ void SubstituteDimExprBasedOnConstraints(pir::Operation* op) {
                                                substituted_shape_or_data);
       }
     });
+    if (op->num_regions() > 0) {
+      return;
+    }
     if (op->num_results() > 0) {
       pir::shape::SetShapeAttrForOp(
           op, shape_analysis->GetShapeOrDataForValue(op->result(0)));
@@ -194,7 +201,7 @@ class SubstituteDimExprBasedOnConstraintsPass : public pir::Pass {
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
-    return op->isa<cinn::dialect::GroupOp>() && op->num_regions() > 0;
+    return op->num_regions() > 0;
   }
 };
 
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
@@ -310,21 +310,21 @@ bool MatmulOpInferSymbolicShape(
 
   if ((ndims_x == ndims_y) && ndims_x >= 2) {
     if (transpose_x_attr == false && transpose_y_attr == false) {
-      shape_analysis->CreateDimExprBuilder().CstrEq(x_dims[ndims_x - 1],
-                                                    y_dims[ndims_x - 2]);
+      shape_analysis->DimExprBuilder().CstrEq(x_dims[ndims_x - 1],
+                                              y_dims[ndims_x - 2]);
     } else if (transpose_x_attr == false && transpose_y_attr == true) {
-      shape_analysis->CreateDimExprBuilder().CstrEq(x_dims[ndims_x - 1],
-                                                    y_dims[ndims_x - 1]);
+      shape_analysis->DimExprBuilder().CstrEq(x_dims[ndims_x - 1],
+                                              y_dims[ndims_x - 1]);
     } else if (transpose_x_attr == true && transpose_y_attr == false) {
-      shape_analysis->CreateDimExprBuilder().CstrEq(x_dims[ndims_x - 2],
-                                                    y_dims[ndims_x - 2]);
+      shape_analysis->DimExprBuilder().CstrEq(x_dims[ndims_x - 2],
+                                              y_dims[ndims_x - 2]);
     } else {
-      shape_analysis->CreateDimExprBuilder().CstrEq(x_dims[ndims_x - 2],
-                                                    y_dims[ndims_x - 1]);
+      shape_analysis->DimExprBuilder().CstrEq(x_dims[ndims_x - 2],
+                                              y_dims[ndims_x - 1]);
     }
 
     for (size_t i = 0; i < ndims_x - 2; ++i) {
-      shape_analysis->CreateDimExprBuilder().CstrEq(x_dims[i], y_dims[i]);
+      shape_analysis->DimExprBuilder().CstrEq(x_dims[i], y_dims[i]);
     }
   }
   return true;
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.cc
@@ -86,7 +86,7 @@ bool ConcatOpInferSymbolicShape(
   };
 
   VLOG(3) << "constraints size:"
-          << shape_analysis->CreateDimExprBuilder().constraints().size();
+          << shape_analysis->DimExprBuilder().constraints().size();
 
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(GetOutDimExprs())};
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.cc
@@ -104,8 +104,8 @@ void BuildCstrEqForTensorListAlongAxis(
     const symbol::TensorListShapeOrDataDimExprs &shape_data_list,
     int axis) {
   for (size_t i = 1; i < shape_data_list.size(); ++i) {
-    shape_analysis->CreateDimExprBuilder().CstrEq(
-        shape_data_list[0].shape()[axis], shape_data_list[i].shape()[axis]);
+    shape_analysis->DimExprBuilder().CstrEq(shape_data_list[0].shape()[axis],
+                                            shape_data_list[i].shape()[axis]);
   }
 }
 
@@ -114,7 +114,7 @@ void BuildCstrEqForTensorListAlongAxis(
     const std::vector<pir::Value> &values,
     int axis) {
   for (size_t i = 1; i < values.size(); ++i) {
-    shape_analysis->CreateDimExprBuilder().CstrEq(
+    shape_analysis->DimExprBuilder().CstrEq(
         shape_analysis->GetShapeOrDataForValue(values[0]).shape()[axis],
         shape_analysis->GetShapeOrDataForValue(values[i]).shape()[axis]);
   }
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
@@ -601,8 +601,8 @@ bool SplitOpInferSymbolicShape(pir::Operation *op,
     const bool &all_sections_sym_not_minus_one =
         All(sections_sym, IsNotMinusOne);
     if (all_sections_sym_not_minus_one) {
-      shape_analysis->CreateDimExprBuilder().CstrEq(x_dims_sym[axis],
-                                                    sum_exclude_minus_one);
+      shape_analysis->DimExprBuilder().CstrEq(x_dims_sym[axis],
+                                              sum_exclude_minus_one);
     }
 
     symbol::TensorListShapeOrDataDimExprs shape_data_list;
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -738,6 +738,46 @@ bool WhileOp::InferSymbolicShape(
 
   pir::InferSymExprForBlock(body(), shape_analysis);
 
+  // add constraints for args
+  const auto &body_args = block_args();
+  for (size_t i = 0; i < body_args.size(); ++i) {
+    const auto &input_arg_shape =
+        shape_analysis->GetShapeOrDataForValue(body_args[i]).shape();
+    const auto &yield_value_shape =
+        shape_analysis
+            ->GetShapeOrDataForValue(body().back().operand_source(i + 1))
+            .shape();
+    PADDLE_ENFORCE_EQ(input_arg_shape.size(),
+                      yield_value_shape.size(),
+                      phi::errors::InvalidArgument(
+                          "while op's input[%d] rank should equal to "
+                          "output[%d]'s rank, Now the rank of input is %d,"
+                          "the rank of output is %d.",
+                          i,
+                          i + 1,
+                          input_arg_shape.size(),
+                          yield_value_shape.size()));
+    const auto &original_input_shape =
+        shape_analysis->GetShapeOrDataForValue(operand_source(i + 1)).shape();
+    for (size_t j = 0; j < input_arg_shape.size(); ++j) {
+      if (input_arg_shape[j].isa<int64_t>()) {
+        continue;
+      }
+      if (input_arg_shape[j] ==
+          yield_value_shape[j]) {  // Dim isn't changed in while
+        shape_analysis->DimExprBuilder().CstrEq(original_input_shape[j],
+                                                input_arg_shape[j]);
+        continue;
+      }
+      if (original_input_shape.size() == yield_value_shape.size() &&
+          original_input_shape[j] == yield_value_shape[j]) {
+        shape_analysis->DimExprBuilder().CstrEq(original_input_shape[j],
+                                                input_arg_shape[j]);
+        continue;
+      }
+    }
+  }
+
   const auto &last_op = body().back();
   for (size_t i = 1; i < last_op.operands_source().size(); ++i) {
     shape_analysis->SetShapeOrDataForValue(
diff --git a/paddle/pir/include/dialect/shape/utils/shape_analysis.h b/paddle/pir/include/dialect/shape/utils/shape_analysis.h
@@ -41,7 +41,7 @@ class IR_API ShapeConstraintIRAnalysis {
   void SetShapeOrDataForValue(Value val,
                               const symbol::ShapeOrDataDimExprs& shape_or_data);
 
-  symbol::DimExprBuilder CreateDimExprBuilder();
+  symbol::DimExprBuilder DimExprBuilder();
 
   // Used to debug
   void PrintShapeOrDatas() const;
diff --git a/paddle/pir/src/dialect/shape/utils/shape_analysis.cc b/paddle/pir/src/dialect/shape/utils/shape_analysis.cc
@@ -65,7 +65,7 @@ void ShapeConstraintIRAnalysis::SetShapeOrDataForValue(
   }
 }
 
-symbol::DimExprBuilder ShapeConstraintIRAnalysis::CreateDimExprBuilder() {
+symbol::DimExprBuilder ShapeConstraintIRAnalysis::DimExprBuilder() {
   return symbol::DimExprBuilder(&constraints_);
 }
 

Original file line number	Diff line number	Diff line change
`@@ -104,8 +104,8 @@ void BuildCstrEqForTensorListAlongAxis(`
`104`	`104`	`const symbol::TensorListShapeOrDataDimExprs &shape_data_list,`
`105`	`105`	`int axis) {`
`106`	`106`	`for (size_t i = 1; i < shape_data_list.size(); ++i) {`
`107`		`- shape_analysis->CreateDimExprBuilder().CstrEq(`
`108`		`- shape_data_list[0].shape()[axis], shape_data_list[i].shape()[axis]);`
	`107`	`+ shape_analysis->DimExprBuilder().CstrEq(shape_data_list[0].shape()[axis],`
	`108`	`+ shape_data_list[i].shape()[axis]);`
`109`	`109`	`}`
`110`	`110`	`}`
`111`	`111`
`@@ -114,7 +114,7 @@ void BuildCstrEqForTensorListAlongAxis(`
`114`	`114`	`const std::vector<pir::Value> &values,`
`115`	`115`	`int axis) {`
`116`	`116`	`for (size_t i = 1; i < values.size(); ++i) {`
`117`		`- shape_analysis->CreateDimExprBuilder().CstrEq(`
	`117`	`+ shape_analysis->DimExprBuilder().CstrEq(`
`118`	`118`	`shape_analysis->GetShapeOrDataForValue(values[0]).shape()[axis],`
`119`	`119`	`shape_analysis->GetShapeOrDataForValue(values[i]).shape()[axis]);`
`120`	`120`	`}`
Original file line number	Diff line number	Diff line change
`@@ -65,7 +65,7 @@ void ShapeConstraintIRAnalysis::SetShapeOrDataForValue(`
`65`	`65`	`}`
`66`	`66`	`}`
`67`	`67`
`68`		`-symbol::DimExprBuilder ShapeConstraintIRAnalysis::CreateDimExprBuilder() {`
	`68`	`+symbol::DimExprBuilder ShapeConstraintIRAnalysis::DimExprBuilder() {`
`69`	`69`	`return symbol::DimExprBuilder(&constraints_);`
`70`	`70`	`}`
`71`	`71`