[slice] Updating the computation of the index_elementwise output dim (#73184)

zhanghonggeng · web-flow · commit 2c5572658900 · 2025-06-10T19:32:03.000+08:00
diff --git a/paddle/fluid/pybind/slice_utils.h b/paddle/fluid/pybind/slice_utils.h
@@ -102,9 +102,7 @@ static inline std::vector<paddle::Tensor> expand_outplace(
 }
 
 struct AdvancedIndex {
-  AdvancedIndex(paddle::Tensor src,
-                std::vector<paddle::Tensor> indices,
-                bool bool_case);
+  AdvancedIndex(paddle::Tensor src, std::vector<paddle::Tensor> indices);
 
   paddle::Tensor src;
   std::vector<paddle::Tensor> indices;
@@ -145,8 +143,7 @@ inline static paddle::Tensor reshape_indexer(paddle::Tensor* index,
 }
 
 inline AdvancedIndex::AdvancedIndex(paddle::Tensor src,
-                                    std::vector<paddle::Tensor> indices_list,
-                                    bool bool_case = false) {
+                                    std::vector<paddle::Tensor> indices_list) {
   uint32_t element_size_bytes = phi::SizeOf(src.dtype());
   int64_t dims_before = 0, dims_after = 0, dims_indexed = 0;
   std::vector<int64_t> shape_vec = common::vectorize<int64_t>(src.dims());
@@ -165,10 +162,6 @@ inline AdvancedIndex::AdvancedIndex(paddle::Tensor src,
     } else {
       dims_indexed++;
       replacement_shape = common::vectorize<int64_t>(indices_list[dim].dims());
-      if (bool_case && !replacement_shape.empty() &&
-          replacement_shape.back() == 1) {
-        replacement_shape.pop_back();
-      }
 
       idx_shape_vec.push_back(shape_vec[dim]);
       idx_stride_vec.push_back(stride_vec[dim] * element_size_bytes);
@@ -654,23 +647,9 @@ static std::vector<paddle::Tensor> PrepareIndices(
   for (int j = 0; j < bool_2_idx.shape()[1]; ++j) {
     paddle::Tensor sliced_tensor =
         slice_ad_func(bool_2_idx, {1}, {j}, {j + 1}, {1}, {});
-
-    // Calculate the required dimensionality
-    int64_t original_ndim =
-        tensor.shape().size() - bool_index.shape().size() + 1;
-    int64_t sliced_ndim = sliced_tensor.shape().size();
-    int64_t num_ones_to_add = original_ndim - sliced_ndim;
-
-    // Reshape the tensor by adding 1s if needed
-    if (num_ones_to_add > 0) {
-      std::vector<int64_t> new_shape = sliced_tensor.shape();
-      for (int64_t k = 0; k < num_ones_to_add; ++k) {
-        new_shape.push_back(1);
-      }
-      sliced_tensor = reshape_ad_func(sliced_tensor, new_shape);
-    }
-
-    indices.emplace_back(sliced_tensor);
+    paddle::Tensor sliced_tensor_c = sliced_tensor.contiguous();
+    sliced_tensor_c.reshape({sliced_tensor.dims()[0]});
+    indices.emplace_back(sliced_tensor_c);
   }
   return indices;
 }
@@ -715,17 +694,24 @@ static paddle::Tensor getValueForBoolTensor(const paddle::Tensor& tensor,
 
   auto bool_2_idx = nonzero_ad_func(bool_index);
 #ifdef PADDLE_WITH_CUDA
-  auto indices = PrepareIndices(tensor, bool_2_idx, bool_index);
-  AdvancedIndex ad = AdvancedIndex(tensor, indices, true);
-
-  return index_elementwise_get_ad_func(tensor,
-                                       ad.indices,
-                                       ad.src_sizes,
-                                       ad.src_strides,
-                                       ad.indexed_sizes,
-                                       ad.indexed_strides);
-#else
+  if (tensor.is_gpu()) {
+    auto indices = PrepareIndices(tensor, bool_2_idx, bool_index);
+    while (indices.size() < static_cast<size_t>(tensor.dims().size())) {
+      indices.emplace_back();
+    }
+
+    AdvancedIndex ad = AdvancedIndex(tensor, indices);
 
+    return index_elementwise_get_ad_func(tensor,
+                                         ad.indices,
+                                         ad.src_sizes,
+                                         ad.src_strides,
+                                         ad.indexed_sizes,
+                                         ad.indexed_strides);
+  } else {
+    return gather_nd_ad_func(tensor, bool_2_idx);
+  }
+#else
   return gather_nd_ad_func(tensor, bool_2_idx);
 #endif
 }
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -2133,51 +2133,6 @@ void GatherNdInferMeta(const MetaTensor& x,
   out->set_dtype(x.dtype());
 }
 
-void IndexElementwisePutInferMeta(const MetaTensor& x,
-                                  const std::vector<const MetaTensor*>& index,
-                                  const MetaTensor& value,
-                                  const std::vector<int64_t>& input_dims,
-                                  const std::vector<int64_t>& input_strides,
-                                  const std::vector<int64_t>& index_dims,
-                                  const std::vector<int64_t>& index_strides,
-                                  MetaTensor* out) {
-  out->set_dims(x.dims());
-  out->set_dtype(x.dtype());
-}
-
-void IndexElementwiseGetInferMeta(const MetaTensor& x,
-                                  const std::vector<const MetaTensor*>& index,
-                                  const std::vector<int64_t>& input_dims,
-                                  const std::vector<int64_t>& input_strides,
-                                  const std::vector<int64_t>& index_dims,
-                                  const std::vector<int64_t>& index_stride,
-                                  MetaTensor* out) {
-  const auto& x_dims = x.dims();
-
-  PADDLE_ENFORCE_LE(
-      index_dims.size(),
-      x_dims.size(),
-      common::errors::InvalidArgument(
-          "Input(index).rank should be less than or equal to Input(x).rank"));
-
-  for (size_t i = 0; i < index_dims.size(); ++i) {
-    PADDLE_ENFORCE_EQ(index_dims[i],
-                      x_dims[i],
-                      common::errors::InvalidArgument(
-                          "Input(index).shape should match the first k "
-                          "dimensions of Input(x).shape"));
-  }
-  std::vector<int64_t> result_dims;
-  result_dims.push_back(-1);
-
-  int64_t k = index_dims.size();
-  for (int64_t i = k; i < static_cast<int64_t>(x_dims.size()); ++i) {
-    result_dims.push_back(x_dims[i]);
-  }
-  out->set_dims(common::make_ddim(result_dims));
-  out->set_dtype(x.dtype());
-}
-
 void GatherTreeMeta(const MetaTensor& ids,
                     const MetaTensor& parents,
                     MetaTensor* out) {
@@ -2591,6 +2546,51 @@ void IndexAddInferMeta(const MetaTensor& x,
   output->share_lod(x);
 }
 
+void IndexElementwisePutInferMeta(const MetaTensor& x,
+                                  const std::vector<const MetaTensor*>& index,
+                                  const MetaTensor& value,
+                                  const std::vector<int64_t>& input_dims,
+                                  const std::vector<int64_t>& input_strides,
+                                  const std::vector<int64_t>& index_dims,
+                                  const std::vector<int64_t>& index_strides,
+                                  MetaTensor* out) {
+  out->set_dims(x.dims());
+  out->set_dtype(x.dtype());
+}
+
+void IndexElementwiseGetInferMeta(const MetaTensor& x,
+                                  const std::vector<const MetaTensor*>& index,
+                                  const std::vector<int64_t>& input_dims,
+                                  const std::vector<int64_t>& input_strides,
+                                  const std::vector<int64_t>& index_dims,
+                                  const std::vector<int64_t>& index_stride,
+                                  MetaTensor* out) {
+  const auto& x_dims = x.dims();
+
+  PADDLE_ENFORCE_LE(
+      index_dims.size(),
+      x_dims.size(),
+      common::errors::InvalidArgument(
+          "Input(index).rank should be less than or equal to Input(x).rank"));
+
+  for (size_t i = 0; i < index_dims.size(); ++i) {
+    PADDLE_ENFORCE_EQ(index_dims[i],
+                      x_dims[i],
+                      common::errors::InvalidArgument(
+                          "Input(index).shape should match the first k "
+                          "dimensions of Input(x).shape"));
+  }
+  std::vector<int64_t> result_dims;
+  result_dims.push_back(-1);
+
+  int64_t k = index_dims.size();
+  for (int64_t i = k; i < static_cast<int64_t>(x_dims.size()); ++i) {
+    result_dims.push_back(x_dims[i]);
+  }
+  out->set_dims(common::make_ddim(result_dims));
+  out->set_dtype(x.dtype());
+}
+
 void KronInferMeta(const MetaTensor& x, const MetaTensor& y, MetaTensor* out) {
   auto dim_x = x.dims();
   auto dim_y = y.dims();
diff --git a/paddle/phi/infermeta/binary.h b/paddle/phi/infermeta/binary.h
@@ -401,23 +401,6 @@ void GatherNdInferMeta(const MetaTensor& x,
                        const MetaTensor& index,
                        MetaTensor* out);
 
-void IndexElementwisePutInferMeta(const MetaTensor& x,
-                                  const std::vector<const MetaTensor*>& index,
-                                  const MetaTensor& value,
-                                  const std::vector<int64_t>& input_dims,
-                                  const std::vector<int64_t>& input_strides,
-                                  const std::vector<int64_t>& index_dims,
-                                  const std::vector<int64_t>& index_strides,
-                                  MetaTensor* out);
-
-void IndexElementwiseGetInferMeta(const MetaTensor& x,
-                                  const std::vector<const MetaTensor*>& index,
-                                  const std::vector<int64_t>& input_dims,
-                                  const std::vector<int64_t>& input_strides,
-                                  const std::vector<int64_t>& index_dims,
-                                  const std::vector<int64_t>& index_stride,
-                                  MetaTensor* out);
-
 void GatherTreeMeta(const MetaTensor& ids,
                     const MetaTensor& parents,
                     MetaTensor* out);
@@ -481,6 +464,23 @@ void IndexAddInferMeta(const MetaTensor& x,
                        int axis,
                        MetaTensor* output);
 
+void IndexElementwisePutInferMeta(const MetaTensor& x,
+                                  const std::vector<const MetaTensor*>& index,
+                                  const MetaTensor& value,
+                                  const std::vector<int64_t>& input_dims,
+                                  const std::vector<int64_t>& input_strides,
+                                  const std::vector<int64_t>& index_dims,
+                                  const std::vector<int64_t>& index_strides,
+                                  MetaTensor* out);
+
+void IndexElementwiseGetInferMeta(const MetaTensor& x,
+                                  const std::vector<const MetaTensor*>& index,
+                                  const std::vector<int64_t>& input_dims,
+                                  const std::vector<int64_t>& input_strides,
+                                  const std::vector<int64_t>& index_dims,
+                                  const std::vector<int64_t>& index_stride,
+                                  MetaTensor* out);
+
 void KronInferMeta(const MetaTensor& x, const MetaTensor& y, MetaTensor* out);
 
 void LegacyCropInferMeta(const MetaTensor& x,
diff --git a/paddle/phi/kernels/gpu/index_elementwise_get_grad_kernel.cu b/paddle/phi/kernels/gpu/index_elementwise_get_grad_kernel.cu
@@ -90,9 +90,12 @@ void GPUIndexElementwisePutKernel(const phi::GPUContext& ctx,
       desired_shape.size(), desired_shape.data(), strides_vector);
 
   const int64_t N = numel;
-  PADDLE_ENFORCE(N >= 0 && N <= std::numeric_limits<int32_t>::max(),
-                 "Output numel be in the range [0, "
-                 "std::numeric_limits<int32_t>::max()]");
+  PADDLE_ENFORCE_GE(
+      N, 0, common::errors::InvalidArgument("Output numel must >= 0"));
+  PADDLE_ENFORCE_LE(
+      N,
+      std::numeric_limits<int32_t>::max(),
+      common::errors::InvalidArgument("Output numel must <= INT32_MAX"));
 
   constexpr int nt = 128;
   constexpr int vt = 4;
@@ -116,8 +119,10 @@ void GPUIndexElementwisePutKernel(const phi::GPUContext& ctx,
         for (int i = 0; i < num_indices; i++) {
           int64_t index =
               *reinterpret_cast<int64_t*>(index_ptrs[i] + offsets[2]);
-          PADDLE_ENFORCE(-sizes[i] <= index && index < sizes[i],
-                         "index out of bounds");
+          PADDLE_ENFORCE(-sizes[i] <= index,
+                         "index is less than the lower bound");
+          PADDLE_ENFORCE(index < sizes[i],
+                         "index is greater than or equal to the upper bound");
           if (index < 0) {
             index += sizes[i];
           }
diff --git a/paddle/phi/kernels/gpu/index_elementwise_get_kernel.cu b/paddle/phi/kernels/gpu/index_elementwise_get_kernel.cu
@@ -60,9 +60,10 @@ void GPUIndexElementwiseGetKernel(const phi::GPUContext& ctx,
                            &numel);
 
   const int64_t* template_stride = strides_array[2];
-  PADDLE_ENFORCE(
-      template_stride != nullptr,
-      "strides_array[2] should not be nullptr in GPUIndexElementwiseGetKernel");
+  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;
@@ -87,9 +88,12 @@ void GPUIndexElementwiseGetKernel(const phi::GPUContext& ctx,
       desired_shape.size(), desired_shape.data(), strides_vector);
 
   const int64_t N = output->numel();
-  PADDLE_ENFORCE(N >= 0 && N <= std::numeric_limits<int32_t>::max(),
-                 "Output numel be in the range [0, "
-                 "std::numeric_limits<int32_t>::max()]");
+  PADDLE_ENFORCE_GE(
+      N, 0, common::errors::InvalidArgument("Output numel must >= 0"));
+  PADDLE_ENFORCE_LE(
+      N,
+      std::numeric_limits<int32_t>::max(),
+      common::errors::InvalidArgument("Output numel must <= INT32_MAX"));
   constexpr int nt = 128;
   constexpr int vt = 4;
   const dim3 block(nt);
@@ -112,8 +116,10 @@ void GPUIndexElementwiseGetKernel(const phi::GPUContext& ctx,
         for (int i = 0; i < num_indices; i++) {
           int64_t index =
               *reinterpret_cast<int64_t*>(index_ptrs[i] + offsets[2]);
-          PADDLE_ENFORCE(-sizes[i] <= index && index < sizes[i],
-                         "index out of bounds");
+          PADDLE_ENFORCE(-sizes[i] <= index,
+                         "index is less than the lower bound");
+          PADDLE_ENFORCE(index < sizes[i],
+                         "index is greater than or equal to the upper bound");
           if (index < 0) {
             index += sizes[i];
           }
@@ -145,11 +151,10 @@ void IndexElementwiseGetKernel(const Context& ctx,
           phi::DataType::INT32,
           phi::DataType::INT64));
 
-  size_t total_nonzero = index[0]->numel();
   auto out_dims = out->dims();
   if (out_dims.size() > 0) {
-    out_dims[0] = total_nonzero;
-    out->Resize(out_dims);
+    std::vector<int64_t> output_dims(input_dims);
+    out->Resize(phi::make_ddim(output_dims));
   }
 
   ctx.template Alloc<T>(out);
