【Infer Symbolic Shape No.204】【BUAA】martix_nms (#67607)

tlxd · web-flow · commit 1f8e5c1509b2 · 2024-08-30T10:32:04.000+08:00
* Update binary_infer_sym.cc

* Update binary_infer_sym.cc

* Update binary_infer_sym.cc

* Update ops.yaml binary_infer_sym.cc  binary_infer_sym.h

* Update binary_infer_sym.cc

* Update binary_infer_sym.cc

* Update binary_infer_sym.cc binary.cc test_matrix_nms_op.py
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
@@ -1174,12 +1174,46 @@ bool MatmulOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
-// bool MatrixNmsOpInferSymbolicShape(pir::Operation *op,
-//                                    pir::InferSymbolicShapeContext
-//                                    *infer_context) {
-//   // pass
-//   return true;
-// }
+bool MatrixNmsOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &bboxes_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &scores_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const std::vector<symbol::DimExpr> &bboxes_shape =
+      bboxes_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &scores_shape =
+      scores_shape_or_data.shape();
+
+  infer_context->AddEqualCstr(bboxes_shape[2], symbol::DimExpr(4));
+  infer_context->AddEqualCstr(bboxes_shape[1], scores_shape[2]);
+
+  symbol::DimExpr num_kept = infer_context->GetNextSymName();
+
+  std::vector<symbol::DimExpr> out_shape = {
+      num_kept, bboxes_shape[2] + symbol::DimExpr(2)};
+  std::vector<symbol::DimExpr> index_shape = {num_kept, symbol::DimExpr(1)};
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(index_shape)});
+
+  if (op->result(2)) {
+    std::vector<symbol::DimExpr> roisnum_shape = {scores_shape[0]};
+    infer_context->SetShapeOrDataForValue(
+        op->result(2),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(roisnum_shape)});
+  }
+
+  return true;
+}
 
 bool MarginCrossEntropyOpInferSymbolicShape(
     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
@@ -68,7 +68,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(LuUnpack)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixRankTol)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedSelect)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Matmul)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixNms)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixNms)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MarginCrossEntropy)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatmulWithFlatten)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mv)
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -3034,12 +3034,12 @@ void MatrixNMSInferMeta(const MetaTensor& bboxes,
             box_dims[1],
             score_dims[2]));
   }
-  out->set_dims({box_dims[1], box_dims[2] + 2});
+  out->set_dims({-1, box_dims[2] + 2});
   out->set_dtype(bboxes.dtype());
-  index->set_dims({box_dims[1], 1});
+  index->set_dims({-1, 1});
   index->set_dtype(phi::DataType::INT32);
   if (roisnum != nullptr) {
-    roisnum->set_dims({-1});
+    roisnum->set_dims({score_dims[0]});
     roisnum->set_dtype(phi::DataType::INT32);
   }
 }
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3114,7 +3114,7 @@
   optional : roisnum
   kernel :
     func : matrix_nms
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : matrix_power
   args : (Tensor x, int n)
diff --git a/test/legacy_test/test_matrix_nms_op.py b/test/legacy_test/test_matrix_nms_op.py
@@ -296,7 +296,7 @@ def setUp(self):
         }
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_symbol_infer=False)
 
 
 class TestMatrixNMSOpNoOutput(TestMatrixNMSOp):

Original file line number	Diff line number	Diff line change
`@@ -3034,12 +3034,12 @@ void MatrixNMSInferMeta(const MetaTensor& bboxes,`
`3034`	`3034`	`box_dims[1],`
`3035`	`3035`	`score_dims[2]));`
`3036`	`3036`	`}`
`3037`		`- out->set_dims({box_dims[1], box_dims[2] + 2});`
	`3037`	`+ out->set_dims({-1, box_dims[2] + 2});`
`3038`	`3038`	`out->set_dtype(bboxes.dtype());`
`3039`		`- index->set_dims({box_dims[1], 1});`
	`3039`	`+ index->set_dims({-1, 1});`
`3040`	`3040`	`index->set_dtype(phi::DataType::INT32);`
`3041`	`3041`	`if (roisnum != nullptr) {`
`3042`		`- roisnum->set_dims({-1});`
	`3042`	`+ roisnum->set_dims({score_dims[0]});`
`3043`	`3043`	`roisnum->set_dtype(phi::DataType::INT32);`
`3044`	`3044`	`}`
`3045`	`3045`	`}`
Original file line number	Diff line number	Diff line change
`@@ -296,7 +296,7 @@ def setUp(self):`
`296`	`296`	`}`
`297`	`297`
`298`	`298`	`def test_check_output(self):`
`299`		`- self.check_output(check_pir=True)`
	`299`	`+ self.check_output(check_pir=True, check_symbol_infer=False)`
`300`	`300`
`301`	`301`
`302`	`302`	`class TestMatrixNMSOpNoOutput(TestMatrixNMSOp):`