【Infer Symbolic Shape No.240】【BUAA】unpool3d (#67494)

tlxd · web-flow · commit 95c5e5d88e23 · 2024-08-19T13:55:00.000+08:00
* Update binary_infer_sym.cc binary_infer_sym.h ops.yaml

* Update binary_infer_sym.cc tdm_child + yolo_box

* Update binary_infer_sym.cc
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
@@ -1348,20 +1348,20 @@ bool TdmChildOpInferSymbolicShape(
     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> &input_dims = x_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &input_shape = x_shape_or_data.shape();
   int child_nums = op->attribute<pir::Int32Attribute>("child_nums").data();
 
-  std::vector<symbol::DimExpr> output_dims = input_dims;
-  output_dims.push_back(symbol::DimExpr(child_nums));
+  std::vector<symbol::DimExpr> output_shape = input_shape;
+  output_shape.push_back(symbol::DimExpr(child_nums));
 
   infer_context->SetShapeOrDataForValue(
       op->result(0),
       symbol::ShapeOrDataDimExprs{
-          symbol::TensorShapeOrDataDimExprs(output_dims)});
+          symbol::TensorShapeOrDataDimExprs(output_shape)});
   infer_context->SetShapeOrDataForValue(
       op->result(1),
       symbol::ShapeOrDataDimExprs{
-          symbol::TensorShapeOrDataDimExprs(output_dims)});
+          symbol::TensorShapeOrDataDimExprs(output_shape)});
 
   return true;
 }
@@ -1373,12 +1373,38 @@ bool TdmChildOpInferSymbolicShape(
 //   return true;
 // }
 
-// bool Unpool3dOpInferSymbolicShape(pir::Operation *op,
-//                                   pir::InferSymbolicShapeContext
-//                                   *infer_context) {
-//   // pass
-//   return true;
-// }
+bool Unpool3dOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &indices_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &indices_shape =
+      indices_shape_or_data.shape();
+
+  for (size_t i = 0; i < x_shape.size(); ++i) {
+    infer_context->AddEqualCstr(x_shape[i], indices_shape[i]);
+  }
+
+  const std::vector<int> &ksize = details::GetVectorAttr<int>(op, "ksize");
+  const std::vector<int> &output_size =
+      details::GetVectorAttr<int>(op, "output_size");
+
+  std::vector<symbol::DimExpr> output_shape = {x_shape[0], x_shape[1]};
+  for (int i = 0; i < static_cast<int>(ksize.size()); ++i) {
+    output_shape.emplace_back(symbol::DimExpr(output_size[i]));
+  }
+
+  if (op->result(0)) {
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(output_shape)});
+  }
+
+  return true;
+}
 
 // bool UnpoolOpInferSymbolicShape(pir::Operation *op,
 //                                 pir::InferSymbolicShapeContext
@@ -1405,6 +1431,15 @@ bool YoloBoxOpInferSymbolicShape(
   const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
   int anchor_num = static_cast<int>(anchors.size() / 2);
 
+  bool iou_aware = op->attribute<pir::BoolAttribute>("iou_aware").data();
+  if (iou_aware) {
+    infer_context->AddEqualCstr(x_shape[1],
+                                symbol::DimExpr(anchor_num * (6 + class_num)));
+  } else {
+    infer_context->AddEqualCstr(x_shape[1],
+                                symbol::DimExpr(anchor_num * (5 + class_num)));
+  }
+
   symbol::DimExpr box_num = symbol::DimExpr(0);
   box_num = x_shape[2] * x_shape[3] * symbol::DimExpr(anchor_num);
 
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
@@ -84,7 +84,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(TakeAlongAxis)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TopPSampling)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TdmChild)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(TriangularSolve)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unpool3d)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unpool3d)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unpool)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightDequantize)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBox)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4966,7 +4966,7 @@
     func: unpool3d
     data_type: x
   backward: unpool3d_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : unsqueeze
   args : (Tensor x, IntArray axis = {})
