[PHI] The gathernd and scatternd kernels replace the Dispatch kernel with a switch case. (#72741)

zhanghonggeng · web-flow · commit 0df410c83adb · 2025-05-19T20:13:15.000+08:00
diff --git a/paddle/phi/kernels/funcs/gather.cu.h b/paddle/phi/kernels/funcs/gather.cu.h
@@ -37,18 +37,18 @@ __global__ void GatherNdCUDAKernel(const T* input,
                                    size_t remain_size,
                                    size_t slice_size,
                                    size_t end_size) {
-  int64_t total_size = remain_size * slice_size;
-  int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * VecSize;
+  int total_size = remain_size * slice_size;
+  int idx = (blockIdx.x * blockDim.x + threadIdx.x) * VecSize;
   int64_t stride = blockDim.x * gridDim.x * VecSize;
 
 #pragma unroll
   for (; idx < total_size; idx += stride) {
-    int64_t indices_i = idx / slice_size;
-    int64_t slice_i = idx % slice_size;
+    int indices_i = idx / slice_size;
+    int slice_i = idx % slice_size;
     int64_t gather_i = 0;
     int64_t temp = slice_size;
 #pragma unroll
-    for (int64_t j = end_size - 1; j >= 0; --j) {
+    for (int j = end_size - 1; j >= 0; --j) {
       auto index_value = indices[indices_i * end_size + j];
       PADDLE_ENFORCE(
           index_value >= -input_dims[j] && index_value < input_dims[j],
@@ -75,70 +75,6 @@ __global__ void GatherNdCUDAKernel(const T* input,
   }
 }
 
-template <typename T, typename IndexT, int VecSize>
-void DispatchGatherNdKernel(const phi::GPUContext& ctx,
-                            const T* p_input,
-                            const Dim<DDim::kMaxRank>& g_input_dims,
-                            const IndexT* p_index,
-                            T* p_output,
-                            int64_t remain_numel,
-                            int64_t slice_size,
-                            int64_t end_size,
-                            int vec_size,
-                            const phi::backends::gpu::GpuLaunchConfig& config) {
-  if (vec_size == VecSize) {
-    auto stream = ctx.stream();
-    GatherNdCUDAKernel<T, IndexT, VecSize>
-        <<<config.block_per_grid, config.thread_per_block, 0, stream>>>(
-            p_input,
-            g_input_dims,
-            p_index,
-            p_output,
-            remain_numel,
-            slice_size,
-            end_size);
-  } else {
-    DispatchGatherNdKernel<T, IndexT, VecSize / 2>(ctx,
-                                                   p_input,
-                                                   g_input_dims,
-                                                   p_index,
-                                                   p_output,
-                                                   remain_numel,
-                                                   slice_size,
-                                                   end_size,
-                                                   vec_size,
-                                                   config);
-  }
-}
-
-template <typename T, typename IndexT>
-void DispatchGatherNdKernel(const phi::GPUContext& ctx,
-                            const T* p_input,
-                            const Dim<DDim::kMaxRank>& g_input_dims,
-                            const IndexT* p_index,
-                            T* p_output,
-                            int64_t remain_numel,
-                            int64_t slice_size,
-                            int64_t end_size,
-                            int vec_size,
-                            const phi::backends::gpu::GpuLaunchConfig& config) {
-  if (vec_size == 1) {
-    auto stream = ctx.stream();
-    GatherNdCUDAKernel<T, IndexT, 1>
-        <<<config.block_per_grid, config.thread_per_block, 0, stream>>>(
-            p_input,
-            g_input_dims,
-            p_index,
-            p_output,
-            remain_numel,
-            slice_size,
-            end_size);
-  } else {
-    PADDLE_THROW(common::errors::Unimplemented(
-        "Unsupported vectorized size: %d", vec_size));
-  }
-}
-
 template <typename T, typename IndexT = int>
 void GPUGatherNd(const phi::GPUContext& ctx,
                  const DenseTensor& input,
@@ -183,16 +119,29 @@ void GPUGatherNd(const phi::GPUContext& ctx,
   auto config = phi::backends::gpu::GetGpuLaunchConfig1D(
       ctx, remain_numel * slice_size, vec_size * loop_count);
 
-  DispatchGatherNdKernel<T, IndexT, 4>(ctx,
-                                       p_input,
-                                       g_input_dims,
-                                       p_index,
-                                       p_output,
-                                       remain_numel,
-                                       slice_size,
-                                       end_size,
-                                       vec_size,
-                                       config);
+  auto stream = ctx.stream();
+
+  switch (vec_size) {
+#define CASE_VEC_SIZE(__Sz)                                              \
+  case __Sz:                                                             \
+    GatherNdCUDAKernel<T, IndexT, __Sz>                                  \
+        <<<config.block_per_grid, config.thread_per_block, 0, stream>>>( \
+            p_input,                                                     \
+            g_input_dims,                                                \
+            p_index,                                                     \
+            p_output,                                                    \
+            remain_numel,                                                \
+            slice_size,                                                  \
+            end_size);                                                   \
+    break
+    CASE_VEC_SIZE(4);
+    CASE_VEC_SIZE(2);
+    CASE_VEC_SIZE(1);
+#undef CASE_VEC_SIZE
+    default:
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Unsupported vectorized size: %d", vec_size));
+  }
 }
 
 template <typename T, typename U, int VecSize>
diff --git a/paddle/phi/kernels/funcs/scatter.cu.h b/paddle/phi/kernels/funcs/scatter.cu.h
@@ -104,19 +104,19 @@ __global__ void ScatterNdCUDAKernel(const T* update,
                                     size_t remain_size,
                                     size_t slice_size,
                                     size_t end_size) {
-  int64_t total_size = remain_size * slice_size;
-  int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * VecSize;
+  int total_size = remain_size * slice_size;
+  int idx = (blockIdx.x * blockDim.x + threadIdx.x) * VecSize;
   int64_t stride = blockDim.x * gridDim.x * VecSize;
 
 #pragma unroll
   for (; idx < total_size; idx += stride) {
-    int64_t indices_i = idx / slice_size;
-    int64_t slice_i = idx % slice_size;
+    int indices_i = idx / slice_size;
+    int slice_i = idx % slice_size;
     int64_t gather_i = 0;
     int64_t temp = slice_size;
 
 #pragma unroll
-    for (int64_t j = end_size - 1; j >= 0; --j) {
+    for (int j = end_size - 1; j >= 0; --j) {
       IndexT index_value = indices[indices_i * end_size + j];
       PADDLE_ENFORCE(
           index_value >= -output_dims[j] && index_value < output_dims[j],
@@ -270,72 +270,6 @@ void GPUScatterGradForX(const phi::GPUContext& ctx,
       p_index, p_output, dst_dims[0], index_size, slice_size);
 }
 
-template <typename T, typename IndexT, int VecSize>
-void DispatchScatterNdKernel(
-    const phi::GPUContext& ctx,
-    const T* p_update,
-    const IndexT* p_index,
-    T* p_output,
-    const Dim<DDim::kMaxRank>& g_output_dims,
-    int64_t remain_numel,
-    int64_t slice_size,
-    int64_t end_size,
-    int vec_size,
-    const phi::backends::gpu::GpuLaunchConfig& config) {
-  if (vec_size == VecSize) {
-    auto stream = ctx.stream();
-    ScatterNdCUDAKernel<T, IndexT, VecSize>
-        <<<config.block_per_grid, config.thread_per_block, 0, stream>>>(
-            p_update,
-            p_index,
-            p_output,
-            g_output_dims,
-            remain_numel,
-            slice_size,
-            end_size);
-  } else {
-    DispatchScatterNdKernel<T, IndexT, VecSize / 2>(ctx,
-                                                    p_update,
-                                                    p_index,
-                                                    p_output,
-                                                    g_output_dims,
-                                                    remain_numel,
-                                                    slice_size,
-                                                    end_size,
-                                                    vec_size,
-                                                    config);
-  }
-}
-
-template <typename T, typename IndexT>
-void DispatchScatterNdKernel(
-    const phi::GPUContext& ctx,
-    const T* p_update,
-    const IndexT* p_index,
-    T* p_output,
-    const Dim<DDim::kMaxRank>& g_output_dims,
-    int64_t remain_numel,
-    int64_t slice_size,
-    int64_t end_size,
-    int vec_size,
-    const phi::backends::gpu::GpuLaunchConfig& config) {
-  if (vec_size == 1) {
-    auto stream = ctx.stream();
-    ScatterNdCUDAKernel<T, IndexT, 1>
-        <<<config.block_per_grid, config.thread_per_block, 0, stream>>>(
-            p_update,
-            p_index,
-            p_output,
-            g_output_dims,
-            remain_numel,
-            slice_size,
-            end_size);
-  } else {
-    PADDLE_THROW(common::errors::Unimplemented(
-        "Unsupported vectorized size: %d", vec_size));
-  }
-}
-
 template <typename T, typename IndexT = int>
 void GPUScatterNdAdd(const phi::GPUContext& ctx,
                      const DenseTensor& update,
@@ -379,16 +313,29 @@ void GPUScatterNdAdd(const phi::GPUContext& ctx,
   auto config = phi::backends::gpu::GetGpuLaunchConfig1D(
       ctx, remain_numel * slice_size, vec_size * loop_count);
 
-  DispatchScatterNdKernel<T, IndexT, 4>(ctx,
-                                        p_update,
-                                        p_index,
-                                        p_output,
-                                        g_output_dims,
-                                        remain_numel,
-                                        slice_size,
-                                        end_size,
-                                        vec_size,
-                                        config);
+  auto stream = ctx.stream();
+
+  switch (vec_size) {
+#define CASE_VEC_SIZE(__Sz)                                              \
+  case __Sz:                                                             \
+    ScatterNdCUDAKernel<T, IndexT, __Sz>                                 \
+        <<<config.block_per_grid, config.thread_per_block, 0, stream>>>( \
+            p_update,                                                    \
+            p_index,                                                     \
+            p_output,                                                    \
+            g_output_dims,                                               \
+            remain_numel,                                                \
+            slice_size,                                                  \
+            end_size);                                                   \
+    break
+    CASE_VEC_SIZE(4);
+    CASE_VEC_SIZE(2);
+    CASE_VEC_SIZE(1);
+#undef CASE_VEC_SIZE
+    default:
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Unsupported vectorized size: %d", vec_size));
+  }
 }
 
 }  // namespace funcs
