【Infer Symbolic Shape No.71,72,75,85】【BUAA】lu, lu_unpack, mean, one_hot (#67206)

tlxd · web-flow · commit d212fac10762 · 2024-08-19T10:56:29.000+08:00
* Update binary_infer_sym.cc binary_infer_sym.h unary_infer_sym.cc unary_infer_sym.h ops.yaml

* Update binary_infer_sym.cc

* Update unary_infer_sym.cc

* Update unary_infer_sym.cc

* Update binary_infer_sym.cc

* Update unary_infer_sym.cc

* Update unary_infer_sym.cc

* Update unary_infer_sym.cc MeanOp

* Update unary_infer_sym.cc MeanOp

* Update unary_infer_sym.cc MeanOp and OneHotOp

* Update unary_infer_sym.cc OneHotOp

* Update unary_infer_sym.cc OneHotOp

* Update unary_infer_sym.cc

* Update unary_infer_sym.cc MeanOp

* Update unary_infer_sym.cc MeanOp

* Update unary_infer_sym.cc MeanOp

* 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
@@ -826,6 +826,61 @@ bool KronOpInferSymbolicShape(pir::Operation *op,
 //   return true;
 // }
 
+bool LuUnpackOpInferSymbolicShape(
+    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();
+
+  bool unpack_ludata =
+      op->attribute<pir::BoolAttribute>("unpack_ludata").data();
+  bool unpack_pivots =
+      op->attribute<pir::BoolAttribute>("unpack_pivots").data();
+
+  PADDLE_ENFORCE_GE(
+      x_rank,
+      2,
+      common::errors::InvalidArgument(
+          "The rank of input must be greater than or equal to 2."));
+
+  const auto &m = x_shape[x_rank - 1];
+  const auto &n = x_shape[x_rank - 2];
+
+  if (unpack_pivots) {
+    std::vector<symbol::DimExpr> p_shape = x_shape;
+    p_shape[x_rank - 1] = n;
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(p_shape)});
+  }
+
+  if (unpack_ludata) {
+    std::vector<symbol::DimExpr> l_shape = x_shape;
+    std::vector<symbol::DimExpr> u_shape = x_shape;
+    if (m.isa<int64_t>() && n.isa<int64_t>()) {
+      int64_t m_value = static_cast<int64_t>(m.Get<std::int64_t>());
+      int64_t n_value = static_cast<int64_t>(n.Get<std::int64_t>());
+      if (m_value >= n_value) {
+        l_shape[x_rank - 1] = n;
+      } else {
+        u_shape[x_rank - 2] = m;
+      }
+    }
+    infer_context->SetShapeOrDataForValue(
+        op->result(1),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(l_shape)});
+    infer_context->SetShapeOrDataForValue(
+        op->result(2),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(u_shape)});
+  }
+
+  return true;
+}
+
 // bool MatrixRankTolOpInferSymbolicShape(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
@@ -59,6 +59,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSelectStrided)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(KldivLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lstsq)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(LuUnpack)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixRankTol)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedSelect)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Matmul)
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
@@ -1357,6 +1357,61 @@ bool LogsumexpOpInferSymbolicShape(
   return details::ReduceInferDim(op, infer_context, axis, keepdim, reduce_all);
 }
 
+bool LuOpInferSymbolicShape(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 be greater than or equal to 2."));
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)});
+
+  const auto &m = x_shape[x_rank - 1];
+  const auto &n = x_shape[x_rank - 2];
+  symbol::DimExprBuilder builder;
+  symbol::DimExpr min_mn = builder.Min(m, n);
+
+  if (x_rank == 2) {
+    infer_context->SetShapeOrDataForValue(
+        op->result(2),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs({symbol::DimExpr(1)})});
+  } else {
+    std::vector<symbol::DimExpr> infos_shape(x_shape.begin(),
+                                             x_shape.begin() + x_rank - 2);
+    infer_context->SetShapeOrDataForValue(
+        op->result(2),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(infos_shape)});
+  }
+
+  bool pivot = op->attribute<pir::BoolAttribute>("pivot").data();
+  if (pivot) {
+    std::vector<symbol::DimExpr> pivots_shape(x_shape.begin(),
+                                              x_shape.begin() + x_rank - 1);
+    pivots_shape[x_rank - 2] = min_mn;
+    infer_context->SetShapeOrDataForValue(
+        op->result(1),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(pivots_shape)});
+  }
+
+  return true;
+}
+
+bool Lu_OpInferSymbolicShape(pir::Operation *op,
+                             pir::InferSymbolicShapeContext *infer_context) {
+  return LuOpInferSymbolicShape(op, infer_context);
+}
+
 bool MaxOpInferSymbolicShape(pir::Operation *op,
                              pir::InferSymbolicShapeContext *infer_context) {
   bool keepdim = GetBoolAttr(op, "keepdim");
@@ -1596,6 +1651,23 @@ bool MaxPool3dWithIndexOpInferSymbolicShape(
   return MaxPoolWithIndexOpInferSymbolicShape(op, infer_context);
 }
 
+bool MeanOpInferSymbolicShape(pir::Operation *op,
+                              pir::InferSymbolicShapeContext *infer_context) {
+  bool keepdim = GetBoolAttr(op, "keepdim");
+
+  const auto &attributes = op->attributes();
+  std::vector<int64_t> axis;
+  if (attributes.find("axis") != attributes.end()) {
+    axis = op->attribute<paddle::dialect::IntArrayAttribute>("axis")
+               .data()
+               .GetData();
+  }
+
+  bool reduce_all = axis.size() == 0 ? true : false;
+
+  return details::ReduceInferDim(op, infer_context, axis, keepdim, reduce_all);
+}
+
 bool MeanAllOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =
@@ -1734,6 +1806,46 @@ bool NumelOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool OneHotOpInferSymbolicShape(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 &attributes = op->attributes();
+  int64_t num_classes;
+  if (op->operand_source(1)) {
+    const auto &num_classes_shape_or_date =
+        infer_context->GetShapeOrDataForValue(op->operand_source(1));
+    if (attributes.find("num_classes") != attributes.end()) {
+      num_classes = op->attribute<pir::Int64Attribute>("num_classes").data();
+    } else if (num_classes_shape_or_date.data().has_value()) {
+      num_classes =
+          num_classes_shape_or_date.data().value().at(0).Get<int64_t>();
+    } else {
+      PADDLE_ENFORCE_EQ(!num_classes_shape_or_date.data().has_value(),
+                        true,
+                        common::errors::InvalidArgument(
+                            "The depth should have data! Please check."));
+    }
+  }
+
+  const std::vector<symbol::DimExpr> &out_shape = [&] {
+    std::vector<symbol::DimExpr> out_shape = x_shape;
+    out_shape.push_back(symbol::DimExpr(num_classes));
+    return out_shape;
+  }();
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0), symbol::TensorShapeOrDataDimExprs(out_shape));
+
+  return true;
+}
+
+// bool P_NormOpInferSymbolicShape(pir::Operation *op,
+//                                 pir::InferSymbolicShapeContext
+//                                 *infer_context) {
+
 // bool OverlapAddOpInferSymbolicShape(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
@@ -82,10 +82,13 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(L1Norm_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(LpPool2d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logcumsumexp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsumexp)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lu)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lu_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mode)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Max)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Maxout)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Min)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mean)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MeanAll)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixPower)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixRank)
@@ -96,6 +99,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaxPool3dWithIndex)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Norm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Nonzero)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Numel)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(OneHot)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(P_Norm)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(OverlapAdd)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PixelShuffle)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PixelUnshuffle)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3037,6 +3037,7 @@
     data_type : x
   inplace : (x -> out)
   backward : lu_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : lu_unpack
   args : (Tensor x, Tensor y, bool unpack_ludata = true, bool unpack_pivots = true)
@@ -3047,6 +3048,7 @@
     func : lu_unpack
     data_type : x
   backward : lu_unpack_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : margin_cross_entropy
   args : (Tensor logits, Tensor label, bool return_softmax = false, int ring_id = 0, int rank = 0, int nranks = 1, float margin1 = 1.0f, float margin2 = 0.5f, float margin3 = 0.0f, float scale = 64.0f)
@@ -3179,6 +3181,7 @@
   kernel :
     func : mean
   backward : mean_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : mean_all
   args : (Tensor x)
@@ -3470,6 +3473,7 @@
   kernel :
     func : one_hot
   traits : paddle::dialect::ForwardOnlyTrait
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : ones
   args : (IntArray shape, DataType dtype=DataType::FLOAT32, Place place=CPUPlace())
