【Infer Symbolic Shape No.124】【BUAA】Add bmm, changed 3 files

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc

@@ -207,11 +207,63 @@ bool Binomial_OpInferSymbolicShape(
 //   return true;
 // }
 
-// bool BmmOpInferSymbolicShape(pir::Operation *op,
-//                              pir::InferSymbolicShapeContext *infer_context) {
-//   // pass
-//   return true;
-// }
+bool BmmOpInferSymbolicShape(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 std::vector<symbol::DimExpr> &x_dims = x_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &y_dims = y_shape_or_data.shape();
+  std::size_t x_ndims = x_dims.size();
+  std::size_t y_ndims = y_dims.size();
+
+  PADDLE_ENFORCE_EQ(
+      x_ndims,
+      3,
+      common::errors::InvalidArgument("Input(X) of BmmOp must be 3-dimensional "
+                                      "in BmmOp, but received X's shape: [%d].",
+                                      x_ndims));
+  PADDLE_ENFORCE_EQ(
+      y_ndims,
+      3,
+      common::errors::InvalidArgument("Input(Y) of BmmOp must be 3-dimensional "
+                                      "in BmmOp, but received Y's shape: [%d].",
+                                      y_ndims));
+
+  auto cal_shape_fn = [](const symbol::DimExpr &x,
+                         const symbol::DimExpr &y,
+                         const std::string &error_str) -> symbol::DimExpr {
+    if (x == -1) {
+      return y;
+    } else if (y == -1) {
+      return x;
+    }
+    PADDLE_ENFORCE_EQ(x, y, common::errors::InvalidArgument(error_str, x, y));
+    return x;
+  };
+
+  cal_shape_fn(x_dims[2],
+               y_dims[1],
+               "Input(X)'s width must be equal with Input(Y)'s height in "
+               "BmmOp, but receive X's width: [%d], Y's height: [%d].");
+  symbol::DimExpr batch_size = cal_shape_fn(
+      x_dims[0],
+      y_dims[0],
+      "Input(X) and Input(Y) must have the same batch size in BmmOp, but "
+      "received X's batch size: [%d], Y's batch size [%d]");
+  symbol::DimExpr out_height = x_dims[1];
+  symbol::DimExpr out_width = y_dims[2];
+
+  std::vector<symbol::DimExpr> out_dims = {batch_size, out_height, out_width};
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(out_dims)});
+
+  return true;
+}
 
 // bool CholeskySolveOpInferSymbolicShape(pir::Operation *op,
 //                                        pir::InferSymbolicShapeContext

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h

@@ -26,7 +26,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(BoxClip)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Binomial)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Binomial_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bincount)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bmm)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bmm)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(CholeskySolve)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(CtcAlign)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv2d)

paddle/phi/ops/yaml/ops.yaml

@@ -644,6 +644,7 @@
   kernel :
     func : bmm
   backward : bmm_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : box_clip
   args: (Tensor input, Tensor im_info)