refine big tensor index use (#72276)

wanghuancoder · web-flow · commit 35c03d2ed48c · 2025-04-16T14:02:25.000+08:00
diff --git a/paddle/phi/kernels/gpu/arg_min_max_kernel.cu b/paddle/phi/kernels/gpu/arg_min_max_kernel.cu
@@ -55,22 +55,26 @@ using KeyValuePair = cub::KeyValuePair<K, V>;
   FIXED_BLOCK_DIM_CASE_BASE(4, ##__VA_ARGS__);  \
   FIXED_BLOCK_DIM_CASE_BASE(3, ##__VA_ARGS__);
 
-template <typename T, typename IndType, class Reducer, size_t BlockDim>
+template <typename T,
+          typename IndType,
+          class Reducer,
+          size_t BlockDim,
+          typename IndexType>
 __global__ void ArgCUDAKernel(const int64_t height,     // n * h
                               const int64_t width,      // c
                               const int64_t post_size,  // h
                               const Reducer reducer,
                               const T init,
                               const T* in,
                               IndType* out) {
-  typedef cub::BlockReduce<KeyValuePair<int64_t, T>, BlockDim> BlockReduce;
+  typedef cub::BlockReduce<KeyValuePair<IndexType, T>, BlockDim> BlockReduce;
   __shared__ typename BlockReduce::TempStorage temp_storage;
 
-  for (int64_t idx = blockIdx.x; idx < height; idx += gridDim.x) {
-    KeyValuePair<int64_t, T> kv_pair = {-1, init};
-    int64_t h = idx / post_size;
-    int64_t w = idx % post_size;
-    for (int64_t k = threadIdx.x; k < width; k += blockDim.x) {
+  for (IndexType idx = blockIdx.x; idx < height; idx += gridDim.x) {
+    KeyValuePair<IndexType, T> kv_pair = {-1, init};
+    IndexType h = idx / post_size;
+    IndexType w = idx % post_size;
+    for (IndexType k = threadIdx.x; k < width; k += blockDim.x) {
       kv_pair =
           reducer({k, in[h * width * post_size + k * post_size + w]}, kv_pair);
     }
@@ -82,7 +86,7 @@ __global__ void ArgCUDAKernel(const int64_t height,     // n * h
   }
 }
 
-template <typename T, typename IndType, class Reducer>
+template <typename T, typename IndType, class Reducer, typename IndexType>
 void ComputeFullArg(const phi::GPUContext& dev_ctx,
                     const DenseTensor& input,
                     DenseTensor* indices,
@@ -119,27 +123,29 @@ void ComputeFullArg(const phi::GPUContext& dev_ctx,
 
   if (typeid(Reducer) == typeid(cub::ArgMax)) {
     switch (ComputeBlockSize(width)) {
-      FIXED_BLOCK_DIM_CASE(ArgCUDAKernel<T, IndType, Reducer, kBlockDim>
-                           <<<grid_size, kBlockDim, 0, cu_stream>>>(
-                               height,
-                               width,
-                               post,
-                               Reducer(),
-                               std::numeric_limits<T>::lowest(),
-                               in_data,
-                               out_data));
+      FIXED_BLOCK_DIM_CASE(
+          ArgCUDAKernel<T, IndType, Reducer, kBlockDim, IndexType>
+          <<<grid_size, kBlockDim, 0, cu_stream>>>(
+              height,
+              width,
+              post,
+              Reducer(),
+              std::numeric_limits<T>::lowest(),
+              in_data,
+              out_data));
     }
   } else {
     switch (ComputeBlockSize(width)) {
-      FIXED_BLOCK_DIM_CASE(ArgCUDAKernel<T, IndType, Reducer, kBlockDim>
-                           <<<grid_size, kBlockDim, 0, cu_stream>>>(
-                               height,
-                               width,
-                               post,
-                               Reducer(),
-                               std::numeric_limits<T>::max(),
-                               in_data,
-                               out_data));
+      FIXED_BLOCK_DIM_CASE(
+          ArgCUDAKernel<T, IndType, Reducer, kBlockDim, IndexType>
+          <<<grid_size, kBlockDim, 0, cu_stream>>>(
+              height,
+              width,
+              post,
+              Reducer(),
+              std::numeric_limits<T>::max(),
+              in_data,
+              out_data));
     }
   }
 }
@@ -199,7 +205,13 @@ struct VisitDataCudaArgMinMaxFunctor {
       post *= x_dims[i];
     }
 
-    ComputeFullArg<T, IndType, Reducer>(dev_ctx, x, out, pre, post, n);
+    if (x_dims[new_axis] > std::numeric_limits<int32_t>::max()) {
+      ComputeFullArg<T, IndType, Reducer, int64_t>(
+          dev_ctx, x, out, pre, post, n);
+    } else {
+      ComputeFullArg<T, IndType, Reducer, int32_t>(
+          dev_ctx, x, out, pre, post, n);
+    }
   }
 };
 
