[CINN][BUAA] Add infer_symbol_shape for dot, diag (#67161)

Lans1ot · web-flow · commit 69c60af6b4ed · 2024-08-23T14:51:44.000+08:00
* 修改算子

* 复原测试文件

* 修改数据类型

* 进行修改

* 尝试添加新符号推断

* 删除无用变量

* 修改代码

* 修改代码

* 修改分支语句错误

* 根据review进行修改
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
@@ -428,11 +428,41 @@ bool CrossOpInferSymbolicShape(pir::Operation *op,
 //   return true;
 // }
 
-// bool DotOpInferSymbolicShape(pir::Operation *op,
-//                              pir::InferSymbolicShapeContext *infer_context) {
-//   // pass
-//   return true;
-// }
+bool DotOpInferSymbolicShape(pir::Operation *op,
+                             pir::InferSymbolicShapeContext *infer_context) {
+  const auto x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  auto x_shape = x_shape_or_data.shape();
+  const size_t x_rank = x_shape.size();
+  PADDLE_ENFORCE_EQ(
+      1 == x_rank || 2 == x_rank,
+      true,
+      common::errors::InvalidArgument(
+          "ShapeError: The dimensions of input tensor X (%u) should be 1 or 2",
+          x_rank));
+
+  const auto y_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const auto y_shape = y_shape_or_data.shape();
+  auto y_rank = y_shape.size();
+  PADDLE_ENFORCE_EQ(x_rank == y_rank,
+                    true,
+                    common::errors::InvalidArgument(
+                        "ShapeError: The rank of input tensor Y (%u) must "
+                        "match that of input tensor X(%u).",
+                        y_rank,
+                        x_rank));
+  for (size_t i = 0; i < x_rank; ++i) {
+    infer_context->AddEqualCstr(x_shape[i], y_shape[i]);
+  }
+  // Dot OP require both inputs should have the same shape
+
+  x_shape.erase(x_shape.end() - 1);
+  auto output_shape =
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)};
+  infer_context->SetShapeOrDataForValue(op->result(0), output_shape);
+  return true;
+}
 
 bool DropoutOpInferSymbolicShape(
     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
@@ -37,7 +37,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv3d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cross)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Correlation)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(DepthwiseConv)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dot)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dot)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dropout)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Embedding)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(EqualAll)
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
@@ -625,12 +625,53 @@ bool DecodeJpegOpInferSymbolicShape(
   return true;
 }
 
-// bool DiagOpInferSymbolicShape(pir::Operation *op,
-//                               pir::InferSymbolicShapeContext *infer_context)
-//                               {
-//   // pass
-//   return true;
-// }
+bool DiagOpInferSymbolicShape(pir::Operation *op,
+                              pir::InferSymbolicShapeContext *infer_context) {
+  const auto x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto x_shape = x_shape_or_data.shape();
+  const int offset_data = op->attribute<pir::Int32Attribute>("offset").data();
+  auto offset = symbol::DimExpr(offset_data);
+
+  if (x_shape.size() <= 1) {
+    symbol::DimExpr size_ =
+        (x_shape.size() == 1UL ? x_shape[0] : symbol::DimExpr(1)) +
+        symbol::DimExpr(std::abs(offset_data));
+    infer_context->SetShapeOrDataForValue(
+        op->result(0), symbol::TensorShapeOrDataDimExprs({size_, size_}));
+  } else if (x_shape.size() == 2UL) {
+    if (x_shape[0].isa<int64_t>() && x_shape[1].isa<int64_t>()) {
+      int64_t size_ = 0;
+      if (offset_data >= 0) {
+        if (x_shape[0].dyn_cast<int64_t>() >
+            x_shape[1].dyn_cast<int64_t>() - offset_data) {
+          size_ = x_shape[0].dyn_cast<int64_t>();
+        } else {
+          size_ = x_shape[1].dyn_cast<int64_t>() - offset_data;
+        }
+      } else {
+        if (x_shape[0].dyn_cast<int64_t>() + offset_data <
+            x_shape[1].dyn_cast<int64_t>()) {
+          size_ = x_shape[0].dyn_cast<int64_t>() + offset_data;
+        } else {
+          size_ = x_shape[1].dyn_cast<int64_t>();
+        }
+      }
+      infer_context->SetShapeOrDataForValue(
+          op->result(0), symbol::TensorShapeOrDataDimExprs({size_}));
+    } else {
+      symbol::DimExpr out_unknown =
+          infer_context->GetNextSymName();  // unknown until runtime
+      infer_context->SetShapeOrDataForValue(
+          op->result(0), symbol::TensorShapeOrDataDimExprs({out_unknown}));
+    }
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "diag only support 1D/2D matrix, but input has %u dims",
+        x_shape.size()));
+  }
+  return true;
+}
 
 bool DiagEmbedOpInferSymbolicShape(
     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
@@ -47,7 +47,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumsum_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ChannelShuffle)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Crop)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DecodeJpeg)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Diag)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Diag)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DiagEmbed)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Diagonal)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DistributeFpnProposals)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1384,6 +1384,7 @@
   kernel :
     func : diag
   backward : diag_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : diag_embed
   args : (Tensor input, int offset = 0, int dim1 = -2, int dim2 = -1)
@@ -1455,6 +1456,7 @@
     func : dot
     data_type : x
   backward : dot_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : dpsgd
   args: (Tensor param, Tensor grad, Tensor learning_rate, float clip = 10.0f, float batch_size = 16.0f, float sigma = 1.0f, int seed = 0)
