【Infer Symbolic Shape No.176】【BUAA】Add unpool op (#67818)

Caogration · web-flow · commit 3f4837390a5b · 2024-09-19T10:32:58.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
@@ -1821,12 +1821,51 @@ bool Unpool3dOpInferSymbolicShape(
   return true;
 }
 
-// bool UnpoolOpInferSymbolicShape(pir::Operation *op,
-//                                 pir::InferSymbolicShapeContext
-//                                 *infer_context) {
-//   // pass
-//   return true;
-// }
+bool UnpoolOpInferSymbolicShape(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 auto &attributes = op->attributes();
+  std::vector<symbol::DimExpr> output_size;
+  if (attributes.find("output_size") != attributes.end()) {
+    std::vector<int64_t> output_size_int =
+        details::GetVectorAttr<int64_t>(op, "output_size");
+    for (size_t i = 0; i < output_size_int.size(); i++) {
+      output_size.emplace_back(symbol::DimExpr(output_size_int[i]));
+    }
+  } else if (op->operand_source(2)) {
+    const auto &shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(2));
+    output_size =
+        details::GetOrCreateExprVecFromData(shape_or_data, infer_context);
+  }
+
+  std::vector<symbol::DimExpr> output_shape = {x_shape[0], x_shape[1]};
+  for (size_t i = 0; i < ksize.size(); ++i) {
+    output_shape.emplace_back(output_size[i]);
+  }
+
+  if (op->result(0)) {
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(output_shape)});
+  }
+
+  return true;
+}
 
 // bool WeightDequantizeOpInferSymbolicShape(pir::Operation *op,
 //                                           pir::InferSymbolicShapeContext
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
@@ -92,7 +92,7 @@ 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(Unpool)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unpool)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightDequantize)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBox)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBoxHead)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -5009,6 +5009,7 @@
     func: unpool
     data_type: x
   backward: unpool_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : unpool3d
   args: (Tensor x, Tensor indices, int[] ksize, int[] strides={1,1,1}, int[] paddings={0,0,0}, int[] output_size={0,0,0}, str data_format="NCDHW")
