[Infer Symbolic Shape No.27][BUAA]edit_distance (#67117)

MufanColin · web-flow · commit 3aaa0270ffc7 · 2024-08-12T10:39:45.000+08:00
* Fixd edit_distance op

* Added one more constraint

* Added one more constraint

* Resolved suggested changes

* Added symbol:: before DimExpr

* Fixed errors

* Removed two unnecessary variables and added .shape()

* Fixed typos

* Revert to a previous version

* added .shape()

* added .shape

* added .shape()

* Try to resolve type conflicts

* Try to resolve type conflicts

* Try to resolve type conflicts

* Try to resolve type conflicts
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -812,6 +812,67 @@ bool ConcatOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool EditDistanceOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &hyps_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &refs_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+
+  const auto &hyps_dims = hyps_shape_or_data.shape();
+  const auto &refs_dims = refs_shape_or_data.shape();
+
+  PADDLE_ENFORCE_EQ(
+      hyps_dims.size(),
+      2,
+      phi::errors::InvalidArgument(
+          "Input(Hyps) must be a 2-D Tensor, but received rank %u.",
+          hyps_dims.size()));
+  PADDLE_ENFORCE_EQ(
+      refs_dims.size(),
+      2,
+      phi::errors::InvalidArgument(
+          "Input(Refs) must be a 2-D Tensor, but received rank %u.",
+          refs_dims.size()));
+
+  infer_context->AddEqualCstr(hyps_dims[0], refs_dims[0]);
+
+  bool has_lengths = op->operand_source(2) && op->operand_source(3);
+  if (has_lengths) {
+    const auto &hypslength_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(2));
+    const auto &refslength_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(3));
+
+    infer_context->AddEqualCstr(hypslength_shape_or_data.shape()[0],
+                                hyps_dims[0]);
+
+    infer_context->AddEqualCstr(refslength_shape_or_data.shape()[0],
+                                refs_dims[0]);
+
+    infer_context->AddEqualCstr(hypslength_shape_or_data.shape()[0],
+                                refslength_shape_or_data.shape()[0]);
+  } else {
+    symbol::DimExpr one = symbol::DimExpr(1);
+
+    infer_context->AddEqualCstr(hyps_dims[1], one);
+    infer_context->AddEqualCstr(refs_dims[1], one);
+  }
+
+  symbol::ShapeOrDataDimExprs refs_shape_or_data_exprs(
+      symbol::TensorShapeOrDataDimExprs(
+          std::vector<symbol::DimExpr>{refs_dims}));
+  infer_context->SetShapeOrDataForValue(op->result(0),
+                                        refs_shape_or_data_exprs);
+
+  symbol::ShapeOrDataDimExprs single_dim_expr(symbol::TensorShapeOrDataDimExprs(
+      std::vector<symbol::DimExpr>{symbol::DimExpr(1)}));
+  infer_context->SetShapeOrDataForValue(
+      op->result(1), symbol::ShapeOrDataDimExprs({single_dim_expr}));
+
+  return true;
+}
+
 bool FakeQuantizeMovingAverageAbsMaxOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const symbol::ShapeOrDataDimExprs &x_shape =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
@@ -43,6 +43,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeQuantizeMovingAverageAbsMax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeQuantizeMovingAverageAbsMax_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(CoalesceTensor)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(CoalesceTensor_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(EditDistance)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FullWithTensor)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttn)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttnUnpadded)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1469,6 +1469,7 @@
     func : edit_distance
     data_type : DataType::FLOAT32
   optional : hypslength, refslength
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : eig
   args: (Tensor x)
