【Infer Symbolic Shape No.124】【BUAA】Add bmm, changed 3 files (#67730)

Whsjrczr · web-flow · commit a3cfda8d5d39 · 2024-08-27T17:00:33.000+08:00
* add 3 files

* update .cc

* undo change in .cc file

* update num_samples in .cc

* attribute

* update

* update

* update

* update

* update

* update

* update

* updata

* updata

* update

* update

* update

* update

* update

* update

* u

* batch_function

* bincount

* update batchfc

* update batchfc

* update EQ

* update {-1}

* update binary with output_size

* undo change

* {-1}

* add Bincount

* update class_center_sample

* delete batch_fc

* update batchnormop

* update batchnormop

* changed bn

* changed bn

* unduo _

* Update DimExpr

* Update unary_infer_sym.cc

* out_unknown

* Expr -&gt; Exprs

* {{out_dims}} -&gt; out_dims

* update dimexpr

* restore changes

* undo

* undo some changes

* add bmm
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
@@ -271,11 +271,47 @@ bool BincountOpInferSymbolicShape(
   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
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
@@ -27,7 +27,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)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -652,6 +652,7 @@
   kernel :
     func : bmm
   backward : bmm_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : box_clip
   args: (Tensor input, Tensor im_info)
