【BUAA】【Infer Symbolic Shape】add mv, shadow_feed, share_data_ (#66956)

uanu2002 · web-flow · commit c6de16c62236 · 2024-08-07T10:18:35.000+08:00
* add mv, shadow_feed, share_data_

* fix sharedata

* fix codestyle

* fix
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
@@ -720,6 +720,33 @@ bool MatmulOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool MvOpInferSymbolicShape(pir::Operation *op,
+                            pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &vec_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  PADDLE_ENFORCE_EQ(
+      x_shape_or_data.shape().size(),
+      2,
+      phi::errors::InvalidArgument("The rank of input X should be 2, but is %d",
+                                   x_shape_or_data.shape().size()));
+  PADDLE_ENFORCE_EQ(vec_shape_or_data.shape().size(),
+                    1,
+                    phi::errors::InvalidArgument(
+                        "The rank of input Vec should be 1, but is %d",
+                        vec_shape_or_data.shape().size()));
+  infer_context->AddEqualCstr(x_shape_or_data.shape()[1],
+                              vec_shape_or_data.shape()[0]);
+
+  std::vector<symbol::DimExpr> out_shape = {x_shape_or_data.shape()[0]};
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+  return true;
+}
+
 // bool PullBoxSparseOpInferSymbolicShape(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
@@ -51,6 +51,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixRankTol)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedSelect)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Matmul)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mv)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Nextafter)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullBoxSparse)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullGpuPsSparse)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc
@@ -140,6 +140,8 @@ OP_SAME_OPERANDS_AND_RESULT(ScatterNdAdd)
 OP_SAME_OPERANDS_AND_RESULT(Scatter)
 OP_SAME_OPERANDS_AND_RESULT(Scatter_)
 OP_SAME_OPERANDS_AND_RESULT(Select)
+OP_SAME_OPERANDS_AND_RESULT(ShadowFeed)
+OP_SAME_OPERANDS_AND_RESULT(ShareData_)
 OP_SAME_OPERANDS_AND_RESULT(Sign)
 OP_SAME_OPERANDS_AND_RESULT(Sin)
 OP_SAME_OPERANDS_AND_RESULT(Sin_)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h
@@ -133,6 +133,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(ScatterNdAdd)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Scatter)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Scatter_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Select)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShadowFeed)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShareData_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sign)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sin)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sin_)
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -810,6 +810,7 @@
   kernel:
     func: shadow_feed
     param: [x]
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : shadow_feed_tensors
   args : (Tensor[] x)
@@ -832,6 +833,7 @@
     param: [x]
   inplace : (x -> out)
   traits : paddle::dialect::ForwardOnlyTrait
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : soft_relu
   args : (Tensor x, float threshold = 20.0f)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3286,6 +3286,7 @@
   kernel :
     func : mv
   backward : mv_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : nadam_
   args : (Tensor param, Tensor grad, Tensor learning_rate, Tensor momentum_decay_pow, Tensor beta2_pow, Tensor mu_product, Tensor moment1, Tensor moment2, Tensor master_param, float beta1 = 0.9f, float beta2 = 0.999f, float epsilon = 1.0e-8f, float momentum_decay = 0.004f, bool multi_precision = false)
