【slice】modify Set value cuda kernel (stride slice+ stride copy) (#72838)

xiaoguoguo626807 · web-flow · commit aeb8275d8d51 · 2025-05-23T11:29:41.000+08:00
* dpa3 fix

* modify tuple pop check

* add set_value.cu

* fix ci wrong case

* modify windows bug, input dim was change in win

* modify setvalue 0-size part
diff --git a/paddle/phi/kernels/funcs/slice_utils.h b/paddle/phi/kernels/funcs/slice_utils.h
@@ -335,6 +335,67 @@ inline DDim GetDecreasedDims(const DDim slice_dims,
   return decreased_dims;
 }
 
+template <typename T = int64_t>
+inline void GetDecreasedDimsAndStrides(const std::vector<T> slice_dims,
+                                       const std::vector<T> slice_strides,
+                                       const std::vector<T>& decrease_axes,
+                                       const std::vector<T>& none_axes,
+                                       std::vector<T>* new_dims,
+                                       std::vector<T>* new_strides,
+                                       std::vector<T>* infer_flags = nullptr) {
+  std::vector<uint8_t> decrease_flag(slice_dims.size(), 0);
+
+  if (none_axes.size() > 0) {
+    size_t none_axes_cur = 0, decrease_axes_cur = 0;
+    for (int i = 0; i < slice_dims.size(); ++i) {
+      while (none_axes_cur < none_axes.size() &&
+             none_axes[none_axes_cur] <= i) {
+        new_dims->push_back(1);
+        new_strides->push_back(slice_strides[i]);
+        none_axes_cur++;
+      }
+      if (decrease_axes_cur < decrease_axes.size() &&
+          decrease_axes[decrease_axes_cur] == i) {
+        decrease_axes_cur++;
+      } else {
+        new_dims->push_back(slice_dims[i]);
+        new_strides->push_back(slice_strides[i]);
+      }
+    }
+    while (none_axes_cur < none_axes.size()) {
+      new_dims->push_back(1);
+      new_strides->push_back(slice_strides[-1]);
+      none_axes_cur++;
+    }
+  } else if (decrease_axes.size() > 0) {
+    for (size_t i = 0; i < decrease_axes.size(); ++i) {
+      T axis = decrease_axes[i];
+      decrease_flag[axis] = 1;
+      if (infer_flags && (*infer_flags)[i] != -1) {
+        PADDLE_ENFORCE_EQ(slice_dims[axis],
+                          1,
+                          common::errors::InvalidArgument(
+                              "Decrease dim should be 1, but now received %d",
+                              slice_dims[axis]));
+      }
+    }
+
+    for (int i = 0; i < slice_dims.size(); ++i) {
+      if (decrease_flag[i] == 0) {
+        new_dims->push_back(slice_dims[i]);
+        new_strides->push_back(slice_strides[i]);
+      }
+    }
+  } else {
+    for (int i = 0; i < slice_dims.size(); ++i) {
+      new_dims->push_back(slice_dims[i]);
+      new_strides->push_back(slice_strides[i]);
+    }
+  }
+
+  return;
+}
+
 template <typename T = int64_t>
 inline void CheckAndUpdateSparseSliceAttrs(const DDim in_dims,
                                            std::vector<T>* axes,
diff --git a/paddle/phi/kernels/gpu/set_value_kernel.cu b/paddle/phi/kernels/gpu/set_value_kernel.cu
@@ -13,16 +13,162 @@
 // limitations under the License.
 
 #include "paddle/phi/kernels/set_value_kernel.h"
-
+#include <chrono>
+#include <iostream>
+#include <type_traits>
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/common/complex.h"
+#include "paddle/phi/common/int_array.h"
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/expand_kernel.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/slice_utils.h"
 #include "paddle/phi/kernels/impl/set_value_kernel_impl.h"
+#include "paddle/phi/kernels/strided_copy_kernel.h"
+namespace phi {
+template <typename T, typename Context>
+void SetTensorValueKernelV2(const Context& dev_ctx,
+                            const DenseTensor& in,
+                            const DenseTensor& value,
+                            const IntArray& starts,
+                            const IntArray& ends,
+                            const IntArray& steps,
+                            const std::vector<int64_t>& axes,
+                            const std::vector<int64_t>& decrease_axes,
+                            const std::vector<int64_t>& none_axes,
+                            DenseTensor* out) {
+  auto in_dims = in.dims();
+  auto meta = in.meta();
+  std::vector<int64_t> starts_local = starts.GetData();
+  std::vector<int64_t> ends_local = ends.GetData();
+  std::vector<int64_t> steps_local = steps.GetData();
+  phi::funcs::CheckAndUpdateSliceAttrs(
+      in_dims, axes, &starts_local, &ends_local, &steps_local);
+
+  std::vector<int64_t> output_dims = common::vectorize<int64_t>(in.dims());
+  std::vector<int64_t> output_stride = common::vectorize<int64_t>(in.strides());
+  int64_t output_offset = static_cast<int64_t>(in.offset());
+  for (size_t i = 0; i < axes.size(); ++i) {
+    int64_t axis_size = in.dims()[axes[i]];
+    if (axis_size < 0) {
+      continue;
+    }
+
+    int64_t step_size = std::abs(steps_local[i]);
+
+    auto out_dim =
+        (std::abs(ends_local[i] - starts_local[i]) + step_size - 1) / step_size;
+    output_offset += static_cast<int>(starts_local[i] * output_stride[axes[i]] *
+                                      SizeOf(out->dtype()));
+    output_dims[axes[i]] = out_dim;
+    output_stride[axes[i]] *= steps_local[i];
+  }
+
+  // generate new shape
+  std::vector<int64_t> new_out_shape;
+  std::vector<int64_t> new_out_stride;
+  funcs::GetDecreasedDimsAndStrides(output_dims,
+                                    output_stride,
+                                    decrease_axes,
+                                    none_axes,
+                                    &new_out_shape,
+                                    &new_out_stride);
+
+  if (product(phi::make_ddim(new_out_shape)) <= 0) {
+    out->ResetHolder(in.Holder());
+    out->ShareInplaceVersionCounterWith(in);
+    return;
+  }
+
+  CheckIsDimsMatch(phi::make_ddim(new_out_shape), value.dims());
+  auto value_dims = phi::vectorize<int64_t>(value.dims());
+  DenseTensor value_tensor = Empty<T>(dev_ctx, IntArray{value_dims});
+  value_tensor = value;
+  auto it = value_dims.begin();
+  while (it != value_dims.end() && *it == 1) {
+    it = value_dims.erase(it);
+  }
+  if (value_dims.empty()) value_dims.push_back(1);
+  value_tensor.Resize(phi::make_ddim(value_dims));
+
+  if (new_out_shape.empty()) new_out_shape.push_back(1);
+
+  DenseTensor expand_tensor = Empty<T>(dev_ctx, IntArray{new_out_shape});
+  ExpandKernel<T, Context>(
+      dev_ctx, value_tensor, IntArray{new_out_shape}, &expand_tensor);
+
+  out->ResetHolder(in.Holder());
+  out->ShareInplaceVersionCounterWith(in);
+  StridedCopyKernel<T, Context>(dev_ctx,
+                                expand_tensor,
+                                new_out_shape,
+                                new_out_stride,
+                                output_offset,
+                                out);
+  out->set_meta(meta);
+}
+
+template <typename T, typename Context>
+void SetValueKernelV2(const Context& dev_ctx,
+                      const DenseTensor& in,
+                      const IntArray& starts,
+                      const IntArray& ends,
+                      const IntArray& steps,
+                      const std::vector<int64_t>& axes,
+                      const std::vector<int64_t>& decrease_axes,
+                      const std::vector<int64_t>& none_axes,
+                      const std::vector<int64_t>& shape,
+                      const std::vector<Scalar>& values,
+                      DenseTensor* out) {
+  // auto t1 = std::chrono::high_resolution_clock::now();
+  std::vector<T> assign_values;
+  assign_values.reserve(values.size());
+  for (const auto& val : values) {
+    assign_values.push_back(val.to<T>());
+  }
+
+  bool is_full_set_one_value = false;
+  std::vector<int64_t> starts_local = starts.GetData();
+  std::vector<int64_t> ends_local = ends.GetData();
+  std::vector<int64_t> steps_local = steps.GetData();
+  if (starts_local.empty() && ends_local.empty() && steps_local.empty() &&
+      shape.size() == 1 && shape[0] == 1 && assign_values.size() == 1) {
+    is_full_set_one_value = true;
+  }
+
+  if (is_full_set_one_value && std::is_same<T, float>::value) {
+    dev_ctx.template Alloc<T>(out);
+    phi::funcs::set_constant(
+        dev_ctx, out, static_cast<float>(assign_values[0]));
+    return;
+  }
+
+  DenseTensor value_tensor = Empty<T>(dev_ctx, shape);
+  phi::TensorFromVector(assign_values, dev_ctx, &value_tensor);
+  value_tensor.Resize(common::make_ddim(shape));
+
+  SetTensorValueKernelV2<T, Context>(dev_ctx,
+                                     in,
+                                     value_tensor,
+                                     starts,
+                                     ends,
+                                     steps,
+                                     axes,
+                                     decrease_axes,
+                                     none_axes,
+                                     out);
+}
+
+}  // namespace phi
 
 PD_REGISTER_KERNEL(set_value,
                    GPU,
                    ALL_LAYOUT,
-                   phi::SetValueKernel,
+                   phi::SetValueKernelV2,
                    float,
                    double,
                    int,
@@ -38,7 +184,7 @@ PD_REGISTER_KERNEL(set_value,
 PD_REGISTER_KERNEL(set_value_with_tensor,
                    GPU,
                    ALL_LAYOUT,
-                   phi::SetTensorValueKernel,
+                   phi::SetTensorValueKernelV2,
                    float,
                    double,
                    int,
