【BUAA】【Infer Symbolic Shape No.211,214】Add prior_box and qr (#67705)

Fripping · web-flow · commit 00685cce02f0 · 2024-09-02T10:12:36.000+08:00
* a

* update cinn api

* update cinn api

* update

* Delete paddle/phi/ops/yaml/.ipynb_checkpoints directory

* Delete paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/.ipynb_checkpoints directory

* Update unary_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
@@ -1339,12 +1339,77 @@ bool MvOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
-// bool PriorBoxOpInferSymbolicShape(pir::Operation *op,
-//                                   pir::InferSymbolicShapeContext
-//                                   *infer_context) {
-//   // pass
-//   return true;
-// }
+bool PriorBoxOpInferSymbolicShape(
+    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_shape = x_shape_or_data.shape();
+  const auto &image_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const std::vector<symbol::DimExpr> &image_shape = image_shape_or_data.shape();
+
+  PADDLE_ENFORCE_EQ(
+      x_shape.size(),
+      4,
+      common::errors::InvalidArgument(
+          "The Input(Input) of Op(PriorBoxOp) should be a 4-D Tensor "
+          "and data format is NCHW. But received Input's dimensions = %d, "
+          "shape = [%s].",
+          x_shape.size(),
+          x_shape));
+  PADDLE_ENFORCE_EQ(
+      image_shape.size(),
+      4,
+      common::errors::InvalidArgument(
+          "The Input(Image) of Op(PriorBoxOp) should be a 4-D Tensor "
+          "and data format is NCHW. But received Image's dimensions = %d, "
+          "shape = [%s].",
+          image_shape.size(),
+          image_shape));
+
+  const std::vector<float> &aspect_ratios =
+      paddle::dialect::details::GetVectorAttr<float>(op, "aspect_ratios");
+  bool flip = op->attribute<pir::BoolAttribute>("flip").data();
+  std::vector<float> aspect_ratios_vec;
+  constexpr float epsilon = 1e-6;
+  aspect_ratios_vec.push_back(1.0f);
+  for (auto ar : aspect_ratios) {
+    bool already_exist = false;
+    for (auto item : aspect_ratios_vec) {
+      if (fabs(ar - item) < epsilon) {
+        already_exist = true;
+        break;
+      }
+    }
+    if (!already_exist) {
+      aspect_ratios_vec.push_back(ar);
+      if (flip) {
+        aspect_ratios_vec.push_back(1.0f / ar);
+      }
+    }
+  }
+  const std::vector<float> &min_sizes =
+      paddle::dialect::details::GetVectorAttr<float>(op, "min_sizes");
+  const std::vector<float> &max_sizes =
+      paddle::dialect::details::GetVectorAttr<float>(op, "max_sizes");
+  int num_priors = aspect_ratios_vec.size() * min_sizes.size();
+  if (!max_sizes.empty()) {
+    num_priors += max_sizes.size();
+  }
+
+  std::vector<symbol::DimExpr> out_shape = {
+      x_shape[2], x_shape[3], symbol::DimExpr(num_priors), symbol::DimExpr(4)};
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+
+  return true;
+}
 
 // bool PullBoxSparseOpInferSymbolicShape(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
@@ -72,7 +72,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixNms)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MarginCrossEntropy)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatmulWithFlatten)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mv)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(PriorBox)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(PriorBox)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullBoxSparse)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullGpuPsSparse)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullSparseV2)
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
@@ -2502,11 +2502,68 @@ bool ProdOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
-// bool QrOpInferSymbolicShape(pir::Operation *op,
-//                             pir::InferSymbolicShapeContext *infer_context) {
-//   // pass
-//   return true;
-// }
+bool QrOpInferSymbolicShape(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_shape = x_shape_or_data.shape();
+  int x_rank = x_shape.size();
+
+  PADDLE_ENFORCE_GE(
+      x_rank,
+      2,
+      common::errors::InvalidArgument("the rank of input must greater than 2"));
+
+  bool compute_q = false;
+  bool reduced = false;
+  const std::string &mode = op->attribute<pir::StrAttribute>("mode").AsString();
+  if (mode == "reduced") {
+    compute_q = true;
+    reduced = true;
+  } else if (mode == "complete") {
+    compute_q = true;
+    reduced = false;
+  } else if (mode == "r") {
+    compute_q = false;
+    reduced = true;
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "QR received unrecognized mode '%s'"
+        " but expected one of 'reduced' (default), 'r', or 'complete'",
+        mode));
+  }
+
+  symbol::DimExpr m = x_shape[x_rank - 2];
+  symbol::DimExpr n = x_shape[x_rank - 1];
+  symbol::DimExprBuilder builder;
+  symbol::DimExpr min_mn = builder.Min(m, n);
+
+  if (compute_q) {
+    symbol::DimExpr k = reduced ? min_mn : m;
+    std::vector<symbol::DimExpr> q_shape = x_shape;
+    q_shape[x_rank - 1] = k;
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(q_shape)});
+  } else {
+    std::vector<symbol::DimExpr> q_shape = {0};
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(q_shape)});
+  }
+
+  symbol::DimExpr k = reduced ? min_mn : m;
+  std::vector<symbol::DimExpr> r_shape = x_shape;
+  r_shape[x_rank - 2] = k;
+  r_shape[x_rank - 1] = n;
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(r_shape)});
+
+  return true;
+}
 
 bool RepeatInterleaveOpInferSymbolicShape(
     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
@@ -112,7 +112,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Pad3d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Pool2d)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Pool)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Prod)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Qr)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Qr)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RepeatInterleave)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Reshape)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Reshape_)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3666,7 +3666,7 @@
     func : prior_box
     data_type : input
   traits : paddle::dialect::ForwardOnlyTrait
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : prod
   args : (Tensor x, IntArray axis, bool keepdim, bool reduce_all)
@@ -3736,7 +3736,7 @@
   kernel :
     func : qr
   backward : qr_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : radam_
   args : (Tensor param, Tensor grad, Tensor learning_rate, Tensor beta1_pow, Tensor beta2_pow, Tensor rho, Tensor moment1, Tensor moment2, Tensor master_param, float beta1 = 0.9f, float beta2 = 0.999f, float epsilon = 1.0e-8f, bool multi_precision = false)
