add infer symbol shape (#65623)

zyfncg · web-flow · commit 98f7ececc8b6 · 2024-07-03T10:59:35.000+08:00
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
@@ -516,6 +516,14 @@ bool AccuracyCheckOpInferSymbolicShape(
   return true;
 }
 
+bool ReduceAsOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::ShapeOrDataDimExprs &target_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  infer_context->SetShapeOrDataForValue(op->result(0), target_shape);
+  return true;
+}
+
 bool TakeAlongAxisOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   // input
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h
@@ -31,6 +31,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSample)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedSelect)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Matmul)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReduceAs)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Searchsorted)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TakeAlongAxis)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TopPSampling)
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -232,6 +232,60 @@ std::vector<pir::Type> AddNOp::InferMeta(
   return argument_outputs;
 }
 
+namespace details {
+bool AddNOpInferSymbolicShape(pir::Operation *op,
+                              pir::InferSymbolicShapeContext *infer_context) {
+  const auto &input_list_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  PADDLE_ENFORCE_EQ(
+      input_list_shape.isa<symbol::TensorListShapeOrDataDimExprs>(),
+      true,
+      common::errors::InvalidArgument(
+          "The type of inputs shape should be TensorListShapeOrDataDimExprs"));
+  const auto &inputs_shape =
+      input_list_shape.dyn_cast<symbol::TensorListShapeOrDataDimExprs>();
+  PADDLE_ENFORCE_GT(
+      inputs_shape.size(),
+      0,
+      common::errors::InvalidArgument(
+          "The input tensor X's dimensions of AddNOp "
+          "should be larger than 0. But received X's dimensions %d.",
+          inputs_shape.size()));
+  symbol::TensorShapeOrDataDimExprs candidate_shape = inputs_shape.front();
+  size_t candidate_idx = 0;
+  for (size_t i = 1; i < inputs_shape.size(); ++i) {
+    // 0D tensor
+    if (inputs_shape[i].shape().size() == 0) {
+      continue;
+    }
+    if (candidate_shape.shape().size() == 0) {
+      candidate_shape = inputs_shape[i];
+      candidate_idx = i;
+      continue;
+    }
+    PADDLE_ENFORCE_EQ(candidate_shape,
+                      inputs_shape[i],
+                      common::errors::InvalidArgument(
+                          "The input tensor X of AddNOp must"
+                          " have same shape. But received X[%d]'s shape = "
+                          "[%s], X[%d]'s shape = [%s].",
+                          candidate_idx,
+                          candidate_shape,
+                          i,
+                          inputs_shape[i]));
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0), symbol::ShapeOrDataDimExprs{candidate_shape});
+
+  return true;
+}
+
+}  // namespace details
+
+bool AddNOp::InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context) {
+  return details::AddNOpInferSymbolicShape(this->operation(), infer_context);
+}
+
 OpInfoTuple AddN_Op::GetOpInfo() {
   std::vector<paddle::dialect::OpInputInfo> inputs = {
       paddle::dialect::OpInputInfo(
@@ -380,6 +434,11 @@ std::vector<pir::Type> AddN_Op::InferMeta(
   return argument_outputs;
 }
 
+bool AddN_Op::InferSymbolicShape(
+    pir::InferSymbolicShapeContext *infer_context) {
+  return details::AddNOpInferSymbolicShape(this->operation(), infer_context);
+}
+
 OpInfoTuple AddNArrayOp::GetOpInfo() {
   std::vector<paddle::dialect::OpInputInfo> inputs = {
       OpInputInfo("inputs",
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.h b/paddle/fluid/pir/dialect/operator/ir/manual_op.h
@@ -38,6 +38,7 @@ namespace dialect {
 class AddNOp : public pir::Op<AddNOp,
                               paddle::dialect::OpYamlInfoInterface,
                               paddle::dialect::InferMetaInterface,
+                              paddle::dialect::InferSymbolicShapeInterface,
                               paddle::dialect::VjpInterface,
                               paddle::dialect::DecompInterface> {
  public:
@@ -57,6 +58,7 @@ class AddNOp : public pir::Op<AddNOp,
   static std::vector<pir::Type> InferMeta(
       const std::vector<pir::Value> &input_values,
       pir::AttributeMap *p_attributes);
+  bool InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
 
   static std::vector<std::vector<pir::Value>> Vjp(
       pir::Operation *op,
@@ -70,6 +72,7 @@ class AddNOp : public pir::Op<AddNOp,
 class AddN_Op : public pir::Op<AddN_Op,
                                paddle::dialect::OpYamlInfoInterface,
                                paddle::dialect::InferMetaInterface,
+                               paddle::dialect::InferSymbolicShapeInterface,
                                paddle::dialect::InplaceTrait> {
  public:
   using Op::Op;
@@ -89,6 +92,7 @@ class AddN_Op : public pir::Op<AddN_Op,
   static std::vector<pir::Type> InferMeta(
       const std::vector<pir::Value> &input_values,
       pir::AttributeMap *p_attributes);
+  bool InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
 };
 
 class AddNArrayOp : public pir::Op<AddNArrayOp,
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -2450,8 +2450,7 @@ void InferSymbolicShapePass(
     pir::Program &program) {                          // NOLINT
   pir::IrContext *ctx = pir::IrContext::Instance();
   ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
-  if (pir::shape::HasDynamicShape(program) &&
-      FLAGS_pir_apply_shape_optimization_pass) {
+  if (FLAGS_pir_apply_shape_optimization_pass) {
     pass_manager->AddPass(pir::CreateShapeOptimizationPass());
   }
 }
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3634,6 +3634,7 @@
     func : reduce_as
     data_type : x
   backward : reduce_as_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : reindex_graph
   args : (Tensor x, Tensor neighbors, Tensor count, Tensor hashtable_value, Tensor hashtable_index)

Original file line number	Diff line number	Diff line change
`@@ -2450,8 +2450,7 @@ void InferSymbolicShapePass(`
`2450`	`2450`	`pir::Program &program) { // NOLINT`
`2451`	`2451`	`pir::IrContext *ctx = pir::IrContext::Instance();`
`2452`	`2452`	`ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();`
`2453`		`- if (pir::shape::HasDynamicShape(program) &&`
`2454`		`- FLAGS_pir_apply_shape_optimization_pass) {`
	`2453`	`+ if (FLAGS_pir_apply_shape_optimization_pass) {`
`2455`	`2454`	`pass_manager->AddPass(pir::CreateShapeOptimizationPass());`
`2456`	`2455`	`}`
`2457`	`2456`	`}`