[BUAA][Infer Symbolic Shape]add full_batch_size_like&temporal_shift&squared_l2_norm (#67503)

Wizard-ZP · web-flow · commit 1fcca97d137d · 2024-09-05T19:42:39.000+08:00
* 优化了三个文件

* debug

* debug

* debug

* debug
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -459,6 +459,66 @@ bool AsStridedOpInferSymbolicShape(
   return true;
 }
 
+bool BatchSizeLikeOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::ShapeOrDataDimExprs &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &x_dims = x_shape_or_data.shape();
+
+  std::vector<int> shape =
+      paddle::dialect::details::GetVectorAttr<int>(op, "shape");
+  int x_batch_size_dim =
+      op->attribute<pir::Int32Attribute>("x_batch_size_dim").data();
+  int out_batch_size_dim =
+      op->attribute<pir::Int32Attribute>("out_batch_size_dim").data();
+
+  PADDLE_ENFORCE_GT(
+      shape.size(),
+      0,
+      common::errors::InvalidArgument(
+          "Shape size must be larger than 0, but received: %d.", shape.size()));
+
+  std::vector<symbol::DimExpr> shape_exprs(shape.size(), 0);
+  std::transform(shape.begin(), shape.end(), shape_exprs.begin(), [](int dim) {
+    return symbol::DimExpr(dim);
+  });
+
+  PADDLE_ENFORCE_GE(
+      x_batch_size_dim,
+      0,
+      common::errors::InvalidArgument(
+          "Input dimension index must be larger than or equal to 0."));
+  size_t input_dim_size = x_dims.size();
+  PADDLE_ENFORCE_GE(
+      out_batch_size_dim,
+      0,
+      common::errors::InvalidArgument(
+          "Output dimension index must be larger than or equal to 0."));
+  PADDLE_ENFORCE(static_cast<int>(input_dim_size) > x_batch_size_dim ||
+                     static_cast<int>(input_dim_size) == -1,
+                 common::errors::InvalidArgument(
+                     "Input dimension size must be larger than "
+                     "input dimension index, but received input "
+                     "dimension size: %s, input dimension index: %s.",
+                     static_cast<int>(input_dim_size),
+                     x_batch_size_dim));
+
+  size_t output_dim_size = shape.size();
+  PADDLE_ENFORCE_GT(
+      output_dim_size,
+      out_batch_size_dim,
+      common::errors::InvalidArgument(
+          "Output dimension size must be larger than output dimension index."));
+  shape_exprs[out_batch_size_dim] = x_dims[x_batch_size_dim];
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(shape_exprs)});
+
+  return true;
+}
+
 bool BipartiteMatchOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &dist_mat_shape =
@@ -979,6 +1039,11 @@ bool EigvalshOpInferSymbolicShape(
   return EighOpInferSymbolicShape(op, infer_context);
 }
 
+bool FullBatchSizeLikeOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return BatchSizeLikeOpInferSymbolicShape(op, infer_context);
+}
+
 bool FractionalMaxPoolOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =
@@ -3118,6 +3183,60 @@ bool SetValueWithTensor_OpInferSymbolicShape(
 //   return true;
 // }
 
+bool SquaredL2NormOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  auto dtype = infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> batch_dims;
+  batch_dims.push_back(symbol::DimExpr(1));
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs(
+          symbol::TensorShapeOrDataDimExprs(batch_dims)));
+
+  return true;
+}
+
+bool TemporalShiftOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::ShapeOrDataDimExprs &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &x_dims = x_shape_or_data.shape();
+
+  PADDLE_ENFORCE_EQ(
+      x_dims.size(),
+      4,
+      common::errors::InvalidArgument("Input(X) rank should be 4 in shape of "
+                                      "[N*T, C, H, W], but received X rank(%d)",
+                                      x_dims.size()));
+
+  int seg_num = op->attribute<pir::Int32Attribute>("seg_num").data();
+  float shift_ratio = op->attribute<pir::FloatAttribute>("shift_ratio").data();
+
+  PADDLE_ENFORCE_GT(
+      seg_num,
+      0,
+      common::errors::InvalidArgument(
+          "Attr(seg_num) should be greater than 0, but received %d", seg_num));
+  PADDLE_ENFORCE_GT(
+      shift_ratio,
+      0.0f,
+      common::errors::InvalidArgument(
+          "Attr(shift_ratio) should be greater than 0, but received %f",
+          shift_ratio));
+  PADDLE_ENFORCE_LT(
+      shift_ratio,
+      0.5f,
+      common::errors::InvalidArgument(
+          "Attr(shift_ratio) should be less than 0.5, but received %f",
+          shift_ratio));
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs(symbol::TensorShapeOrDataDimExprs(x_dims)));
+
+  return true;
+}
+
 bool TileOpInferSymbolicShape(pir::Operation *op,
                               pir::InferSymbolicShapeContext *infer_context) {
   pir::Value operand_x = op->operand_source(0);
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -33,6 +33,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(AsStrided)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(AllReduce_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Barrier)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BipartiteMatch)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(BatchSizeLike)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cast)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cast_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cholesky)
@@ -72,6 +73,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flatten)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flatten_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Fold)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeQuantizeDequantizeAbsMax)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FullBatchSizeLike)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FrobeniusNorm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Inverse)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GumbelSoftmax)
@@ -126,6 +128,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(ShuffleChannel)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Slice)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Split)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SplitWithNum)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(SquaredL2Norm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Squeeze)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Squeeze_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(StridedSlice)
@@ -137,6 +140,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(SetValueWithTensor)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SetValueWithTensor_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(TensorUnfold)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Trace)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(TemporalShift)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tile)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Topk)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TopkV1)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -2075,6 +2075,7 @@
     param : [input, shape, value, dtype, input_dim_idx, output_dim_idx]
     data_type : dtype
     backend : place
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : full_int_array
   args : (int64_t[] value, DataType dtype=DataType::FLOAT32, Place place=CPUPlace())
@@ -4530,6 +4531,7 @@
   kernel :
     func : squared_l2_norm
   backward : squared_l2_norm_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : squeeze
   args : (Tensor x, IntArray axis={})
@@ -4739,6 +4741,7 @@
     func : temporal_shift
     data_type : x
   backward : temporal_shift_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : tensor_unfold
   args : (Tensor input, int64_t axis, int64_t size, int64_t step)
