[CINN] 【Infer Symbolic Shape BUAA 】Add match_matrix_tensor op (#68177)

crazyxiaoxi · Jeff114514 · web-flow · commit f2381590e573 · 2024-09-18T11:22:56.000+08:00
* first

* fix name

---------

Co-authored-by: Jeff114514 &lt;928430299@qq.com&gt;
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
@@ -2187,6 +2187,58 @@ bool NearestInterpOpInferSymbolicShape(
   return BicubicInterpOpInferSymbolicShape(op, infer_context);
 }
 
+bool MatchMatrixTensorOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &y_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const auto &w_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(2));
+
+  const auto &x_shape = x_shape_or_data.shape();
+  const auto &y_shape = y_shape_or_data.shape();
+  const auto &w_shape = w_shape_or_data.shape();
+
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    2,
+                    common::errors::InvalidArgument(
+                        "The dimensions of Input(X) should be equal to 2, "
+                        "but received %d.",
+                        x_shape.size()));
+  PADDLE_ENFORCE_EQ(y_shape.size(),
+                    2,
+                    common::errors::InvalidArgument(
+                        "The dimensions of Input(Y) should be equal to 2, "
+                        "but received %d.",
+                        y_shape.size()));
+  PADDLE_ENFORCE_EQ(w_shape.size(),
+                    3,
+                    common::errors::InvalidArgument(
+                        "The dimensions of Input(W) should be equal to 3, "
+                        "but received %d.",
+                        w_shape.size()));
+
+  infer_context->AddEqualCstr(w_shape[0], x_shape[1]);
+  int dim_t = op->attribute<pir::Int32Attribute>("dim_t").data();
+  infer_context->AddEqualCstr(w_shape[1], symbol::DimExpr(dim_t));
+  infer_context->AddEqualCstr(w_shape[2], y_shape[1]);
+
+  std::vector<symbol::DimExpr> out_dims = {infer_context->GetNextSymName(),
+                                           symbol::DimExpr(1)};
+  std::vector<symbol::DimExpr> tmp_dims = {infer_context->GetNextSymName(),
+                                           symbol::DimExpr(1)};
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(out_dims)});
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(tmp_dims)});
+
+  return true;
+}
+
 bool MaskedMultiheadAttention_OpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const symbol::ShapeOrDataDimExprs &x_shape_or_data =
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
@@ -78,6 +78,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Linspace)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(LinearInterp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lstm)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatchMatrixTensor)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedMultiheadAttention_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedAdam)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedAdam_)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3119,7 +3119,7 @@
   kernel:
     func: match_matrix_tensor
   backward: match_matrix_tensor_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : matrix_nms
   args : (Tensor bboxes, Tensor scores, float score_threshold, int nms_top_k, int keep_top_k, float post_threshold=0., bool use_gaussian = false, float gaussian_sigma = 2., int background_label = 0, bool normalized = true)
