[buaa]【Infer Symbolic Shape 2 BUAA No.2】Add fractionalMaxPool op (#67210)

crazyxiaoxi · web-flow · commit d103b9140376 · 2024-08-13T11:06:44.000+08:00
* test

together

fix

* define

try

fix

* fractional
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
@@ -761,19 +761,51 @@ bool EigvalshOpInferSymbolicShape(
   return EighOpInferSymbolicShape(op, infer_context);
 }
 
-// bool FractionalMaxPool2DOpInferSymbolicShape(pir::Operation *op,
-//                                              pir::InferSymbolicShapeContext
-//                                              *infer_context) {
-//   // pass
-//   return true;
-// }
+bool FractionalMaxPoolOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const std::vector<symbol::DimExpr> &x_dims = x_shape_or_data.shape();
 
-// bool FractionalMaxPool3DOpInferSymbolicShape(pir::Operation *op,
-//                                              pir::InferSymbolicShapeContext
-//                                              *infer_context) {
-//   // pass
-//   return true;
-// }
+  PADDLE_ENFORCE_EQ(
+      (x_dims.size() == 4 || x_dims.size() == 5),
+      true,
+      phi::errors::InvalidArgument(
+          "Pooling input should be 4-D or 5-D tensor but received %dD-Tensor",
+          x_dims.size()));
+
+  std::vector<int> output_size =
+      paddle::dialect::details::GetVectorAttr<int>(op, "output_size");
+  std::vector<int> kernel_size =
+      paddle::dialect::details::GetVectorAttr<int>(op, "kernel_size");
+
+  std::vector<symbol::DimExpr> output_shape = {x_dims[0], x_dims[1]};
+  for (const auto &dim : output_size) {
+    output_shape.emplace_back(symbol::DimExpr(dim));
+  }
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_shape)});
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_shape)});
+
+  return true;
+}
+
+bool FractionalMaxPool3dOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return FractionalMaxPoolOpInferSymbolicShape(op, infer_context);
+}
+
+bool FractionalMaxPool2dOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return FractionalMaxPoolOpInferSymbolicShape(op, infer_context);
+}
 
 bool FakeChannelWiseQuantizeAbsMaxOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -53,11 +53,12 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(DistributeFpnProposals)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Eig)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Eigh)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Eigvalsh)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FractionalMaxPool2d)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FractionalMaxPool3d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Eigvals)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FractionalMaxPool2D)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FractionalMaxPool3D)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeChannelWiseQuantizeAbsMax)
-
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2c)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftC2r)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FftR2c)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1962,6 +1962,7 @@
   kernel :
     func : fractional_max_pool2d
   backward : fractional_max_pool2d_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : fractional_max_pool3d
   args : (Tensor x, int[] output_size, int[] kernel_size = {0, 0, 0}, float random_u = 0.0, bool return_mask = true)
@@ -1971,6 +1972,7 @@
   kernel :
     func : fractional_max_pool3d
   backward : fractional_max_pool3d_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : frame
   args : (Tensor x, int frame_length, int hop_length, int axis=-1)
