[PIR][DynamicShape] Fix bug in InferSymbolicShape ElementWiseBinary (#62455)

lanxianghit · web-flow · commit f2d1f4d35e58 · 2024-03-08T10:42:36.000+08:00
* Fix bug in InferSymbolicShape ElementWiseBinary

* fix bug in fuse pass

* optimize error message

* fix typo

* fix more bugs
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc
@@ -21,6 +21,7 @@
 #include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
 #include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/common/ddim.h"
+#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
@@ -107,8 +108,12 @@ bool MakeGenerateShapeOpAttribute(
     std::vector<pir::Attribute>* output_dim_expr_attrs,
     GenerateShapeOp::SymbolBindings* symbol_bindings) {
   const auto& shape_or_data_dim_exprs = ShapeOrDataDimExprs4Value(output_shape);
-  CHECK(shape_or_data_dim_exprs.data().has_value());
-  const auto& out_dim_exprs = shape_or_data_dim_exprs.data().value();
+  ExprVec data_vec =
+      paddle::dialect::details::GetExprVecFromData(shape_or_data_dim_exprs);
+  // CHECK(shape_or_data_dim_exprs.data().has_value());
+  CHECK(data_vec.size());
+  // const auto& out_dim_exprs = shape_or_data_dim_exprs.data().value();
+  const auto& out_dim_exprs = data_vec;
   return MakeGenerateShapeOpAttribute(ir_context,
                                       ShapeOrDataDimExprs4Value,
                                       out_dim_exprs,
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
@@ -23,7 +23,9 @@ bool InferSymbolicShapeElementWiseBinary(
   // 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()->isa<paddle::dialect::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()) {
     shape_0 = x_shapeordata.data().value();
   } else {
@@ -34,7 +36,9 @@ bool InferSymbolicShapeElementWiseBinary(
       shape_analysis->GetShapeOrDataForValue(op->operand_source(1));
   std::vector<symbol::DimExpr> shape_1;
   bool y_from_fullop =
-      op->operand_source(1).defining_op()->isa<paddle::dialect::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()) {
     shape_1 = y_shapeordata.data().value();
   } else {
@@ -224,4 +228,12 @@ bool Remainder_OpInferSymbolicShape(
   return InferSymbolicShapeElementWiseBinary(op, shape_analysis);
 }
 
+bool SubtractOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  return InferSymbolicShapeElementWiseBinary(op, shape_analysis);
+}
+bool Subtract_OpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  return InferSymbolicShapeElementWiseBinary(op, shape_analysis);
+}
 }  // namespace paddle::dialect
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.h
@@ -53,5 +53,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(NotEqual)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(NotEqual_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Remainder)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Remainder_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Subtract)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Subtract_)
 
 }  // namespace paddle::dialect
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h
@@ -75,6 +75,22 @@ std::vector<T> GetVectorAttr(const ::pir::Operation *op,
   return vec_res;
 }
 
+inline ExprVec GetExprVecFromData(const ShapeOrData &shapeordata) {
+  if (shapeordata.isa<TensorListExprs>()) {
+    ExprVec result;
+    TensorListExprs list =
+        shapeordata.dyn_cast<symbol::TensorListShapeOrDataDimExprs>();
+    for (size_t i = 0; i < list.size(); i++) {
+      for (auto expr : list[i].data().value()) {
+        result.emplace_back(expr);
+      }
+    }
+    return result;
+  } else {
+    return shapeordata.data().value();
+  }
+}
+
 std::optional<std::vector<int64_t>> VecExpr2Int64(const ExprVec &expr_vec);
 
 ExprVec VecInt642Expr(const std::vector<int64_t> &int_vec);
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc
@@ -289,6 +289,21 @@ bool ConcatOpInferSymbolicShape(
   axis = axis >= 0 ? axis : std::max(int64_t(0), int64_t(axis + rank));
 
   if (shape_data_list[0].data().has_value()) {
+    if (rank == 1) {
+      ExprVec data = details::GetExprVecFromData(
+          shape_analysis->GetShapeOrDataForValue(operand_source));
+      const std::vector<symbol::DimExpr> shape{std::int64_t(data.size())};
+      symbol::ShapeOrDataDimExprs shape_data{
+          symbol::TensorShapeOrDataDimExprs(shape, data)};
+      pir::Value res = op->result(0);
+      shape_analysis->SetShapeOrDataForValue(res, shape_data);
+
+      return true;
+    } else {
+      PADDLE_THROW(phi::errors::Unimplemented(
+          op->name() +
+          " 's InferSymbolicShape can NOT deal with rank > 1 now."));
+    }
     std::vector<symbol::DimExpr> data;
     data.reserve(shape_data_list.size());
     for (auto &data_elem : shape_data_list) {
@@ -436,9 +451,9 @@ bool SqueezeOpInferSymbolicShape(
         if (in_dims_sym[current] == 1) {
           should_squeeze[current] = true;
         } else if (!in_dims_sym[current].Has<std::int64_t>()) {
-          PADDLE_THROW(
-              phi::errors::Unimplemented("SqueezeOpInferSymbolicShape CAN NOT "
-                                         "deal with symbol in axis now"));
+          should_squeeze[current] = true;
+        } else {
+          should_squeeze[current] = true;
         }
       }
     }
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.cc
@@ -379,14 +379,7 @@ bool Sinh_OpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return SameOperandsAndResultShape(op, shape_analysis);
 }
-bool SubtractOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  return SameOperandsAndResultShape(op, shape_analysis);
-}
-bool Subtract_OpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  return SameOperandsAndResultShape(op, shape_analysis);
-}
+
 bool TanOpInferSymbolicShape(pir::Operation *op,
                              pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return SameOperandsAndResultShape(op, shape_analysis);
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.h
@@ -105,8 +105,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sin)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sin_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sinh)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sinh_)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Subtract)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Subtract_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tan)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tan_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tanh)
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
@@ -165,6 +165,7 @@ bool Cumsum_OpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return CumsumOpInferSymbolicShape(op, shape_analysis);
 }
+
 bool DiagEmbedOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   pir::Value operand_source = op->operand_source(0);
@@ -280,6 +281,7 @@ bool KthvalueOpInferSymbolicShape(
   shape_analysis->SetShapeOrDataForValue(op->result(1), shape_data);
   return true;
 }
+
 bool ReshapeOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   pir::Value operand_source = op->operand_source(0);
@@ -329,10 +331,11 @@ bool ReshapeOpInferSymbolicShape(
     const auto &numel =
         GetProduct(original_shape, [](const auto &) { return true; });
 
+    ExprVec target_shape = details::GetExprVecFromData(operand_shape_or_data);
     const auto &product_exclude_minus_one =
-        GetProduct(operand_shape_or_data.data().value(), IsNotMinusOne);
+        GetProduct(target_shape, IsNotMinusOne);
 
-    const auto &input_dims = operand_shape_or_data.data().value();
+    const auto &input_dims = target_shape;
 
     std::vector<symbol::DimExpr> out_dims;
     out_dims.reserve(input_dims.size());
diff --git a/paddle/fluid/pir/transforms/shape_optimization_pass.cc b/paddle/fluid/pir/transforms/shape_optimization_pass.cc
@@ -16,6 +16,7 @@
 #include "paddle/common/flags.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 #include "paddle/pir/include/core/dialect.h"
+#include "paddle/pir/include/core/ir_printer.h"
 #include "paddle/pir/include/dialect/shape/ir/shape_attribute.h"
 #include "paddle/pir/include/dialect/shape/ir/shape_dialect.h"
 #include "paddle/pir/include/pass/pass_manager.h"
@@ -31,22 +32,79 @@ namespace {
 using PassPipelineRunner =
     std::function<bool(pir::PassManager&, pir::ModuleOp)>;
 
-void PrintProgram(pir::ModuleOp m, std::string mgs) {
+void PrintProgram(pir::ModuleOp m, std::string msg) {
   ShapeConstraintIRAnalysis& shape_analysis =
       ShapeAnalysisManager::Instance().Get(m.program());
-  VLOG(vlog_level) << "===================== " << mgs
-                   << " =====================\n"
-                   << pir::CustomPrintHelper(*m.program(),
-                                             shape_analysis.PrintHook());
+  if (VLOG_IS_ON(vlog_level)) {
+    std::cerr << "===================== [ShapeDialect]" << msg
+              << " =====================\n"
+              << pir::CustomPrintHelper(*m.program(),
+                                        shape_analysis.PrintHook())
+              << std::endl;
+  }
+}
+
+std::string PrintOperationWithNoRegion(Operation* op) {
+  std::ostringstream os;
+  pir::IrPrinter printer(os);
+
+  // print OpResults
+  os << "(";
+  auto num_op_result = op->num_results();
+  for (size_t idx = 0; idx < num_op_result; idx++) {
+    os << "%op_" << op->id() << "_" << idx;
+    if (idx < num_op_result - 1) os << ", ";
+  }
+  os << ")";
+
+  os << " =";
+
+  // print OpName & OpId
+  os << " \"" << op->name() << "(op_" << op->id() << ")"
+     << "\"";
+
+  // print OpOperands
+  os << " (";
+  auto num_op_operands = op->num_operands();
+  for (size_t idx = 0; idx < num_op_operands; idx++) {
+    const pir::Value& input = op->operand_source(idx);
+    if (input.defining_op()) {
+      os << "op_" << input.defining_op()->id() << "_"
+         << input.dyn_cast<pir::OpResult>().index();
+    } else {
+      os << "op_NULL";
+    }
+    if (idx < num_op_operands - 1) os << ", ";
+  }
+  os << ")";
+
+  printer.PrintAttributeMap(op);
+  os << " :";
+
+  // PrintOpSignature
+  printer.PrintOperandsType(op);
+  os << " -> ";
+
+  printer.PrintOpReturnType(op);
+
+  return os.str();
+}
+
+void PrintOpInfo(pir::Operation* op) {
+  if (VLOG_IS_ON(vlog_level)) {
+    VLOG(vlog_level) << op->name() << "(op_id: op_" << op->id()
+                     << ", num_results=" << op->num_results() << ")"
+                     << " has InferSymbolicShapeInterface.\n\t"
+                     << PrintOperationWithNoRegion(op);
+  }
 }
 
 void DebugPrintOpInfo(
     pir::Operation* op,
     pir::ShapeConstraintIRAnalysis* shape_analysis = nullptr) {
+  std::ostringstream print_stream;
   for (auto& res : op->results()) {
-    std::ostringstream print_stream;
-
-    print_stream << "  result(" << res.dyn_cast<pir::OpResult>().index() << ") "
+    print_stream << "\tresult(" << res.dyn_cast<pir::OpResult>().index() << ") "
                  << "ShapeOrData: {";
 
     if (shape_analysis != nullptr) {
@@ -78,8 +136,10 @@ void DebugPrintOpInfo(
 
       print_stream << "]";
     }
-    print_stream << " }";
-    VLOG(vlog_level) << print_stream.str();
+    print_stream << " }\n";
+  }
+  if (VLOG_IS_ON(vlog_level)) {
+    std::cerr << print_stream.str();
   }
 }
 
@@ -131,8 +191,7 @@ void InferSymExprForBlock(const Block& block,
     auto infer_symbolic_shape_interface =
         op.dyn_cast<paddle::dialect::InferSymbolicShapeInterface>();
     if (infer_symbolic_shape_interface) {
-      VLOG(vlog_level) << op.name() << "(op_id: op_" << op.id() << ")"
-                       << " has InferSymbolicShapeInterface.";
+      PrintOpInfo(&op);
       PADDLE_ENFORCE_EQ(
           infer_symbolic_shape_interface.InferSymbolicShape(shape_analysis),
           true,
