[PIR][DynamicShape] More logic on shape or data selection in InferSymbolicShape (#62569)

lanxianghit · web-flow · commit a179608f6027 · 2024-03-11T19:03:28.000+08:00
* More logic on shape or data selection in InferSymbolicShape
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.cc
@@ -15,18 +15,23 @@
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 
+bool ShouldUseData(pir::Value val) {
+  if (!val.defining_op()) return false;
+  if (val.defining_op()->isa<paddle::dialect::ShapeOp>()) {
+    return true;
+  }
+  return false;
+}
+
 bool InferSymbolicShapeElementWiseBinary(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   const auto &x_shapeordata =
       shape_analysis->GetShapeOrDataForValue(op->operand_source(0));
   std::vector<symbol::DimExpr> shape_0;
   // For ElementWiseBinary ops, if the input tensor is from full op, the value
   // of fullop is useless, only the shape need doing broadcast
-  bool x_from_fullop =
-      op->operand_source(0).defining_op()
-          ? op->operand_source(0).defining_op()->isa<paddle::dialect::FullOp>()
-          : false;
-  if (!x_from_fullop && x_shapeordata.data().has_value()) {
+  if (ShouldUseData(op->operand_source(0)) &&
+      x_shapeordata.data().has_value()) {
     shape_0 = x_shapeordata.data().value();
   } else {
     shape_0 = x_shapeordata.shape();
@@ -35,11 +40,8 @@ bool InferSymbolicShapeElementWiseBinary(
   const auto &y_shapeordata =
       shape_analysis->GetShapeOrDataForValue(op->operand_source(1));
   std::vector<symbol::DimExpr> shape_1;
-  bool y_from_fullop =
-      op->operand_source(1).defining_op()
-          ? op->operand_source(1).defining_op()->isa<paddle::dialect::FullOp>()
-          : false;
-  if (!y_from_fullop && y_shapeordata.data().has_value()) {
+  if (ShouldUseData(op->operand_source(1)) &&
+      y_shapeordata.data().has_value()) {
     shape_1 = y_shapeordata.data().value();
   } else {
     shape_1 = y_shapeordata.shape();
