[release/3.4][Operator Mechanism] fix BroadcastConfig do not support big_tensor

paddle/phi/kernels/funcs/broadcast_function.h

@@ -15,6 +15,7 @@ limitations under the License. */
 #pragma once
 
 #include <sstream>
+#include "paddle/common/enforce.h"
 #include "paddle/phi/kernels/funcs/elementwise_base.h"
 
 #if defined(__NVCC__) || defined(__HIPCC__) || defined(__xpu__)
@@ -47,7 +48,12 @@ struct BroadcastTypeClassifier {
                           int axis) {
     numel = (*outs)[0]->numel();
 
-#ifndef PADDLE_WITH_XPU_KP
+#ifdef PADDLE_WITH_XPU_KP
+    // datamover_primitives_xpu2.h::BroadcastConfig (built inside
+    // InitBroadcastConfigs below) still computes strides/numel with 32-bit
+    // int, so the INT_MAX check must happen before that call, not after.
+    PADDLE_ENFORCE_LE_INT_MAX(numel, "BroadcastKernel numel (XPU)");
+#else
     for (size_t i = 0; i < ins.size(); ++i) {
       bool is_same_dim = ins[i]->numel() == numel;
       if (is_same_dim) {
@@ -120,9 +126,9 @@ struct BroadcastDataLoader {
                                                ArgsT *args,
                                                const Array2 &configs,
                                                const Array3 &use_broadcast,
-                                               const int block_offset,
+                                               uint64_t block_offset,
                                                const int num,
-                                               const uint32_t numel,
+                                               uint64_t numel,
                                                int read_lens) {
     using Type = std::tuple_element_t<Index, ArgsT>;
 #ifdef PADDLE_WITH_XPU_KP
@@ -178,16 +184,17 @@ struct BroadcastDataLoader<Index, VecSize, true, kElementwise> {
                                                ArgsT *args,
                                                const Array2 &configs,
                                                const Array3 &use_broadcast,
-                                               const int block_offset,
+                                               uint64_t block_offset,
                                                const int num,
-                                               const uint32_t numel,
+                                               uint64_t numel,
                                                int read_lens) {
     using Type = std::tuple_element_t<Index, ArgsT>;
-    int thread_offset = threadIdx.x * VecSize + block_offset;
+    uint64_t thread_offset =
+        static_cast<uint64_t>(threadIdx.x) * VecSize + block_offset;
 #pragma unroll
     for (int idx = 0; idx < VecSize; ++idx) {
       std::get<Index>(args[idx]) = static_cast<Type>(1);
-      int index = thread_offset + idx;
+      uint64_t index = thread_offset + idx;
       if (index < numel) {
         std::get<Index>(args[idx]) =
             reinterpret_cast<const _ptr_ Type *>(ins[Index])[index];
@@ -204,9 +211,9 @@ struct BroadcastDataLoader<Index, VecSize, false, kElementwise> {
                                                ArgsT *args,
                                                const Array2 &configs,
                                                const Array3 &use_broadcast,
-                                               const int block_offset,
+                                               uint64_t block_offset,
                                                const int num,
-                                               const uint32_t numel,
+                                               uint64_t numel,
                                                int read_lens) {
     using Type = std::tuple_element_t<Index, ArgsT>;
     using VecType = phi::kps::details::VectorType<Type, VecSize>;
@@ -242,7 +249,7 @@ struct BroadcastDataSetter {
   template <typename Array, typename ArgsT>
   static __device__ __forceinline__ void Apply(const Array &ins,
                                                ArgsT *args,
-                                               uint32_t index_bc[][VecSize]) {
+                                               uint64_t index_bc[][VecSize]) {
     using Type = std::tuple_element_t<Index, ArgsT>;
 #pragma unroll
     for (int k = 0; k < VecSize; ++k) {
@@ -294,10 +301,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 uint64_t numel,
     const Array<kps::details::BroadcastConfig, Arity> &configs,
     uint32_t num,
-    uint32_t block_offset,
+    uint64_t block_offset,
     int read_lens,
     Functor func) {
   using Traits = funcs::FunctionTraits<Functor>;
@@ -310,12 +317,13 @@ __device__ void VectorizedBroadcastKernelImpl(
       ins, args, configs, use_broadcast, block_offset, num, numel, read_lens);
 #else
   if (LoadType == kBroadcast) {
-    uint32_t index_bc[Arity][VecSize] = {0};
+    uint64_t index_bc[Arity][VecSize] = {0};
     Unroller<BroadcastDataInit, VecSize, Arity>::step(args);
-    uint32_t thread_offset = block_offset + threadIdx.x * VecSize;
+    uint64_t thread_offset =
+        block_offset + static_cast<uint64_t>(threadIdx.x) * VecSize;
 #pragma unroll
     for (int k = 0; k < VecSize; ++k) {
-      uint32_t idx = thread_offset + k;
+      uint64_t idx = thread_offset + k;
       if (IsBoundary && idx == numel) break;
 #pragma unroll
       for (int i = 0; i < DDim::kMaxRank; ++i) {
@@ -353,15 +361,16 @@ __global__ void VectorizedBroadcastKernel(
     Array<const _ptr_ char *__restrict__, Arity> ins,
     Array<_ptr_ OutT *, NumOuts> outs,
     Array<bool, Arity> use_broadcast,
-    uint32_t numel,
+    uint64_t numel,
     Array<kps::details::BroadcastConfig, Arity> configs,
-    uint32_t main_offset,
-    uint32_t tail_tid,
+    uint64_t main_offset,
+    uint64_t tail_tid,
     int read_lens,
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
-  int64_t block_offset = BLOCK_ID_X * BLOCK_NUM_X * read_lens;
-  int64_t stride = BLOCK_NUM_X * GRID_NUM_X * read_lens;
+  uint64_t block_offset =
+      static_cast<uint64_t>(BLOCK_ID_X) * BLOCK_NUM_X * read_lens;
+  uint64_t stride = static_cast<uint64_t>(BLOCK_NUM_X) * GRID_NUM_X * read_lens;
   for (; block_offset < main_offset; block_offset += stride) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
@@ -379,8 +388,8 @@ __global__ void VectorizedBroadcastKernel(
                                             read_lens,
                                             func);
   }
-  int64_t num = numel - block_offset;
-  if (num > 0) {
+  if (block_offset < numel) {
+    uint64_t num = numel - block_offset;
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
                                   Arity,
@@ -398,7 +407,8 @@ __global__ void VectorizedBroadcastKernel(
                                             func);
   }
 #else
-  int64_t block_offset = BLOCK_ID_X * BLOCK_NUM_X * VecSize;
+  uint64_t block_offset =
+      static_cast<uint64_t>(BLOCK_ID_X) * BLOCK_NUM_X * VecSize;
   if (block_offset < main_offset) {
     VectorizedBroadcastKernelImpl<OutT,
                                   Functor,
@@ -441,13 +451,14 @@ void LaunchBroadcastKernel(
     const BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts> &classifier,
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
-  int numel = classifier.numel;
+  const int64_t numel = classifier.numel;
   const int threads = 64;
   const int blocks = 8;
   int read_lens = configs[0].buf_len;
   auto stream = dev_ctx.x_context()->xpu_stream;
-  uint32_t main_offset = (numel / (read_lens * threads)) * read_lens * threads;
-  uint32_t tail_tid = numel % (read_lens * threads);
+  const int64_t block_len = static_cast<int64_t>(read_lens) * threads;
+  const int64_t main_offset = (numel / block_len) * block_len;
+  const int64_t tail_tid = numel % block_len;
 
   VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, false>
       <<<blocks, threads, 0, stream>>>(classifier.ins_data,
@@ -460,14 +471,15 @@ void LaunchBroadcastKernel(
                                        read_lens,
                                        func);
 #else
-  const int64_t &numel = classifier.numel;
+  const int64_t numel = classifier.numel;
   auto gpu_config =
       phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, VecSize);
   auto stream = dev_ctx.stream();
   auto threads = gpu_config.GetBlockSize();
   auto blocks = gpu_config.block_per_grid;
-  uint32_t main_offset = (numel / (VecSize * threads)) * VecSize * threads;
-  uint32_t tail_tid = numel % (VecSize * threads);
+  const int64_t block_len = static_cast<int64_t>(VecSize) * threads;
+  const int64_t main_offset = (numel / block_len) * block_len;
+  const int64_t tail_tid = numel % block_len;
 
   if (classifier.all_elementwise) {
     VectorizedBroadcastKernel<Functor,

paddle/phi/kernels/funcs/elementwise_base.h

@@ -679,7 +679,7 @@ struct ElementwiseWriteDataCallerBc {
   __device__ __forceinline__ void operator()(
       Array<_ptr_ OutT *, NumOuts> outs,
       ConditionalT<OutT, NumOuts> src[VecSize],
-      kps::IndexType block_offset,
+      int64_t block_offset,
       int num,
       int read_lens) {
     OutT dst[NumOuts][VecSize];
@@ -702,7 +702,7 @@ template <typename OutT, int VecSize, bool IsBoundary>
 struct ElementwiseWriteDataCallerBc<OutT, VecSize, IsBoundary, 1> {
   __device__ __forceinline__ void operator()(Array<_ptr_ OutT *, 1> outs,
                                              OutT src[VecSize],
-                                             kps::IndexType block_offset,
+                                             int64_t block_offset,
                                              int num,
                                              int read_lens) {
     kps::WriteData<OutT, VecSize, 1, IsBoundary>(

paddle/phi/kernels/primitive/datamover_primitives.h

@@ -40,7 +40,7 @@ struct alignas(sizeof(T) * VecSize) VectorType {
  * must be [dim1, dim0].
  */
 struct BroadcastConfig {
-  funcs::FastDivMod<int> divmoders[DDim::kMaxRank];
+  funcs::FastDivMod<int64_t> divmoders[DDim::kMaxRank];
   uint64_t strides[DDim::kMaxRank];
   int rank{0};
 
@@ -51,15 +51,15 @@ struct BroadcastConfig {
                   const std::vector<int64_t>& in_dims,
                   int dim_size) {
     for (int i = 0; i < dim_size; ++i) {
-      divmoders[i] = funcs::FastDivMod<int>(out_dims[i]);
+      divmoders[i] = funcs::FastDivMod<int64_t>(out_dims[i]);
     }
 
     for (int i = 0; i < dim_size; ++i) {
       strides[i] = in_dims[i] == 1 ? 0 : 1;
       strides[i] = (i != 0 && strides[i] != 0)
                        ? std::accumulate(in_dims.begin(),
                                          in_dims.begin() + i,
-                                         1,
+                                         int64_t{1},
                                          std::multiplies<int64_t>())
                        : strides[i];
     }
@@ -412,19 +412,21 @@ template <typename T, int NX, int NY, bool IsBoundary = false>
 __device__ __forceinline__ void ReadDataBc(
     T* dst,
     const T* __restrict__ src,
-    uint32_t block_offset,
+    uint64_t block_offset,
     const details::BroadcastConfig& config,
-    int total_num_output,
+    uint64_t total_num_output,
     int stride_nx,
     int stride_ny) {
-  uint32_t thread_offset = block_offset + threadIdx.x;
-  uint32_t index_src = 0;
+  uint64_t thread_offset = block_offset + static_cast<uint64_t>(threadIdx.x);
+  uint64_t index_src = 0;
 
 #pragma unroll
   for (int ny = 0; ny < NY; ++ny) {
 #pragma unroll
     for (uint32_t nx = 0; nx < NX; ++nx) {
-      uint32_t index_output = thread_offset + ny * stride_ny + nx * stride_nx;
+      uint64_t index_output = thread_offset +
+                              static_cast<uint64_t>(ny) * stride_ny +
+                              static_cast<uint64_t>(nx) * stride_nx;
       index_src = 0;
       if (IsBoundary) {
         if (index_output >= total_num_output) {
@@ -748,16 +750,17 @@ template <typename T, int NX, int NY, bool IsBoundary = false>
 __device__ __forceinline__ void ReadDataBc(
     T* dst,
     const T* __restrict__ src,
-    uint32_t block_offset,
+    uint64_t block_offset,
     const details::BroadcastConfig& config,
-    int total_num_output,
+    uint64_t total_num_output,
     int read_lens = NX) {
-  uint32_t thread_offset = block_offset + threadIdx.x * NX;
-  uint32_t index_src = 0;
+  uint64_t thread_offset =
+      block_offset + static_cast<uint64_t>(threadIdx.x) * NX;
+  uint64_t index_src = 0;
 
 #pragma unroll
   for (uint32_t nx = 0; nx < NX; ++nx) {
-    uint32_t index_output = thread_offset + nx;
+    uint64_t index_output = thread_offset + nx;
     index_src = 0;
     if (IsBoundary) {
       if (index_output >= total_num_output) {
@@ -809,16 +812,17 @@ template <typename T,
 __device__ __forceinline__ void ReadDataBc(
     ArgsT* dst,
     const T* __restrict__ src,
-    uint32_t block_offset,
+    uint64_t block_offset,
     const details::BroadcastConfig& config,
-    int total_num_output,
+    uint64_t total_num_output,
     int read_lens = NX) {
-  uint32_t thread_offset = block_offset + threadIdx.x * NX;
-  uint32_t index_src = 0;
+  uint64_t thread_offset =
+      block_offset + static_cast<uint64_t>(threadIdx.x) * NX;
+  uint64_t index_src = 0;
 
 #pragma unroll
   for (uint32_t nx = 0; nx < NX; ++nx) {
-    uint32_t index_output = thread_offset + nx;
+    uint64_t index_output = thread_offset + nx;
     index_src = 0;
     if (IsBoundary) {
       if (index_output >= total_num_output) {