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

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

@@ -238,11 +238,47 @@ 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_shapes = x_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &y_shapes = y_shape_or_data.shape();
+  std::size_t x_size = x_shapes.size();
+  std::size_t y_size = y_shapes.size();
+
+  PADDLE_ENFORCE_EQ(
+      x_size,
+      3,
+      common::errors::InvalidArgument("Input(X) of BmmOp must be 3-dimensional "
+                                      "in BmmOp, but received X's shape: [%d].",
+                                      x_size));
+  PADDLE_ENFORCE_EQ(
+      y_size,
+      3,
+      common::errors::InvalidArgument("Input(Y) of BmmOp must be 3-dimensional "
+                                      "in BmmOp, but received Y's shape: [%d].",
+                                      y_size));
+
+  infer_context->AddEqualCstr(x_shapes[2], y_shapes[1]);
+  infer_context->AddEqualCstr(x_shapes[0], y_shapes[0]);
+
+  const symbol::DimExpr batch_size = x_shapes[0];
+
+  const symbol::DimExpr out_height = x_shapes[1];
+  const symbol::DimExpr out_width = y_shapes[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,14 +26,18 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(BceLoss_)
 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(Bincount)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bmm)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(CholeskySolve)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(CtcAlign)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv2d)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv2dTranspose)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv3d)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(ConvTranspose)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cross)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dot)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Correlation)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(DepthwiseConv)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dot)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dropout)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Embedding)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(EqualAll)
@@ -45,32 +49,49 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedSoftmaxMask)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(GridSample)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gather)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GatherNd)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(GatherTree)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(HuberLoss)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Histogram)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isclose)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexAdd)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexAdd_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexPut)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexPut_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSelect)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(AccuracyCheck)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSample)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSelectStrided)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(KldivLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lstsq)
+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(MarginCrossEntropy)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatmulWithFlatten)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mv)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(PriorBox)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullBoxSparse)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullGpuPsSparse)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullSparseV2)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(RepeatInterleaveWithTensorIndex)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReduceAs)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Searchsorted)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SegmentPool)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(SequenceMask)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(SequenceMask)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShuffleBatch)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Solve)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Stft)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Swiglu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TakeAlongAxis)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TopPSampling)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TdmChild)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(TriangularSolve)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unpool3d)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unpool)
+// OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightDequantize)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBox)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBoxHead)
 

paddle/phi/ops/yaml/ops.yaml

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