【Infer Symbolic Shape No.175】【BUAA】Add unfold ops (#67751)

Caogration · web-flow · commit dd261ce31ce1 · 2024-09-04T15:03:05.000+08:00
* Update ops.yaml

* Update unary_infer_sym.cc

* Update unary_infer_sym.h

* Update unary_infer_sym.cc with code style

* Update unary_infer_sym.cc without CalcOutputSize

* Update unary_infer_sym.cc with dyn_cast&lt;int64_t&gt;()

* Update unary_infer_sym.cc with code style

* Update unary_infer_sym.cc with 7 comments from the review Gong fixed

* Update unary_infer_sym.cc with code style of &amp; and *
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
@@ -3685,12 +3685,105 @@ bool Unsqueeze_OpInferSymbolicShape(
   return UnsqueezeOpInferSymbolicShape(op, infer_context);
 }
 
-// bool UnfoldOpInferSymbolicShape(pir::Operation *op,
-//                                 pir::InferSymbolicShapeContext
-//                                 *infer_context) {
-//   // pass
-//   return true;
-// }
+bool UnfoldOpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::ShapeOrDataDimExprs &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+
+  const auto &x_shape = x_shape_or_data.shape();
+
+  PADDLE_ENFORCE_EQ(
+      x_shape.size(),
+      4UL,
+      common::errors::InvalidArgument(
+          "Input should be 4-D tensor of format [N, C, H, W], but get %u",
+          x_shape.size()));
+
+  const std::vector<int> &kernel_sizes =
+      paddle::dialect::details::GetVectorAttr<int>(op, "kernel_sizes");
+  const std::vector<int> &strides =
+      paddle::dialect::details::GetVectorAttr<int>(op, "strides");
+  const std::vector<int> &paddings =
+      paddle::dialect::details::GetVectorAttr<int>(op, "paddings");
+  const std::vector<int> &dilations =
+      paddle::dialect::details::GetVectorAttr<int>(op, "dilations");
+
+  PADDLE_ENFORCE_EQ(
+      x_shape.size() - kernel_sizes.size(),
+      2UL,
+      common::errors::InvalidArgument(
+          "The dims of X should be larger than that of kernel_sizes "
+          "by a number of 2, due to the batch size and input channel dim. "
+          "But received dims(X:%u) - dims(kernel_sizes:%u) != 2",
+          x_shape.size(),
+          kernel_sizes.size()));
+  PADDLE_ENFORCE_EQ(
+      strides.size(),
+      kernel_sizes.size(),
+      common::errors::InvalidArgument(
+          "The dims of strides should be the same with that of kernel_sizes. "
+          "But received dims(strides: %u) != dims(kernel_sizes: %u).",
+          strides.size(),
+          kernel_sizes.size()));
+  PADDLE_ENFORCE_EQ(
+      paddings.size(),
+      2 * strides.size(),
+      common::errors::InvalidArgument(
+          "The dims of paddings should be 2 times of that of strides. "
+          "But received dims(paddings: %u) != 2*dims(strides: %u).",
+          paddings.size(),
+          strides.size()));
+  PADDLE_ENFORCE_EQ(
+      strides.size(),
+      dilations.size(),
+      common::errors::InvalidArgument(
+          "The dims of strides should be the same with that of dilations. "
+          "But received dims(strides: %u) != dims(dilations: %u).",
+          strides.size(),
+          dilations.size()));
+
+  std::vector<symbol::DimExpr> out_shapes;
+
+  out_shapes.push_back(x_shape[0]);
+  out_shapes.push_back(x_shape[1] *
+                       symbol::DimExpr(kernel_sizes[0] * kernel_sizes[1]));
+
+  const auto &calculate_output_dim = [&](symbol::DimExpr input_size,
+                                         int padding1,
+                                         int padding2,
+                                         int dilation,
+                                         int kernel_size,
+                                         int stride) {
+    const symbol::DimExpr dkernel_size =
+        symbol::DimExpr(dilation * (kernel_size - 1) + 1);
+    return (input_size + symbol::DimExpr(padding1 + padding2) - dkernel_size) /
+               symbol::DimExpr(stride) +
+           symbol::DimExpr(1);
+  };
+
+  symbol::DimExpr output_height = calculate_output_dim(x_shape[2],
+                                                       paddings[0],
+                                                       paddings[2],
+                                                       dilations[0],
+                                                       kernel_sizes[0],
+                                                       strides[0]);
+  symbol::DimExpr output_width = calculate_output_dim(x_shape[3],
+                                                      paddings[1],
+                                                      paddings[3],
+                                                      dilations[1],
+                                                      kernel_sizes[1],
+                                                      strides[1]);
+
+  symbol::DimExpr output_col_length = output_height * output_width;
+  out_shapes.push_back(output_col_length);
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shapes)});
+
+  return true;
+}
 
 bool UnstackOpInferSymbolicShape(
     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
@@ -151,7 +151,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unique)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(UniqueConsecutive)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unsqueeze)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unsqueeze_)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unfold)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unfold)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unstack)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightQuantize)
 
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4935,6 +4935,7 @@
   kernel :
     func : unfold
   backward : unfold_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : uniform
   args : (IntArray shape,  DataType dtype,  Scalar min,  Scalar max,  int seed, Place place={})
