add a inner loop for index_select_grad_init() in index_select op when dealing...

add a inner loop for index_select_grad_init() in index_select op when dealing with large-shape data (#41563)

* replace for with CUDA_KERNEL_LOOP for index_select_grad_init() in index_select op

* use CUDA_KERNEL_LOOP_TYPE

* fix code style

* replace index_select_grad_init with SetConstant

paddle/phi/kernels/funcs/gather.cu.h

@@ -17,6 +17,7 @@ limitations under the License. */
 #include <vector>
 #include "paddle/fluid/memory/memcpy.h"
 // TODO(paddle-dev): move gpu_primitives.h to phi
+#include "paddle/fluid/platform/device/gpu/gpu_launch_config.h"
 #include "paddle/fluid/platform/device/gpu/gpu_primitives.h"
 #include "paddle/phi/backends/gpu/gpu_launch_config.h"
 #include "paddle/phi/common/place.h"
@@ -110,11 +111,8 @@ void GPUGather(const phi::GPUContext& ctx,
 
   int block = 512;
   int64_t n = slice_size * index_size;
-  int64_t grid = (n + block - 1) / block;
-  unsigned int maxGridDimX = ctx.GetCUDAMaxGridDimSize()[0];
-  if (grid > maxGridDimX) {
-    grid = maxGridDimX;
-  }
+  dim3 grid = dim3((n + block - 1) / block);
+  paddle::platform::LimitGridDim(ctx, &grid);
 
   GatherCUDAKernel<T, IndexT><<<grid, block, 0, ctx.stream()>>>(
       p_src, p_index, p_output, index_size, slice_size);
@@ -155,11 +153,8 @@ void GPUGatherNd(const phi::GPUContext& ctx,
 
   int block = 512;
   int64_t n = slice_size * remain_numel;
-  int64_t grid = (n + block - 1) / block;
-  unsigned int maxGridDimX = ctx.GetCUDAMaxGridDimSize()[0];
-  if (grid > maxGridDimX) {
-    grid = maxGridDimX;
-  }
+  dim3 grid = dim3((n + block - 1) / block);
+  paddle::platform::LimitGridDim(ctx, &grid);
 
   GatherNdCUDAKernel<T, IndexT><<<grid, block, 0, ctx.stream()>>>(p_input,
                                                                   g_input_dims,

paddle/phi/kernels/funcs/scatter.cu.h

@@ -15,6 +15,7 @@ limitations under the License. */
 #pragma once
 #include <unordered_set>
 #include <vector>
+#include "paddle/fluid/platform/device/gpu/gpu_launch_config.h"
 #include "paddle/fluid/platform/device/gpu/gpu_primitives.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/dense_tensor.h"
@@ -155,9 +156,8 @@ void GPUScatterAssign(const phi::GPUContext& ctx,
   // set block and grid num
   int block = 512;
   int64_t n = slice_size * index_size;
-  int64_t grid = (n + block - 1) / block;
-  unsigned int maxGridDimX = ctx.GetCUDAMaxGridDimSize()[0];
-  grid = grid > maxGridDimX ? maxGridDimX : grid;
+  dim3 grid = dim3((n + block - 1) / block);
+  paddle::platform::LimitGridDim(ctx, &grid);
 
   // if not overwrite mode, init data
   if (!overwrite) {
@@ -188,9 +188,8 @@ void GPUScatterGradForX(const phi::GPUContext& ctx,
   int64_t block = 512;
   int64_t n = slice_size * index_size;
   int64_t height = (n + block - 1) / block;
-
-  int64_t max_grid_dimx = ctx.GetCUDAMaxGridDimSize()[0];
-  int64_t grid = height < max_grid_dimx ? height : max_grid_dimx;
+  dim3 grid = dim3((n + block - 1) / block);
+  paddle::platform::LimitGridDim(ctx, &grid);
 
   ScatterInitCUDAKernel<T, IndexT><<<grid, block, 0, ctx.stream()>>>(
       p_index, p_output, index_size, slice_size);
@@ -230,9 +229,8 @@ void GPUScatterNdAdd(const phi::GPUContext& ctx,
 
   int block = 512;
   int64_t n = slice_size * remain_numel;
-  int64_t grid = (n + block - 1) / block;
-  unsigned int maxGridDimX = ctx.GetCUDAMaxGridDimSize()[0];
-  grid = grid > maxGridDimX ? maxGridDimX : grid;
+  dim3 grid = dim3((n + block - 1) / block);
+  paddle::platform::LimitGridDim(ctx, &grid);
 
   ScatterNdCUDAKernel<T, IndexT><<<grid, block, 0, ctx.stream()>>>(
       p_update,

paddle/phi/kernels/gpu/index_select_grad_kernel.cu

@@ -19,6 +19,7 @@
 #include "paddle/phi/backends/gpu/gpu_info.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/utils/data_type.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
 
 DECLARE_bool(cudnn_deterministic);
 
@@ -35,7 +36,7 @@ __global__ void index_select_grad_cuda_kernel(const T* output_grad,
                                               int64_t stride,
                                               int64_t size,
                                               int64_t delta) {
-  CUDA_KERNEL_LOOP(idx, N) {
+  CUDA_KERNEL_LOOP_TYPE(idx, N, int64_t) {
     int64_t pre_idx = idx / (stride * size);
     int64_t dim_idx = idx % (stride * size) / stride;
     IndexT src_dim_idx = index[dim_idx];
@@ -45,15 +46,6 @@ __global__ void index_select_grad_cuda_kernel(const T* output_grad,
   }
 }
 
-template <typename T>
-__global__ void index_select_grad_init(T* input_grad, int64_t N) {
-  int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-  if (idx >= N) {
-    return;
-  }
-  input_grad[idx] = 0.0;
-}
-
 template <typename T, typename Context>
 void IndexSelectGradKernel(const Context& ctx,
                            const DenseTensor& x,
@@ -97,8 +89,8 @@ void IndexSelectGradKernel(const Context& ctx,
   dim3 grid_dim = dim3((numel + block_dim - 1) / block_dim);
   paddle::platform::LimitGridDim(ctx, &grid_dim);
 
-  index_select_grad_init<T><<<grid_dim, block_dim, 0, stream>>>(in_grad_data,
-                                                                numel);
+  phi::funcs::SetConstant<phi::GPUContext, T> index_select_grad_init;
+  index_select_grad_init(ctx, x_grad, static_cast<T>(0));
 
   if (FLAGS_cudnn_deterministic) {
     VLOG(2) << "Run grad kernel of index_select with single thread.";

paddle/phi/kernels/gpu/index_select_kernel.cu

@@ -32,7 +32,7 @@ __global__ void index_select_cuda_kernel(const T* input,
                                          int64_t stride,
                                          int64_t size,
                                          int64_t delta) {
-  CUDA_KERNEL_LOOP(idx, N) {
+  CUDA_KERNEL_LOOP_TYPE(idx, N, int64_t) {
     int64_t pre_idx = idx / (stride * size);
     int64_t dim_idx = idx % (stride * size) / stride;
     IndexT src_dim_idx = index[dim_idx];