[PHI] Fix paddle.take_along_axis for big tensor (#73342)

lshpku · web-flow · commit eb9af7cf6997 · 2025-06-15T16:28:31.000+08:00
diff --git a/paddle/phi/kernels/funcs/gather_scatter_functor.cu b/paddle/phi/kernels/funcs/gather_scatter_functor.cu
@@ -107,17 +107,17 @@ __global__ void ScatterAssignGPUKernel(tensor_t* self_data,
                                        int dim,
                                        const index_t* index_data,
                                        tensor_t* src_data,
-                                       int select_dim_size,
-                                       int self_select_dim_size,
-                                       int src_select_dim_size,
+                                       int64_t select_dim_size,
+                                       int64_t self_select_dim_size,
+                                       int64_t src_select_dim_size,
                                        int64_t outer_dim_size,
                                        int64_t outer_dim_size_self,
                                        int64_t outer_dim_size_src,
                                        int64_t numel,
                                        int64_t numel_data,
                                        const func_t& reduce_op,
                                        int* thread_ids) {
-  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + static_cast<int64_t>(blockIdx.x) * blockDim.x;
   if (tid >= numel) return;
   int64_t i, j, k;  // The i, j, k here is the index of the 3 layers loop
                     // squeezed from the N layers loop.
@@ -199,9 +199,9 @@ __global__ void GatherScatterGPUKernel(tensor_t* self_data,
                                        int dim,
                                        const index_t* index_data,
                                        tensor_t* src_data,
-                                       int select_dim_size,
-                                       int self_select_dim_size,
-                                       int src_select_dim_size,
+                                       int64_t select_dim_size,
+                                       int64_t self_select_dim_size,
+                                       int64_t src_select_dim_size,
                                        int64_t outer_dim_size,
                                        int64_t outer_dim_size_self,
                                        int64_t outer_dim_size_src,
@@ -305,9 +305,9 @@ __global__ void ScatterMeanGPUKernel(tensor_t* self_data,
                                      int dim,
                                      const index_t* index_data,
                                      tensor_t* src_data,
-                                     int select_dim_size,
-                                     int self_select_dim_size,
-                                     int src_select_dim_size,
+                                     int64_t select_dim_size,
+                                     int64_t self_select_dim_size,
+                                     int64_t src_select_dim_size,
                                      int64_t outer_dim_size,
                                      int64_t outer_dim_size_self,
                                      int64_t outer_dim_size_src,
@@ -316,7 +316,7 @@ __global__ void ScatterMeanGPUKernel(tensor_t* self_data,
                                      bool include_self,
                                      const func_t& reduce_op,
                                      int* shared_mem) {
-  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + static_cast<int64_t>(blockIdx.x) * blockDim.x;
   if (tid >= numel) return;
 
   int64_t i, j, k;  // The i, j, k here is the index of the 3 layers loop
@@ -425,10 +425,10 @@ struct gpu_gather_scatter_functor {
     auto index_dims = index.dims();
     auto src_dims = src.dims();
     if (self_size == 0 || src_size == 0 || index_size == 0) return;
-    int select_dim_size = index_dims[dim];
+    int64_t select_dim_size = index_dims[dim];
     // index matrix has different shape with self matrix or src matrix.
-    int self_select_dim_size = self_dims[dim];
-    int src_select_dim_size = src_dims[dim];
+    int64_t self_select_dim_size = self_dims[dim];
+    int64_t src_select_dim_size = src_dims[dim];
     int64_t outer_dim_size_self = 1;
     int64_t outer_dim_size_src = 1;
     int64_t inner_dim_size = 1;
