[PIR][DynamicShape] Add InferSymbolicShape for matmul, max (#61587)

Add InferSymbolicShape for matmul, max, and refine some codes & err msg

Author: lanxianghit

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.cc

@@ -20,6 +20,11 @@
 #include "paddle/pir/core/builtin_type_interfaces.h"
 #include "paddle/pir/dialect/shape/ir/shape_attribute.h"
 
+// to make codes shorter
+using ShapeOrData = symbol::ShapeOrDataDimExprs;
+using TensorExprs = symbol::TensorShapeOrDataDimExprs;
+using TensorListExprs = symbol::TensorListShapeOrDataDimExprs;
+
 template <typename T>
 struct AttributeTrait;
 
@@ -78,9 +83,6 @@ bool SameOperandsAndResultShape(
   symbol::ShapeOrDataDimExprs operand_shape_or_data =
       shape_analysis->GetShapeOrDataForValue(operand_source);
 
-  op->set_attribute("symbolic_shape",
-                    pir::shape::SymbolAttribute::get(pir::IrContext::Instance(),
-                                                     operand_shape_or_data));
   pir::Value res = op->result(0);
   shape_analysis->SetShapeOrDataForValue(res, operand_shape_or_data);
   return true;
@@ -143,9 +145,7 @@ bool InferSymbolicShapeElementWiseBinary(
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(shapes)};
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
+
   return true;
 }
 
@@ -184,9 +184,6 @@ bool DataOpInferSymbolicShape(pir::Operation *op,
 
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(sym_dims)};
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
 
   pir::Value res = op->result(0);
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
@@ -263,9 +260,7 @@ bool ShapeOpInferSymbolicShape(pir::Operation *op,
       sym_shape, operand_shape_or_data.shape())};
 
   shape_analysis->SetShapeOrDataForValue(res, shape_or_data);
-  op->set_attribute("symbolic_shape",
-                    pir::shape::SymbolAttribute::get(pir::IrContext::Instance(),
-                                                     shape_or_data));
+
   return true;
 }
 
@@ -305,9 +300,6 @@ bool StackOpInferSymbolicShape(pir::Operation *op,
   symbol::ShapeOrDataDimExprs shape_data(
       symbol::TensorShapeOrDataDimExprs(out_dims, out_dims_data));
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
   pir::Value res = op->result(0);
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
   return true;
@@ -368,9 +360,7 @@ bool ReduceInferDim(pir::Operation *op,
   pir::Value res = op->result(0);
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(shapes)};
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
+
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
   return true;
 }
@@ -441,9 +431,6 @@ bool ReshapeOpInferSymbolicShape(
 
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(out_dims)};
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
 
   pir::Value res0 = op->result(0);
   pir::Value res1 = op->result(1);
@@ -476,10 +463,6 @@ bool FullIntArrayOpInferSymbolicShape(
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(shape, data)};
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
   pir::Value res = op->result(0);
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
   return true;
@@ -537,10 +520,6 @@ bool SliceOpInferSymbolicShape(pir::Operation *op,
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(sym_shape, out_data)};
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
   return true;
 }
@@ -580,10 +559,6 @@ bool FullOpInferSymbolicShape(pir::Operation *op,
       symbol::TensorShapeOrDataDimExprs(sym_shape)};
   shape_data.SetData(sym_data);
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
   pir::Value res = op->result(0);
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
   return true;
@@ -629,10 +604,6 @@ bool ConcatOpInferSymbolicShape(
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(out_dims)};
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
   pir::Value res = op->result(0);
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
 
@@ -685,10 +656,6 @@ bool GatherNdOpInferSymbolicShape(
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(result_sym_dims)};
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
   pir::Value res = op->result(0);
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
 
@@ -702,7 +669,7 @@ bool PowOpInferSymbolicShape(pir::Operation *op,
 bool Pow_OpInferSymbolicShape(pir::Operation *op,
                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
   PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
   return PowOpInferSymbolicShape(op, shape_analysis);
 }
 
@@ -812,10 +779,6 @@ bool SqueezeOpInferSymbolicShape(
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(output_shape_sym)};
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
   pir::Value res = op->result(0);
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
 
@@ -891,10 +854,6 @@ bool UnsqueezeOpInferSymbolicShape(
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(result_sym_dims)};
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
   pir::Value res = op->result(0);
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
 
@@ -949,10 +908,6 @@ bool TileOpInferSymbolicShape(pir::Operation *op,
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(out_shape)};
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
   pir::Value res = op->result(0);
   shape_analysis->SetShapeOrDataForValue(res, shape_data);
 
@@ -962,7 +917,7 @@ bool TileOpInferSymbolicShape(pir::Operation *op,
 bool TransposeOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
   return true;
 }
 bool Transpose_OpInferSymbolicShape(
@@ -1095,35 +1050,132 @@ bool EmbeddingOpInferSymbolicShape(
 bool SparseWeightEmbeddingOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
   return true;
 }
 
 bool ExpandOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
   return true;
 }
 
 bool MatmulOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+  // x_dims can't be const or ref here, in case to be broadcasted
+  std::vector<symbol::DimExpr> x_dims = [&] {
+    std::vector<symbol::DimExpr> dims;
+    const auto &x_shape_or_data =
+        shape_analysis->GetShapeOrDataForValue(op->operand_source(0));
+    if (x_shape_or_data.data().has_value()) {
+      dims = x_shape_or_data.data().value();
+    } else {
+      dims = x_shape_or_data.shape();
+    }
+    return dims;
+  }();
+
+  // y_dims can't be const or ref here, in case to be broadcasted
+  std::vector<symbol::DimExpr> y_dims = [&] {
+    std::vector<symbol::DimExpr> dims;
+    const auto y_shape_or_data =
+        shape_analysis->GetShapeOrDataForValue(op->operand_source(1));
+    if (y_shape_or_data.data().has_value()) {
+      dims = y_shape_or_data.data().value();
+    } else {
+      dims = y_shape_or_data.shape();
+    }
+    return dims;
+  }();
+
+  size_t ndims_x = x_dims.size();
+  size_t ndims_y = y_dims.size();
+
+  const bool x_broadcasted = [&] {
+    bool broadcasted = false;
+    if (ndims_x == 1) {
+      x_dims.insert(x_dims.begin(), 1);
+      ndims_x = 2;
+      broadcasted = true;
+    }
+    return broadcasted;
+  }();
+
+  const bool y_broadcasted = [&] {
+    bool broadcasted = false;
+    if (ndims_y == 1) {
+      y_dims.emplace_back(1);
+      ndims_y = 2;
+      broadcasted = true;
+    }
+    return broadcasted;
+  }();
+
+  std::vector<symbol::DimExpr> out_dims;
+  if (ndims_x > ndims_y) {
+    out_dims.assign(x_dims.begin(), x_dims.end() - 2);
+  } else if (ndims_x < ndims_y) {
+    out_dims.assign(y_dims.begin(), y_dims.end() - 2);
+  } else {
+    symbol::DimExprBuilder builder{nullptr};
+    for (size_t i = 0; i < ndims_x - 2; ++i) {
+      out_dims.emplace_back(builder.Broadcast(x_dims[i], y_dims[i]));
+    }
+  }
+
+  symbol::DimExpr out_M =
+      op->attributes().at("transpose_x").dyn_cast<pir::BoolAttribute>().data()
+          ? x_dims[ndims_x - 1]
+          : x_dims[ndims_x - 2];
+  symbol::DimExpr out_N =
+      op->attributes().at("transpose_y").dyn_cast<pir::BoolAttribute>().data()
+          ? y_dims[ndims_y - 2]
+          : y_dims[ndims_y - 1];
+  if (!x_broadcasted) {
+    out_dims.emplace_back(out_M);
+  }
+  if (!y_broadcasted) {
+    out_dims.emplace_back(out_N);
+  }
+
+  shape_analysis->SetShapeOrDataForValue(op->result(0),
+                                         ShapeOrData{TensorExprs(out_dims)});
+
   return true;
 }
 
 bool MaxOpInferSymbolicShape(pir::Operation *op,
                              pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
-  return true;
+  bool keepdim =
+      op->attributes().at("keepdim").dyn_cast<pir::BoolAttribute>().data();
+
+  const std::vector<int64_t> axis = [&] {
+    pir::Operation *axis_gen_op = op->operand_source(1).defining_op();
+    std::vector<int64_t> axis_vec;
+    if (axis_gen_op->isa<paddle::dialect::FullIntArrayOp>()) {
+      axis_vec = GetVectorAttr(
+          axis_gen_op->dyn_cast<paddle::dialect::FullIntArrayOp>(), "value");
+    } else {
+      // TODO(lanxianghit): there's other source: pir::VectorType,
+      // paddle::dialect::DenseTensorType, but after PRIM, maybe always
+      // FullIntArrayOp, to be confirmed
+      PADDLE_THROW(
+          phi::errors::Unimplemented("MaxOpInferSymbolicShape: 'axis' only "
+                                     "support FullIntArrayOp's result now."));
+    }
+    return axis_vec;
+  }();
+
+  bool reduce_all = axis.size() == 0 ? true : false;
+
+  return ReduceInferDim(op, shape_analysis, axis, keepdim, reduce_all);
 }
 
 bool TrilOpInferSymbolicShape(pir::Operation *op,
                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
   PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
   return true;
 }
 
@@ -1135,7 +1187,7 @@ bool Tril_OpInferSymbolicShape(pir::Operation *op,
 bool WhereOpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis) {
   PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
   return true;
 }
 
@@ -1189,10 +1241,6 @@ bool SliceOpInferSymbolicShape(pir::Operation *op,
   };
   symbol::ShapeOrDataDimExprs shape_data{GetOutDimExprs()};
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
   shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
   return true;
 }
@@ -1239,10 +1287,6 @@ bool ConcatOpInferSymbolicShape(
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(GetOutDimExprs())};
 
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
   shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
   return true;
 }
@@ -1292,9 +1336,7 @@ bool ReshapeOpInferSymbolicShape(
   symbol::ShapeOrDataDimExprs shape_data{
       symbol::TensorShapeOrDataDimExprs(out_dims)};
   shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
+
   return true;
 }