[buaa]【Infer Symbolic Shape BUAA No.184】Add conv3d_transpose & conv2d_transpose & conv2d_transpose_bias op (#67635)

crazyxiaoxi · web-flow · commit cc9f28e31dfa · 2024-10-18T14:44:14.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
@@ -423,24 +423,160 @@ bool Conv2dOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
-// bool Conv2dTransposeOpInferSymbolicShape(pir::Operation *op,
-//                                          pir::InferSymbolicShapeContext
-//                                          *infer_context) {
-//   // pass
-//   return true;
-// }
-
 bool Conv3dOpInferSymbolicShape(pir::Operation *op,
                                 pir::InferSymbolicShapeContext *infer_context) {
   return Conv2dOpInferSymbolicShape(op, infer_context);
 }
 
-// bool ConvTransposeOpInferSymbolicShape(pir::Operation *op,
-//                                        pir::InferSymbolicShapeContext
-//                                        *infer_context) {
-//   // pass
-//   return true;
-// }
+bool ConvTransposeFunction(pir::Operation *op,
+                           pir::InferSymbolicShapeContext *infer_context,
+                           std::vector<symbol::DimExpr> output_size) {
+  auto x_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  auto filter_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1)).shape();
+
+  std::vector<int> strides =
+      paddle::dialect::details::GetVectorAttr<int>(op, "strides");
+  std::vector<int> paddings =
+      paddle::dialect::details::GetVectorAttr<int>(op, "paddings");
+  std::vector<int> output_padding =
+      paddle::dialect::details::GetVectorAttr<int>(op, "output_padding");
+  std::string padding_algorithm =
+      op->attribute<pir::StrAttribute>("padding_algorithm").AsString();
+  int groups = op->attribute<pir::Int32Attribute>("groups").data();
+  std::vector<int> dilations =
+      paddle::dialect::details::GetVectorAttr<int>(op, "dilations");
+  std::string data_format =
+      op->attribute<pir::StrAttribute>("data_format").AsString();
+
+  std::vector<symbol::DimExpr> new_paddings;
+  for (const auto &i : paddings) {
+    new_paddings.push_back(symbol::DimExpr{i});
+  }
+  std::vector<symbol::DimExpr> new_dilations;
+  for (const auto &i : dilations) {
+    new_dilations.push_back(symbol::DimExpr{i});
+  }
+
+  PADDLE_ENFORCE_EQ(x_shape.size() == 4 || x_shape.size() == 5,
+                    true,
+                    phi::errors::InvalidArgument(
+                        "Input of Op(conv_transpose) should be 4-D or "
+                        "5-D Tensor. But received: %u-D Tensor",
+                        x_shape.size()));
+  PADDLE_ENFORCE_EQ(
+      x_shape.size(),
+      filter_shape.size(),
+      phi::errors::InvalidArgument(
+          "The input's dimension size and filter's dimension size of "
+          "Op (conv_transpose) should be equal. But received: the shape of "
+          "the dimension size of input is [%d],  the dimension size of filter "
+          "is [%d]. ",
+          x_shape.size(),
+          filter_shape.size()));
+
+  int stride_size = static_cast<int>(strides.size());
+  for (int i = 0; i < stride_size; ++i) {
+    PADDLE_ENFORCE_GT(
+        strides[i],
+        0,
+        phi::errors::InvalidArgument(
+            "The strides of Op(conv_transpose) should be greater than 0. "
+            "But received: the strides of Op(conv_transpose) is [%d].",
+            strides[i]));
+  }
+
+  int in_sub_stride_size = x_shape.size() - stride_size;
+
+  PADDLE_ENFORCE_EQ(
+      x_shape.size() - strides.size(),
+      2U,
+      phi::errors::InvalidArgument(
+          "The input's dimension size minus Attr(stride)'s size must "
+          "be equal to 2 for Op(conv_transpose). But received: [%d], the "
+          "input's dimension size is [%d]"
+          "the Attr(stride)'s size is [%d].",
+          in_sub_stride_size,
+          x_shape.size(),
+          strides.size()));
+
+  if (!output_size.empty())
+    PADDLE_ENFORCE_EQ(
+        output_size.size(),
+        strides.size(),
+        phi::errors::InvalidArgument(
+            "The Attr(output_size) and Attr(stride) of Op(conv_transpose) "
+            "should be the same."));
+
+  if (output_padding.size() > 0)
+    PADDLE_ENFORCE_EQ(
+        output_padding.size(),
+        strides.size(),
+        phi::errors::InvalidArgument(
+            "The Attr(output_padding) and Attr(stride) of Op(conv_transpose) "
+            "should be the same."));
+
+  const bool normal_data = (data_format != "NHWC");
+  const symbol::DimExpr C =
+      (normal_data ? x_shape[1] : x_shape[x_shape.size() - 1]);
+
+  infer_context->AddEqualCstr(filter_shape[0], C);
+
+  const std::vector<symbol::DimExpr> x_data_dims =
+      normal_data
+          ? std::vector<symbol::DimExpr>(x_shape.begin() + 2, x_shape.end())
+          : std::vector<symbol::DimExpr>(x_shape.begin() + 1,
+                                         x_shape.end() - 1);
+
+  const std::vector<symbol::DimExpr> filter_data_dims = [&]() {
+    return std::vector<symbol::DimExpr>(filter_shape.begin() + 2,
+                                        filter_shape.end());
+  }();
+  std::vector<symbol::DimExpr> ksize = filter_data_dims;
+  UpdatePaddingAndDilation(&new_paddings,
+                           &new_dilations,
+                           padding_algorithm,
+                           x_data_dims,
+                           strides,
+                           ksize);
+  std::vector<symbol::DimExpr> output_shape({x_shape[0]});
+
+  if (normal_data) {
+    output_shape.push_back(filter_shape[1] * groups);
+  }
+  const int offset = (normal_data ? 2 : 1);  // NHWC
+  for (int i = 0; i < static_cast<int>(strides.size()); ++i) {
+    symbol::DimExpr filter_extent =
+        new_dilations[i] * (filter_shape[i + 2] - 1) + 1;
+    symbol::DimExpr infer_shape;
+    if (x_shape[i + offset].Has<std::int64_t>()) {
+      if (x_shape[i + offset].Get<int64_t>() > 0) {
+        infer_shape = (x_shape[i + offset] - 1) * symbol::DimExpr(strides[i]) -
+                      new_paddings[2 * i] - new_paddings[2 * i + 1] +
+                      filter_extent;
+      }
+    } else {
+      infer_shape = infer_context->GetNextSymName();
+    }
+    if (!output_size.empty()) {
+      output_shape.push_back(output_size[i]);
+    } else if (!output_padding.empty()) {
+      output_shape.push_back(infer_shape + symbol::DimExpr(output_padding[i]));
+    } else {
+      output_shape.push_back(infer_shape);
+    }
+  }
+
+  if (!normal_data) {
+    output_shape.push_back(filter_shape[1] * groups);
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_shape)});
+  return true;
+}
 
 bool CrossOpInferSymbolicShape(pir::Operation *op,
                                pir::InferSymbolicShapeContext *infer_context) {
@@ -480,6 +616,49 @@ bool CrossOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool Conv3dTransposeOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  std::vector<int> out_size =
+      paddle::dialect::details::GetVectorAttr<int>(op, "output_size");
+  std::vector<symbol::DimExpr> output_size;
+  for (const auto &i : out_size) {
+    output_size.emplace_back(symbol::DimExpr{i});
+  }
+  return ConvTransposeFunction(op, infer_context, output_size);
+}
+
+bool Conv2dTransposeOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  //  if ( op->attributes().find("output_size") !=  op->attributes().end()) {
+  if (op->HasAttribute("output_size")) {
+    std::vector<int> out_size =
+        paddle::dialect::details::GetVectorAttr<int>(op, "output_size");
+    std::vector<symbol::DimExpr> output_size;
+    for (const auto &i : out_size) {
+      output_size.emplace_back(symbol::DimExpr{i});
+    }
+    return ConvTransposeFunction(op, infer_context, output_size);
+  } else {
+    const auto &output_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(2));
+    const std::vector<symbol::DimExpr> &output_size =
+        output_shape_or_data.data().has_value()
+            ? output_shape_or_data.data().value()
+            : output_shape_or_data.shape();
+    return ConvTransposeFunction(op, infer_context, output_size);
+  }
+}
+
+bool Conv2dTransposeBiasOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  std::vector<int> out_size =
+      paddle::dialect::details::GetVectorAttr<int>(op, "output_size");
+  std::vector<symbol::DimExpr> output_size;
+  for (const auto &i : out_size) {
+    output_size.emplace_back(symbol::DimExpr{i});
+  }
+  return ConvTransposeFunction(op, infer_context, output_size);
+}
 // bool CorrelationOpInferSymbolicShape(pir::Operation *op,
 //                                      pir::InferSymbolicShapeContext
 //                                      *infer_context) {
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,8 +31,10 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bmm)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(CholeskySolve)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(CtcAlign)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv2d)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv2dTranspose)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv2dTranspose)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv2dTransposeBias)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv3d)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv3dTranspose)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(ConvTranspose)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cross)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Correlation)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1051,7 +1051,7 @@
     func : conv2d_transpose
     data_type : x
   backward : conv2d_transpose_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : conv2d_transpose_bias
   args : (Tensor x, Tensor filter, Tensor bias, int[] strides={1, 1}, int[] paddings={0, 0}, int[] output_padding={}, IntArray output_size={}, str padding_algorithm="EXPLICIT", int groups=1, int[] dilations={1, 1}, str data_format="NCHW")
@@ -1063,7 +1063,7 @@
     func : conv2d_transpose_bias
     data_type : x
   traits : paddle::dialect::ForwardOnlyTrait
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : conv3d
   args : (Tensor input, Tensor filter, int[] strides={1, 1, 1}, int[] paddings={0, 0, 0}, str padding_algorithm="EXPLICIT", int groups=1, int[] dilations={1, 1, 1}, str data_format="NCDHW")
@@ -1085,7 +1085,7 @@
     func : conv3d_transpose
     data_type : x
   backward : conv3d_transpose_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : copy_to
   args : (Tensor x, Place place, bool blocking)
