[PHI] fix broadcast kernel for big tensor (#73105)

liuruyan · web-flow · commit d225dfe8bfd8 · 2025-06-13T10:48:47.000+08:00
* fix bc bug of big tensor

* fix bug
diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -292,10 +292,10 @@ __device__ void VectorizedBroadcastKernelImpl(
     const Array<const _ptr_ char *__restrict__, Arity> &ins,
     Array<_ptr_ OutT *, NumOuts> outs,
     const Array<bool, Arity> &use_broadcast,
-    const uint32_t numel,
+    const int64_t numel,
     const Array<kps::details::BroadcastConfig, Arity> &configs,
-    int num,
-    int block_offset,
+    int64_t num,
+    int64_t block_offset,
     int read_lens,
     Functor func) {
   using Traits = phi::funcs::FunctionTraits<Functor>;
@@ -310,10 +310,10 @@ __device__ void VectorizedBroadcastKernelImpl(
   if (LoadType == kBroadcast) {
     uint32_t index_bc[Arity][VecSize] = {0};
     Unroller<BroadcastDataInit, VecSize, Arity>::step(args);
-    uint32_t thread_offset = block_offset + threadIdx.x * VecSize;
+    int64_t thread_offset = block_offset + threadIdx.x * VecSize;
 #pragma unroll
     for (int k = 0; k < VecSize; ++k) {
-      uint32_t idx = thread_offset + k;
+      int64_t idx = thread_offset + k;
       if (IsBoundary && idx == numel) break;
 #pragma unroll
       for (int i = 0; i < phi::DDim::kMaxRank; ++i) {
@@ -352,15 +352,15 @@ __global__ void VectorizedBroadcastKernel(
     Array<const _ptr_ char *__restrict__, Arity> ins,
     Array<_ptr_ OutT *, NumOuts> outs,
     Array<bool, Arity> use_broadcast,
-    uint32_t numel,
+    int64_t numel,
     Array<kps::details::BroadcastConfig, Arity> configs,
-    int main_offset,
-    int tail_tid,
+    int64_t main_offset,
+    int64_t tail_tid,
     int read_lens,
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
-  int block_offset = BLOCK_ID_X * BLOCK_NUM_X * read_lens;
-  int stride = BLOCK_NUM_X * GRID_NUM_X * read_lens;
+  int64_t block_offset = BLOCK_ID_X * BLOCK_NUM_X * read_lens;
+  int64_t stride = BLOCK_NUM_X * GRID_NUM_X * read_lens;
   for (; block_offset < main_offset; block_offset += stride) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
@@ -378,7 +378,7 @@ __global__ void VectorizedBroadcastKernel(
                                             read_lens,
                                             func);
   }
-  int num = numel - block_offset;
+  int64_t num = numel - block_offset;
   if (num > 0) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
@@ -397,7 +397,7 @@ __global__ void VectorizedBroadcastKernel(
                                             func);
   }
 #else
-  int block_offset = BLOCK_ID_X * BLOCK_NUM_X * VecSize;
+  int64_t block_offset = BLOCK_ID_X * BLOCK_NUM_X * VecSize;
   if (block_offset < main_offset) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
@@ -440,13 +440,13 @@ void LaunchBroadcastKernel(
     const BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts> &classifier,
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
-  int numel = classifier.numel;
+  int64_t numel = classifier.numel;
   const int threads = 64;
   const int blocks = 8;
   int read_lens = configs[0].buf_len;
   auto stream = ctx.x_context()->xpu_stream;
-  int main_offset = (numel / (read_lens * threads)) * read_lens * threads;
-  int tail_tid = numel % (read_lens * threads);
+  int64_t main_offset = (numel / (read_lens * threads)) * read_lens * threads;
+  int64_t tail_tid = numel % (read_lens * threads);
 
   VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, false>
       <<<blocks, threads, 0, stream>>>(classifier.ins_data,
@@ -465,8 +465,8 @@ void LaunchBroadcastKernel(
   auto stream = ctx.stream();
   auto threads = gpu_config.GetBlockSize();
   auto blocks = gpu_config.block_per_grid;
-  int main_offset = (numel / (VecSize * threads)) * VecSize * threads;
-  int tail_tid = numel % (VecSize * threads);
+  int64_t main_offset = (numel / (VecSize * threads)) * VecSize * threads;
+  int64_t tail_tid = numel % (VecSize * threads);
 
   if (classifier.all_elementwise) {
     VectorizedBroadcastKernel<Functor,
diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h
@@ -40,8 +40,8 @@ struct alignas(sizeof(T) * VecSize) VectorType {
  * must be [dim1, dim0].
  */
 struct BroadcastConfig {
-  phi::funcs::FastDivMod<int> divmoders[phi::DDim::kMaxRank];
-  uint32_t strides[phi::DDim::kMaxRank];
+  phi::funcs::FastDivMod<int64_t> divmoders[phi::DDim::kMaxRank];
+  uint64_t strides[phi::DDim::kMaxRank];
   int rank{0};
 
   // BroadcastConfig should be defined on host used on device.
@@ -51,7 +51,7 @@ struct BroadcastConfig {
                   const std::vector<int64_t>& in_dims,
                   int dim_size) {
     for (int i = 0; i < dim_size; ++i) {
-      divmoders[i] = phi::funcs::FastDivMod<int>(out_dims[i]);
+      divmoders[i] = phi::funcs::FastDivMod<int64_t>(out_dims[i]);
     }
 
     for (int i = 0; i < dim_size; ++i) {
