fix gather-scatter int32 out-of-bound (#72905)

Wennie396 · web-flow · commit 34e8aac8d5b8 · 2025-05-29T11:09:41.000+08:00
diff --git a/paddle/phi/kernels/funcs/gather_scatter_functor.cu b/paddle/phi/kernels/funcs/gather_scatter_functor.cu
@@ -94,7 +94,7 @@ class ReduceMin {
 static ReduceMin reduce_min;
 
 __global__ void CudaMemsetAsync(int* dest, int value, size_t size) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid * sizeof(int) >= size) return;
   dest[tid] = value;
 }
@@ -117,7 +117,7 @@ __global__ void ScatterAssignGPUKernel(tensor_t* self_data,
                                        int64_t numel_data,
                                        const func_t& reduce_op,
                                        int* thread_ids) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + 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.
@@ -316,7 +316,7 @@ __global__ void ScatterMeanGPUKernel(tensor_t* self_data,
                                      bool include_self,
                                      const func_t& reduce_op,
                                      int* shared_mem) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + 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
@@ -639,7 +639,7 @@ __global__ void ScatterInputGradGPUKernel(tensor_t* grad_data,
                                           int64_t outer_dim_size_data,
                                           int64_t numel,
                                           int64_t numel_data) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid >= numel) return;
   int64_t i, j, k;
   i = tid / (select_dim_size * outer_dim_size);
@@ -710,7 +710,7 @@ __global__ void ScatterMulInputGradGPUKernel(tensor_t* grad_data,
                                              int64_t numel,
                                              int64_t numel_grad,
                                              int* thread_ids) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid >= numel) return;
   int64_t i, j, k;
   i = tid / (select_dim_size * outer_dim_size);
@@ -746,7 +746,7 @@ __global__ void ScatterMinMaxInputGradGPUKernel(tensor_t* grad_data,
                                                 int64_t numel_grad,
                                                 const std::string& reduce,
                                                 int* shared_mem) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid >= numel) return;
   int64_t i, j, k;
   i = tid / (select_dim_size * outer_dim_size);
@@ -869,7 +869,7 @@ __global__ void ScatterMeanInputGradGPUKernel(tensor_t* grad_data,
                                               int64_t numel,
                                               int64_t numel_grad,
                                               int* shared_mem) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid >= numel) return;
   int64_t i, j, k;
   i = tid / (select_dim_size * outer_dim_size);
@@ -960,7 +960,7 @@ __global__ void ScatterValueGradGPUKernel(tensor_t* grad_data,
                                           int64_t numel,
                                           int64_t numel_data,
                                           int* thread_ids) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid >= numel) return;
 
   int64_t i, j, k;
@@ -1054,7 +1054,7 @@ __global__ void ScatterMeanValueGradGPUKernel(tensor_t* grad_data,
                                               int64_t numel,
                                               int64_t numel_self,
                                               int* shared_mem) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid >= numel) return;
 
   int64_t i, j, k;
@@ -1088,7 +1088,7 @@ __global__ void ScatterAddValueGradGPUKernel(tensor_t* grad_data,
                                              int64_t outer_dim_size_self,
                                              int64_t outer_dim_size_grad,
                                              int64_t numel) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid >= numel) return;
   int64_t i, j, k;
   i = tid / (select_dim_size * outer_dim_size);
@@ -1201,7 +1201,7 @@ __global__ void ScatterMulValueGradGPUKernel(tensor_t* grad_data,
                                              int64_t outer_dim_size_self,
                                              int64_t outer_dim_size_grad,
                                              int64_t numel) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid >= numel) return;
   int64_t i, j, k;
   i = tid / (select_dim_size * outer_dim_size);
@@ -1236,7 +1236,7 @@ __global__ void ScatterMinMaxValueGradGPUKernel(tensor_t* grad_data,
                                                 int64_t numel_self,
                                                 bool include_self,
                                                 int* shared_mem) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid >= numel) return;
   int64_t i, j, k;
   i = tid / (select_dim_size * outer_dim_size);
