fix performence bug (#73501)

liuruyan · web-flow · commit 1dcdc2439ec1 · 2025-06-23T00:29:01.000+08:00
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 int64_t numel,
+    const uint32_t numel,
     const Array<kps::details::BroadcastConfig, Arity> &configs,
-    int64_t num,
-    int64_t block_offset,
+    int num,
+    int 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);
-    int64_t thread_offset = block_offset + threadIdx.x * VecSize;
+    uint32_t thread_offset = block_offset + threadIdx.x * VecSize;
 #pragma unroll
     for (int k = 0; k < VecSize; ++k) {
-      int64_t idx = thread_offset + k;
+      uint32_t idx = thread_offset + k;
       if (IsBoundary && idx == numel) break;
 #pragma unroll
       for (int i = 0; i < phi::DDim::kMaxRank; ++i) {
@@ -352,10 +352,10 @@ __global__ void VectorizedBroadcastKernel(
     Array<const _ptr_ char *__restrict__, Arity> ins,
     Array<_ptr_ OutT *, NumOuts> outs,
     Array<bool, Arity> use_broadcast,
-    int64_t numel,
+    uint32_t numel,
     Array<kps::details::BroadcastConfig, Arity> configs,
-    int64_t main_offset,
-    int64_t tail_tid,
+    int main_offset,
+    int tail_tid,
     int read_lens,
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
@@ -440,13 +440,13 @@ void LaunchBroadcastKernel(
     const BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts> &classifier,
     Functor func) {
 #ifdef PADDLE_WITH_XPU_KP
-  int64_t numel = classifier.numel;
+  int 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;
-  int64_t main_offset = (numel / (read_lens * threads)) * read_lens * threads;
-  int64_t tail_tid = numel % (read_lens * threads);
+  int main_offset = (numel / (read_lens * threads)) * read_lens * threads;
+  int tail_tid = numel % (read_lens * threads);
 
   VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, false>
       <<<blocks, threads, 0, stream>>>(classifier.ins_data,
@@ -459,14 +459,14 @@ void LaunchBroadcastKernel(
                                        read_lens,
                                        func);
 #else
-  const auto &numel = classifier.numel;
+  const int &numel = classifier.numel;
   auto gpu_config =
       phi::backends::gpu::GetGpuLaunchConfig1D(ctx, numel, VecSize);
   auto stream = ctx.stream();
   auto threads = gpu_config.GetBlockSize();
   auto blocks = gpu_config.block_per_grid;
-  int64_t main_offset = (numel / (VecSize * threads)) * VecSize * threads;
-  int64_t tail_tid = numel % (VecSize * threads);
+  int main_offset = (numel / (VecSize * threads)) * VecSize * threads;
+  int 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,7 +40,7 @@ struct alignas(sizeof(T) * VecSize) VectorType {
  * must be [dim1, dim0].
  */
 struct BroadcastConfig {
-  phi::funcs::FastDivMod<int64_t> divmoders[phi::DDim::kMaxRank];
+  phi::funcs::FastDivMod<int> divmoders[phi::DDim::kMaxRank];
   uint64_t strides[phi::DDim::kMaxRank];
   int rank{0};
 
@@ -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<int64_t>(out_dims[i]);
+      divmoders[i] = phi::funcs::FastDivMod<int>(out_dims[i]);
     }
 
     for (int i = 0; i < dim_size; ++i) {
