【Infer Symbolic Shape No.232】Add infer_symbol_shape for StridedSlice (#69911)

AndPuQing · web-flow · commit d72d6ad1e228 · 2024-12-04T20:38:09.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_slice_utils.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_slice_utils.h
@@ -115,6 +115,7 @@ inline ExprVec GetSliceDims(const ExprVec &in_dims,
   for (size_t i = 0; i < axes.size(); ++i) {
     auto out_dim = ends[i] - starts[i];
     int64_t axis = axes[i];
+
     // If in_dims[axis] or ends[i] have symbol, nedd get Min(in_dims[axis] -
     // start[i], ends[i] - start[i] )
     if (!out_dim.isa<int64_t>() &&
@@ -266,4 +267,219 @@ inline ShapeOrData SliceRawInferSymbolicShape(
 
   return out_shape;
 }
+
+inline ExprVec GetStridesSliceDims(
+    const ExprVec &in_dims,
+    const std::vector<int64_t> &axes,
+    const ExprVec &starts_base,
+    const ExprVec &ends_base,
+    const ExprVec &strides_base,
+    std::vector<int64_t> *infer_flags = nullptr) {
+  ExprVec starts = starts_base;
+  ExprVec ends = ends_base;
+  ExprVec strides = strides_base;
+  auto IsMaxInt = [](const symbol::DimExpr &expr) {
+    return expr.isa<int64_t>() &&
+           expr.Get<int64_t>() ==
+               static_cast<int64_t>(std::numeric_limits<int>::max());
+  };
+
+  for (size_t i = 0; i < axes.size(); ++i) {
+    int64_t axis = axes.at(i);
+    int64_t start_i = 0;
+
+    if (starts.at(i).isa<int64_t>()) {
+      if (in_dims.at(axis).isa<int64_t>()) {
+        starts.at(i) =
+            (starts.at(i).Get<int64_t>() > in_dims.at(axis).Get<int64_t>())
+                ? in_dims.at(axis)
+                : starts.at(i);
+        starts.at(i) =
+            (starts.at(i).Get<int64_t>() < -in_dims.at(axis).Get<int64_t>())
+                ? symbol::DimExpr({-1}) * in_dims.at(axis)
+                : starts.at(i);
+      }
+      start_i = starts.at(i).Get<int64_t>();
+    }
+
+    int64_t end_i = 0;
+    if (ends.at(i).isa<int64_t>()) {
+      if (in_dims.at(axis).isa<int64_t>()) {
+        ends[i] = std::min(ends.at(i).Get<int64_t>(),
+                           in_dims.at(axis).Get<int64_t>());
+      }
+      if (ends.at(i).Get<int64_t>() < 0) {
+        ends[i] = ends.at(i) + in_dims.at(axis);
+      }
+      if (ends.at(i).isa<int64_t>()) {
+        end_i = ends.at(i).Get<int64_t>();
+      }
+    }
+
+    ends.at(i) = IsMaxInt(ends.at(i)) ? in_dims.at(axis) : ends.at(i);
+    bool both_negative_or_positive =
+        (start_i >= 0 && end_i >= 0) || (start_i <= 0 && end_i <= 0);
+    bool start_negative_end_positive = start_i <= 0 && end_i >= 0;
+    bool start_positive_end_negative = start_i >= 0 && end_i <= 0;
+
+    if (both_negative_or_positive) {
+      continue;
+    } else if (start_negative_end_positive) {
+      starts.at(i) = starts.at(i) + in_dims.at(axis);
+    } else if (start_positive_end_negative) {
+      starts.at(i) = starts.at(i) - in_dims.at(axis);
+    } else {
+      PADDLE_THROW(common::errors::Fatal("Dead code"));
+    }
+  }
+
+  ExprVec slice_dims(in_dims);
+  PADDLE_ENFORCE_EQ(
+      (axes.size() == starts.size() && axes.size() == ends.size() &&
+       axes.size() == strides.size()),
+      true,
+      common::errors::InvalidArgument(
+          "The size of axes must equal size of starts, ends, and strides."));
+
+  for (size_t i = 0; i < axes.size(); ++i) {
+    auto out_dim = symbol::DimExpr({-1}) * ((starts[i] - ends[i]) / strides[i]);
+    int64_t axis = axes[i];
+
+    if (!out_dim.isa<int64_t>() &&
+        (!in_dims[axis].isa<int64_t>() || !ends[i].isa<int64_t>())) {
+      symbol::List<symbol::DimExpr> min_lists{
+          symbol::DimExpr({-1}) * ((starts[i] - in_dims[axis]) / strides[i]),
+          out_dim};
+
+      slice_dims[axis] =
+          symbol::DimExpr({symbol::Min<symbol::DimExpr>({min_lists})});
+    } else {
+      slice_dims[axis] = out_dim;
+    }
+  }
+
+  return slice_dims;
+}
+
+inline ShapeOrData StridedSliceRawInferSymbolicShape(
+    const pir::Value x,
+    const pir::Value out,
+    const ExprVec &starts_expr,
+    const ExprVec &ends_expr,
+    const ExprVec &strides_expr,
+    const std::vector<int64_t> &axes_raw,
+    const std::vector<int64_t> &infer_flags_raw,
+    const std::vector<int64_t> &decrease_axis,
+    pir::InferSymbolicShapeContext *infer_context) {
+  const auto &in_shapeordata = infer_context->GetShapeOrDataForValue(x);
+  ExprVec starts = starts_expr;
+  ExprVec ends = ends_expr;
+  ExprVec strides = strides_expr;
+  std::vector<int64_t> infer_flags = [&infer_flags_raw, &axes_raw] {
+    return infer_flags_raw.empty() ? std::vector<int64_t>(axes_raw.size(), 1)
+                                   : infer_flags_raw;
+  }();
+
+  const auto &GetShapeDimExprs = [&]() -> symbol::ShapeOrDataDimExprs {
+    const ExprVec &in_dims = in_shapeordata.shape();
+    std::vector<int64_t> axes = FormatSliceAxes(axes_raw, in_dims.size());
+    ExprVec slice_dims =
+        GetStridesSliceDims(in_dims, axes, starts, ends, strides, &infer_flags);
+    ExprVec out_dims = GetDecreasedDims(slice_dims, decrease_axis);
+
+    auto IsOne = [](const symbol::DimExpr &expr) {
+      return expr.isa<int64_t>() && expr.dyn_cast<int64_t>() == 1;
+    };
+    auto IsIntType = [](pir::Value value) {
+      const auto &dtype = value.type().dyn_cast<pir::DenseTensorType>().dtype();
+      return dtype.isa<pir::Int32Type>() || dtype.isa<pir::Int64Type>();
+    };
+    if (IsIntType(x) &&
+        (out_dims.empty() || (out_dims.size() == 1 && IsOne(out_dims[0])))) {
+      return symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+          out_dims,
+          std::vector<symbol::DimExpr>{infer_context->GetNextSymName()})};
+    }
+
+    return symbol::ShapeOrDataDimExprs{
+        symbol::TensorShapeOrDataDimExprs(out_dims)};
+  };
+
+  // When `pd.slice` is operating on a tensor which is produced by a `pd.shape`
+  // op, the result should be written into data.
+  const auto &GetDataDimExprs = [&]() -> symbol::ShapeOrDataDimExprs {
+    std::vector<symbol::DimExpr> out_data;
+
+    // Currently, we DO NOT support the case that any element in `axes` `starts`
+    // or `ends` is a Symbol.
+    auto vec_int64 = details::VecExpr2Int64(starts);
+    PADDLE_ENFORCE_EQ(
+        vec_int64.has_value(),
+        true,
+        common::errors::InvalidArgument(
+            "for slice op, all the elements in `starts` must be int64_t"));
+    std::vector<int64_t> starts_int = vec_int64.value();
+
+    vec_int64 = details::VecExpr2Int64(ends);
+    PADDLE_ENFORCE_EQ(
+        vec_int64.has_value(),
+        true,
+        common::errors::InvalidArgument(
+            "for slice op, all the elements in `ends` must be int64_t"));
+    std::vector<int64_t> ends_int = vec_int64.value();
+
+    vec_int64 = details::VecExpr2Int64(strides);
+    PADDLE_ENFORCE_EQ(
+        vec_int64.has_value(),
+        true,
+        common::errors::InvalidArgument(
+            "for slice op, all the elements in `strides` must be int64_t"));
+
+    const int64_t start =
+        starts_int[0] < 0 ? starts_int[0] + in_shapeordata.data().value().size()
+                          : starts_int[0];
+    const int64_t end = [&]() -> int64_t {
+      if (ends_int[0] < 0) {
+        return ends_int[0] + in_shapeordata.data().value().size();
+      }
+      if (ends_int[0] ==
+          static_cast<int64_t>(std::numeric_limits<int>::max())) {
+        return in_shapeordata.data().value().size();
+      }
+      return ends_int[0];
+    }();
+
+    const int64_t stride = [&]() -> int64_t {
+      if (strides[0].isa<int64_t>()) {
+        return strides[0].Get<int64_t>();
+      }
+      return 1;
+    }();
+
+    for (int64_t i = start; i < end; i += stride) {
+      out_data.push_back(in_shapeordata.data().value().at(i));
+    }
+
+    const ExprVec shape = GetDecreasedDims(
+        ExprVec{static_cast<int64_t>(out_data.size())}, decrease_axis);
+    return symbol::ShapeOrDataDimExprs{
+        symbol::TensorShapeOrDataDimExprs(shape, out_data)};
+  };
+
+  const auto &out_shape = in_shapeordata.data().has_value()
+                              ? GetDataDimExprs()
+                              : GetShapeDimExprs();
+  if (out_shape.data().has_value() && out_shape.shape().empty()) {  // 0D tensor
+    const paddle::dialect::DenseTensorType &tensor_type =
+        out.type().dyn_cast<paddle::dialect::DenseTensorType>();
+    const auto &out_ddim = tensor_type.dims();
+    if (out_ddim.size() == 1 && out_ddim[0] == 1) {  // value is 1D
+      return symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+          std::vector<symbol::DimExpr>{1}, out_shape.data().value())};
+    }
+  }
+
+  return out_shape;
+}
+
 }  // namespace paddle::dialect::slice_utils
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
@@ -3515,12 +3515,42 @@ bool SplitWithNumOpInferSymbolicShape(
   return true;
 }
 
-// bool StridedSliceOpInferSymbolicShape(pir::Operation *op,
-//                                       pir::InferSymbolicShapeContext
-//                                       *infer_context) {
-//   // pass
-//   return true;
-// }
+bool StridedSliceOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  pir::Value operand_source = op->operand_source(0);
+  pir::Value operand_starts = op->operand_source(1);
+  pir::Value operand_ends = op->operand_source(2);
+  pir::Value operand_strides = op->operand_source(3);
+  pir::Value res = op->result(0);
+
+  const symbol::ShapeOrDataDimExprs &starts_shape_data =
+      infer_context->GetShapeOrDataForValue(operand_starts);
+  const symbol::ShapeOrDataDimExprs &ends_shape_data =
+      infer_context->GetShapeOrDataForValue(operand_ends);
+  const symbol::ShapeOrDataDimExprs &strides_shape_data =
+      infer_context->GetShapeOrDataForValue(operand_strides);
+
+  ExprVec starts = slice_utils::GetExprVecFromData(starts_shape_data);
+  ExprVec ends = slice_utils::GetExprVecFromData(ends_shape_data);
+  ExprVec strides = slice_utils::GetExprVecFromData(strides_shape_data);
+
+  std::vector<int32_t> axes_vec = details::GetVectorAttr<int32_t>(op, "axes");
+  std::vector<int64_t> axes_vec_64(axes_vec.begin(), axes_vec.end());
+
+  infer_context->SetShapeOrDataForValue(
+      res,
+      slice_utils::StridedSliceRawInferSymbolicShape(operand_source,
+                                                     res,
+                                                     starts,
+                                                     ends,
+                                                     strides,
+                                                     axes_vec_64,
+                                                     std::vector<int64_t>{},
+                                                     std::vector<int64_t>{},
+                                                     infer_context));
+
+  return true;
+}
 
 bool SumOpInferSymbolicShape(pir::Operation *op,
                              pir::InferSymbolicShapeContext *infer_context) {
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
@@ -139,7 +139,7 @@ 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)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(StridedSlice)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sum)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Svd)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SetValue)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4800,7 +4800,7 @@
   kernel :
     func : strided_slice
   backward : strided_slice_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : sum
   args : (Tensor x, IntArray axis={}, DataType dtype=DataType::UNDEFINED, bool keepdim=false)
diff --git a/python/paddle/base/variable_index.py b/python/paddle/base/variable_index.py
@@ -765,7 +765,16 @@ def get_tensor_with_basic_indexing(
                 stride = attrs['strides']
             if use_strided_slice:
                 # TODO(zoooo0820): support strided_slice_array until PIR API is ready
-
+                if in_pir_mode():
+                    if isinstance(st, (list, tuple)):
+                        if paddle.utils._contain_var(st):
+                            st = paddle.utils.get_int_tensor_list(st)
+                    if isinstance(end, (list, tuple)):
+                        if paddle.utils._contain_var(end):
+                            end = paddle.utils.get_int_tensor_list(end)
+                    if isinstance(stride, (list, tuple)):
+                        if paddle.utils._contain_var(stride):
+                            stride = paddle.utils.get_int_tensor_list(stride)
                 out = paddle._C_ops.strided_slice(x, axes, st, end, stride)
                 if len(decrease_axes) > 0:
                     out = paddle._C_ops.squeeze(out, decrease_axes)
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -5638,7 +5638,22 @@ def strided_slice(
             >>> sliced_2 = paddle.strided_slice(x, axes=axes, starts=[minus_3, 0, 2], ends=ends, strides=strides_2)
             >>> # sliced_2 is x[:, 1:3:1, 0:2:1, 2:4:2].
     """
-    if in_dynamic_or_pir_mode():
+    if in_dynamic_mode():
+        return _C_ops.strided_slice(x, axes, starts, ends, strides)
+    elif in_pir_mode():
+
+        def _convert_to_tensor_list(input):
+            if isinstance(input, paddle.pir.Value):
+                input.stop_gradient = True
+            elif isinstance(input, (list, tuple)):
+                if paddle.utils._contain_var(input):
+                    input = paddle.utils.get_int_tensor_list(input)
+            return input
+
+        starts = _convert_to_tensor_list(starts)
+        ends = _convert_to_tensor_list(ends)
+        strides = _convert_to_tensor_list(strides)
+
         return _C_ops.strided_slice(x, axes, starts, ends, strides)
     else:
         helper = LayerHelper('strided_slice', **locals())
diff --git a/test/legacy_test/test_strided_slice_op.py b/test/legacy_test/test_strided_slice_op.py
