[CINN] Add InferSymbolicShapeInterface for pd_op.strided_slice (#70541)

ooooo-create · AndPuQing · web-flow · commit a1f76a65d4ec · 2025-03-03T14:46:44.000+08:00
* Enhance strided_slice operation with infer_flags and decrease_axis parameters * fix * fix * fix * fix * fix * fix * Revert "fix" This reverts commit 6374d3c. * fix * fix * fix bugs * fix bugs * fix bugs * add tests * refine code * refine code * refine code * refine code * refine code * fix bugs --------- Co-authored-by: PuQing <me@puqing.work>
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
@@ -15,7 +15,7 @@
 #pragma once
 
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h"
-
+#include "paddle/phi/kernels/funcs/slice_utils.h"
 namespace paddle::dialect::slice_utils {
 
 inline bool GetExprVecOfStartEnd(
@@ -301,4 +301,263 @@ inline ShapeOrData SliceRawInferSymbolicShape(
 
   return out_shape;
 }
+inline ExprVec GetStridedSliceDims(
+    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,
+    pir::InferSymbolicShapeContext *infer_context) {
+  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());
+  };
+
+  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) {
+    int64_t axis = axes.at(i);
+    if (in_dims.at(i).isa<int64_t>() && starts.at(i).isa<int64_t>() &&
+        ends.at(i).isa<int64_t>() && strides.at(i).isa<int64_t>()) {
+      int64_t in_dim = in_dims[axis].Get<int64_t>();
+      int64_t start = starts[i].Get<int64_t>();
+      int64_t end = ends[i].Get<int64_t>();
+      int64_t stride = strides[i].Get<int64_t>();
+      bool dummy_zero_dim_out = false;
+      phi::funcs::normalize_interval(
+          start, end, stride, in_dim, &start, &end, &dummy_zero_dim_out);
+      if (end == -in_dim - 1) {
+        end = -1;
+      }
+      int64_t step_size = std::abs(stride);
+      auto out_dim = (std::abs(end - start) + step_size - 1) / step_size;
+      slice_dims[axis] = symbol::DimExpr({out_dim});
+    } else {
+      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) {
+        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.This code should never be reached due to logical."));
+        }
+      }
+
+      symbol::DimExpr out_dim;
+      int64_t stride_int64 = 0;
+      if (strides[i].isa<int64_t>()) {
+        stride_int64 = strides[i].Get<int64_t>();
+        if (stride_int64 > 0) {
+          symbol::List<symbol::DimExpr> unnegativate_lists{
+              (ends[i] - starts[i] - 1 + stride_int64) / stride_int64, 0};
+          out_dim = symbol::DimExpr(
+              {symbol::Max<symbol::DimExpr>({unnegativate_lists})});
+        } else {
+          symbol::List<symbol::DimExpr> unnegativate_lists{
+              (ends[i] - starts[i] + 1 + stride_int64) / stride_int64, 0};
+          out_dim = symbol::DimExpr(
+              {symbol::Max<symbol::DimExpr>({unnegativate_lists})});
+        }
+      } else {
+        out_dim = infer_context->GetNextSymName();
+      }
+      if (!out_dim.isa<int64_t>() &&
+          (!in_dims[axis].isa<int64_t>() || !ends[i].isa<int64_t>())) {
+        if (strides[i].isa<int64_t>()) {
+          if (stride_int64 > 0) {
+            symbol::List<symbol::DimExpr> min_lists{
+                (in_dims[axis] - starts[i] - 1 + stride_int64) / stride_int64,
+                out_dim};
+
+            slice_dims[axis] =
+                symbol::DimExpr({symbol::Min<symbol::DimExpr>({min_lists})});
+          } else {
+            symbol::List<symbol::DimExpr> min_lists{
+                (in_dims[axis] - starts[i] + 1 + stride_int64) / stride_int64,
+                out_dim};
+
+            slice_dims[axis] =
+                symbol::DimExpr({symbol::Min<symbol::DimExpr>({min_lists})});
+          }
+        } else {
+          slice_dims[axis] = out_dim;
+        }
+      } 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 ExprVec &in_dims = in_shapeordata.shape();
+  std::vector<int64_t> axes = FormatSliceAxes(axes_raw, in_dims.size());
+
+  const auto &GetShapeDimExprs = [&]() -> symbol::ShapeOrDataDimExprs {
+    ExprVec slice_dims = GetStridedSliceDims(
+        in_dims, axes, starts, ends, strides, &infer_flags, infer_context);
+    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.strided_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 {
+    PADDLE_ENFORCE_EQ(in_dims.size(),
+                      1,
+                      common::errors::InvalidArgument(
+                          "Currently for strided_slice op, only the rank of "
+                          "shape == 1 is supported."));
+    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);
+    std::vector<int64_t> starts_int = vec_int64.value();
+
+    vec_int64 = details::VecExpr2Int64(ends);
+    std::vector<int64_t> ends_int = vec_int64.value();
+
+    vec_int64 = details::VecExpr2Int64(strides);
+    std::vector<int64_t> strides_int = vec_int64.value();
+    vec_int64 = details::VecExpr2Int64(in_dims);
+    std::vector<int64_t> in_dims_int = vec_int64.value();
+    bool dummy_zero_dim_out = false;
+    phi::funcs::normalize_interval(starts_int[0],
+                                   ends_int[0],
+                                   strides_int[0],
+                                   in_dims_int[0],
+                                   &starts_int[0],
+                                   &ends_int[0],
+                                   &dummy_zero_dim_out);
+    if (ends_int[0] == -in_dims_int[0] - 1) {
+      ends_int[0] = -1;
+    }
+    if (strides_int[0] > 0) {
+      for (int64_t i = starts_int[0]; i < ends_int[0]; i += strides_int[0]) {
+        out_data.push_back(in_shapeordata.data().value().at(i));
+      }
+    } else {
+      for (int64_t i = starts_int[0]; i > ends_int[0]; i += strides_int[0]) {
+        out_data.push_back(in_shapeordata.data().value().at(i));
+      }
+    }
+    const ExprVec shape = GetDecreasedDims(
+        ExprVec{static_cast<int64_t>(out_data.size())}, decrease_axis);
+    if (shape.size() == 1 && shape[0] == 0) {
+      return symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(shape)};
+    }
+    return symbol::ShapeOrDataDimExprs{
+        symbol::TensorShapeOrDataDimExprs(shape, out_data)};
+  };
+  bool starts_ends_all_int =
+      std::all_of(starts_expr.begin(),
+                  starts_expr.end(),
+                  [](const symbol::DimExpr &e) { return e.isa<int64_t>(); }) &&
+      std::all_of(ends_expr.begin(),
+                  ends_expr.end(),
+                  [](const symbol::DimExpr &e) { return e.isa<int64_t>(); }) &&
+      std::all_of(strides_expr.begin(),
+                  strides_expr.end(),
+                  [](const symbol::DimExpr &e) { return e.isa<int64_t>(); });
+  const auto &out_shape =
+      in_shapeordata.data().has_value() && starts_ends_all_int
+          ? 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
@@ -3507,12 +3507,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
@@ -4832,7 +4832,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
@@ -764,7 +764,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
@@ -5678,7 +5678,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
