[Slice] fix IndexGetStride bug (#73237)

zhanghonggeng · web-flow · commit c9cc49e279b4 · 2025-06-11T11:23:47.000+08:00
diff --git a/paddle/phi/kernels/funcs/stride_utils.h b/paddle/phi/kernels/funcs/stride_utils.h
@@ -297,11 +297,11 @@ static inline void IndexPutStride(
     const int64_t index_elesize,
     std::vector<int64_t>* desired_shape,
     std::array<int64_t*, N>* strides_array,
-    int64_t* numel) {
+    int64_t* numel,
+    std::array<std::vector<int64_t>, N>& strides_vec) {  // NOLINT
   int64_t index_size = 1;
   int ndim = output_dims.size();
 
-  std::array<std::vector<int64_t>, N> strides_vec;
   std::vector<int64_t> stride_size;
 
   *desired_shape = compute_shapes({input_dims, output_dims, index_dims});
@@ -341,22 +341,23 @@ static inline void IndexPutStride(
 }
 
 template <int N>
-static inline void IndexGetStride(const std::vector<int64_t> output_dims,
-                                  const std::vector<int64_t> output_strides,
-                                  const int64_t output_elesize,
-                                  const std::vector<int64_t> input_dims,
-                                  const std::vector<int64_t> input_strides,
-                                  const int64_t input_elesize,
-                                  const std::vector<int64_t> index_dims,
-                                  const std::vector<int64_t> index_strides,
-                                  const int64_t index_elesize,
-                                  std::vector<int64_t>* desired_shape,
-                                  std::array<int64_t*, N>* strides_array,
-                                  int64_t* numel) {
+static inline void IndexGetStride(
+    const std::vector<int64_t> output_dims,
+    const std::vector<int64_t> output_strides,
+    const int64_t output_elesize,
+    const std::vector<int64_t> input_dims,
+    const std::vector<int64_t> input_strides,
+    const int64_t input_elesize,
+    const std::vector<int64_t> index_dims,
+    const std::vector<int64_t> index_strides,
+    const int64_t index_elesize,
+    std::vector<int64_t>* desired_shape,
+    std::array<int64_t*, N>* strides_array,
+    int64_t* numel,
+    std::array<std::vector<int64_t>, N>& strides_vec) {  // NOLINT
   int64_t index_size = 1;
   int ndim = output_dims.size();
 
-  std::array<std::vector<int64_t>, N> strides_vec;
   std::vector<int64_t> stride_size;
 
   *desired_shape = compute_shapes({input_dims, output_dims, index_dims});
diff --git a/paddle/phi/kernels/gpu/index_elementwise_get_grad_kernel.cu b/paddle/phi/kernels/gpu/index_elementwise_get_grad_kernel.cu
@@ -47,6 +47,7 @@ void GPUIndexElementwisePutKernel(const phi::GPUContext& ctx,
 
   std::array<int64_t*, 3> strides_array;
   std::vector<int64_t> desired_shape;
+  std::array<std::vector<int64_t>, 3> strides_vec;
 
   funcs::IndexPutStride<3>(input_dims,
                            input_strides,
@@ -59,7 +60,8 @@ void GPUIndexElementwisePutKernel(const phi::GPUContext& ctx,
                            phi::SizeOf(index[0]->dtype()),
                            &desired_shape,
                            &strides_array,
-                           &numel);
+                           &numel,
+                           strides_vec);
 
   const int64_t* template_stride = strides_array[2];
   PADDLE_ENFORCE_NOT_NULL(template_stride,
diff --git a/paddle/phi/kernels/gpu/index_elementwise_get_kernel.cu b/paddle/phi/kernels/gpu/index_elementwise_get_kernel.cu
@@ -45,6 +45,7 @@ void GPUIndexElementwiseGetKernel(const phi::GPUContext& ctx,
 
   std::array<int64_t*, 3> strides_array;
   std::vector<int64_t> desired_shape;
+  std::array<std::vector<int64_t>, 3> strides_vec;
 
   funcs::IndexGetStride<3>(input_dims,
                            input_strides,
@@ -57,35 +58,10 @@ void GPUIndexElementwiseGetKernel(const phi::GPUContext& ctx,
                            phi::SizeOf(index[0]->dtype()),
                            &desired_shape,
                            &strides_array,
-                           &numel);
-
-  const int64_t* template_stride = strides_array[2];
-  PADDLE_ENFORCE_NOT_NULL(
-      template_stride,
-      common::errors::InvalidArgument("strides_array[2] must not be nullptr in "
-                                      "GPUIndexElementwiseGetKernel"));
-
-  size_t stride_size = desired_shape.size();
-  std::vector<std::vector<int64_t>> strides_vector;
-  strides_vector.reserve(num_indices + 2);
-
-  for (int i = 0; i < 2; ++i) {
-    if (i < strides_array.size() && strides_array[i] != nullptr) {
-      strides_vector.emplace_back(strides_array[i],
-                                  strides_array[i] + stride_size);
-    } else {
-      strides_vector.emplace_back(stride_size, 0);
-    }
-  }
-
-  std::vector<int64_t> template_vec(template_stride,
-                                    template_stride + stride_size);
-  for (size_t i = 0; i < num_indices; ++i) {
-    strides_vector.push_back(template_vec);
-  }
-
-  auto offset_calc = funcs::make_offset_calculator<3>(
-      desired_shape.size(), desired_shape.data(), strides_vector);
+                           &numel,
+                           strides_vec);
+  auto offset_calc =
+      funcs::make_offset_calculator_put<3>(desired_shape, strides_array);
 
   const int64_t N = output->numel();
   PADDLE_ENFORCE_GE(
diff --git a/paddle/phi/kernels/gpu/index_elementwise_put_grad_kernel.cu b/paddle/phi/kernels/gpu/index_elementwise_put_grad_kernel.cu
@@ -50,6 +50,8 @@ void GPUIndexElementwisePutGradKernel(
 
   std::array<int64_t*, 3> strides_array;
   std::vector<int64_t> desired_shape;
+  std::array<std::vector<int64_t>, 3> strides_vec;
+
   funcs::IndexPutStride<3>(input_dims,
                            input_strides,
                            phi::SizeOf(output->dtype()),
@@ -61,7 +63,8 @@ void GPUIndexElementwisePutGradKernel(
                            phi::SizeOf(index[0]->dtype()),
                            &desired_shape,
                            &strides_array,
-                           &numel);
+                           &numel,
+                           strides_vec);
   auto offset_calc =
       funcs::make_offset_calculator_put<3>(desired_shape, strides_array);
   const int64_t N = numel;
diff --git a/paddle/phi/kernels/gpu/index_elementwise_put_kernel.cu b/paddle/phi/kernels/gpu/index_elementwise_put_kernel.cu
@@ -46,6 +46,7 @@ void GPUIndexElementwisePutKernel(const phi::GPUContext& dev_ctx,
 
   std::array<int64_t*, 3> strides_array;
   std::vector<int64_t> desired_shape;
+  std::array<std::vector<int64_t>, 3> strides_vec;
 
   funcs::IndexPutStride<3>(input_dims,
                            input_strides,
@@ -58,7 +59,8 @@ void GPUIndexElementwisePutKernel(const phi::GPUContext& dev_ctx,
                            phi::SizeOf(index[0]->dtype()),
                            &desired_shape,
                            &strides_array,
-                           &numel);
+                           &numel,
+                           strides_vec);
 
   auto offset_calc =
       funcs::make_offset_calculator_put<3>(desired_shape, strides_array);
