From 21870fd679fe01a62452d3165e021b6cb74eebf4 Mon Sep 17 00:00:00 2001 From: feixi139 <18210811610@163.com> Date: Wed, 8 Jul 2026 06:15:51 +0000 Subject: [PATCH 01/17] fix bugs fix bugs fix bugs fix bugs --- .../kernels/primitive/datamover_primitives.h | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h index 92f76f1c806f82..3d4c17769f8a33 100644 --- a/paddle/phi/kernels/primitive/datamover_primitives.h +++ b/paddle/phi/kernels/primitive/datamover_primitives.h @@ -41,7 +41,7 @@ struct alignas(sizeof(T) * VecSize) VectorType { * must be [dim1, dim0]. */ struct BroadcastConfig { - funcs::FastDivMod divmoders[DDim::kMaxRank]; + funcs::FastDivMod divmoders[DDim::kMaxRank]; uint64_t strides[DDim::kMaxRank]; int rank{0}; @@ -52,8 +52,7 @@ struct BroadcastConfig { const std::vector& in_dims, int dim_size) { for (int i = 0; i < dim_size; ++i) { - PADDLE_ENFORCE_LE_INT_MAX(out_dims[i], "out_dim"); - divmoders[i] = funcs::FastDivMod(static_cast(out_dims[i])); + divmoders[i] = funcs::FastDivMod(out_dims[i]); } for (int i = 0; i < dim_size; ++i) { @@ -61,7 +60,7 @@ struct BroadcastConfig { strides[i] = (i != 0 && strides[i] != 0) ? std::accumulate(in_dims.begin(), in_dims.begin() + i, - 1, + int64_t{1}, std::multiplies()) : strides[i]; } @@ -420,13 +419,15 @@ __device__ __forceinline__ void ReadDataBc( int stride_nx, int stride_ny) { uint32_t thread_offset = block_offset + threadIdx.x; - uint32_t index_src = 0; + int64_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; + int64_t index_output = thread_offset + + static_cast(ny) * stride_ny + + static_cast(nx) * stride_nx; index_src = 0; if (IsBoundary) { if (index_output >= total_num_output) { @@ -755,11 +756,11 @@ __device__ __forceinline__ void ReadDataBc( int total_num_output, int read_lens = NX) { uint32_t thread_offset = block_offset + threadIdx.x * NX; - uint32_t index_src = 0; + int64_t index_src = 0; #pragma unroll for (uint32_t nx = 0; nx < NX; ++nx) { - uint32_t index_output = thread_offset + nx; + int64_t index_output = thread_offset + nx; index_src = 0; if (IsBoundary) { if (index_output >= total_num_output) { @@ -816,11 +817,11 @@ __device__ __forceinline__ void ReadDataBc( int total_num_output, int read_lens = NX) { uint32_t thread_offset = block_offset + threadIdx.x * NX; - uint32_t index_src = 0; + int64_t index_src = 0; #pragma unroll for (uint32_t nx = 0; nx < NX; ++nx) { - uint32_t index_output = thread_offset + nx; + int64_t index_output = thread_offset + nx; index_src = 0; if (IsBoundary) { if (index_output >= total_num_output) { From dc8c94e977ff1057a55546b9f64f4a0a328b4778 Mon Sep 17 00:00:00 2001 From: feixi139 <18210811610@163.com> Date: Thu, 9 Jul 2026 07:11:51 +0000 Subject: [PATCH 02/17] fix bugs --- .../phi/kernels/primitive/datamover_primitives.h | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h index 3d4c17769f8a33..d5c3f282148e82 100644 --- a/paddle/phi/kernels/primitive/datamover_primitives.h +++ b/paddle/phi/kernels/primitive/datamover_primitives.h @@ -419,15 +419,15 @@ __device__ __forceinline__ void ReadDataBc( int stride_nx, int stride_ny) { uint32_t thread_offset = block_offset + threadIdx.x; - int64_t index_src = 0; + uint64_t index_src = 0; #pragma unroll for (int ny = 0; ny < NY; ++ny) { #pragma unroll for (uint32_t nx = 0; nx < NX; ++nx) { - int64_t index_output = thread_offset + - static_cast(ny) * stride_ny + - static_cast(nx) * stride_nx; + uint64_t index_output = thread_offset + + static_cast(ny) * stride_ny + + static_cast(nx) * stride_nx; index_src = 0; if (IsBoundary) { if (index_output >= total_num_output) { @@ -756,11 +756,11 @@ __device__ __forceinline__ void ReadDataBc( int total_num_output, int read_lens = NX) { uint32_t thread_offset = block_offset + threadIdx.x * NX; - int64_t index_src = 0; + uint64_t index_src = 0; #pragma unroll for (uint32_t nx = 0; nx < NX; ++nx) { - int64_t index_output = thread_offset + nx; + uint64_t index_output = thread_offset + nx; index_src = 0; if (IsBoundary) { if (index_output >= total_num_output) { @@ -817,11 +817,11 @@ __device__ __forceinline__ void ReadDataBc( int total_num_output, int read_lens = NX) { uint32_t thread_offset = block_offset + threadIdx.x * NX; - int64_t index_src = 0; + uint64_t index_src = 0; #pragma unroll for (uint32_t nx = 0; nx < NX; ++nx) { - int64_t index_output = thread_offset + nx; + uint64_t index_output = thread_offset + nx; index_src = 0; if (IsBoundary) { if (index_output >= total_num_output) { From 19a96716c655eea569bfa26733c209207bcdb7cd Mon Sep 17 00:00:00 2001 From: feixi139 <18210811610@163.com> Date: Thu, 9 Jul 2026 08:04:37 +0000 Subject: [PATCH 03/17] fix bugs --- paddle/phi/kernels/primitive/datamover_primitives.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h index d5c3f282148e82..4b51fd9002b3f9 100644 --- a/paddle/phi/kernels/primitive/datamover_primitives.h +++ b/paddle/phi/kernels/primitive/datamover_primitives.h @@ -415,7 +415,7 @@ __device__ __forceinline__ void ReadDataBc( const T* __restrict__ src, uint32_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; @@ -753,7 +753,7 @@ __device__ __forceinline__ void ReadDataBc( const T* __restrict__ src, uint32_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; uint64_t index_src = 0; @@ -814,7 +814,7 @@ __device__ __forceinline__ void ReadDataBc( const T* __restrict__ src, uint32_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; uint64_t index_src = 0; From 16de6c7871e55aab0835918e9430fe41dfc05c57 Mon Sep 17 00:00:00 2001 From: feixi <66367509+feixi139@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:22:19 +0800 Subject: [PATCH 04/17] Update datamover_primitives.h --- paddle/phi/kernels/primitive/datamover_primitives.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h index 4b51fd9002b3f9..5deb29b43b6a0b 100644 --- a/paddle/phi/kernels/primitive/datamover_primitives.h +++ b/paddle/phi/kernels/primitive/datamover_primitives.h @@ -1,5 +1,5 @@ // Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. -// +// // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at From fff4a77ed4d4dccb03237641853879a77dfbcf19 Mon Sep 17 00:00:00 2001 From: feixi <66367509+feixi139@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:36:41 +0800 Subject: [PATCH 05/17] Update datamover_primitives.h --- paddle/phi/kernels/primitive/datamover_primitives.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h index 5deb29b43b6a0b..4b51fd9002b3f9 100644 --- a/paddle/phi/kernels/primitive/datamover_primitives.h +++ b/paddle/phi/kernels/primitive/datamover_primitives.h @@ -1,5 +1,5 @@ // Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. -// +// // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at From 2ee9729964cb7a0ad7661a7c461b3d206b05a275 Mon Sep 17 00:00:00 2001 From: feixi139 <18210811610@163.com> Date: Mon, 13 Jul 2026 07:51:15 +0000 Subject: [PATCH 06/17] fix bugs --- paddle/phi/kernels/funcs/broadcast_function.h | 69 +++++++++---------- paddle/phi/kernels/funcs/elementwise_base.h | 4 +- .../kernels/primitive/datamover_primitives.h | 12 ++-- 3 files changed, 39 insertions(+), 46 deletions(-) diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h index 5596e937d17cf6..63005e3c85bbde 100644 --- a/paddle/phi/kernels/funcs/broadcast_function.h +++ b/paddle/phi/kernels/funcs/broadcast_function.h @@ -119,9 +119,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; #ifdef PADDLE_WITH_XPU_KP @@ -177,16 +177,17 @@ 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; - int thread_offset = threadIdx.x * VecSize + block_offset; + uint64_t thread_offset = + static_cast(threadIdx.x) * VecSize + block_offset; #pragma unroll for (int idx = 0; idx < VecSize; ++idx) { std::get(args[idx]) = static_cast(1); - int index = thread_offset + idx; + uint64_t index = thread_offset + idx; if (index < numel) { std::get(args[idx]) = reinterpret_cast(ins[Index])[index]; @@ -203,9 +204,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; using VecType = phi::kps::details::VectorType; @@ -241,7 +242,7 @@ struct BroadcastDataSetter { template 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; #pragma unroll for (int k = 0; k < VecSize; ++k) { @@ -293,10 +294,10 @@ __device__ void VectorizedBroadcastKernelImpl( const Array &ins, Array<_ptr_ OutT *, NumOuts> outs, const Array &use_broadcast, - const uint32_t numel, + const uint64_t numel, const Array &configs, uint32_t num, - uint32_t block_offset, + uint64_t block_offset, int read_lens, Functor func) { using Traits = funcs::FunctionTraits; @@ -309,12 +310,12 @@ __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::step(args); - uint32_t thread_offset = block_offset + threadIdx.x * VecSize; + uint64_t thread_offset = block_offset + 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) { @@ -352,16 +353,16 @@ __global__ void VectorizedBroadcastKernel( Array ins, Array<_ptr_ OutT *, NumOuts> outs, Array use_broadcast, - uint32_t numel, + uint64_t numel, Array 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 = - static_cast(BLOCK_ID_X) * BLOCK_NUM_X * read_lens; - int64_t stride = static_cast(BLOCK_NUM_X) * GRID_NUM_X * read_lens; + uint64_t block_offset = + static_cast(BLOCK_ID_X) * BLOCK_NUM_X * read_lens; + uint64_t stride = static_cast(BLOCK_NUM_X) * GRID_NUM_X * read_lens; for (; block_offset < main_offset; block_offset += stride) { VectorizedBroadcastKernelImpl 0) { VectorizedBroadcastKernelImpl(BLOCK_ID_X) * BLOCK_NUM_X * VecSize; + uint64_t block_offset = + static_cast(BLOCK_ID_X) * BLOCK_NUM_X * VecSize; if (block_offset < main_offset) { VectorizedBroadcastKernelImpl &classifier, Functor func) { #ifdef PADDLE_WITH_XPU_KP - const int64_t numel_64 = classifier.numel; - PADDLE_ENFORCE_LE_UINT32_MAX(numel_64, "XPU broadcast kernel numel"); - const uint32_t numel = static_cast(numel_64); + 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; const int64_t block_len = static_cast(read_lens) * threads; - const int64_t main_offset_64 = (numel_64 / block_len) * block_len; - const int64_t tail_tid_64 = numel_64 % block_len; - const uint32_t main_offset = static_cast(main_offset_64); - const uint32_t tail_tid = static_cast(tail_tid_64); + const int64_t main_offset = (numel / block_len) * block_len; + const int64_t tail_tid = numel % block_len; VectorizedBroadcastKernel <<>>(classifier.ins_data, @@ -466,19 +463,15 @@ void LaunchBroadcastKernel( read_lens, func); #else - const int64_t numel_64 = classifier.numel; + const int64_t numel = classifier.numel; auto gpu_config = - phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel_64, VecSize); + phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, VecSize); auto stream = dev_ctx.stream(); uint32_t threads = static_cast(gpu_config.GetBlockSize()); auto blocks = gpu_config.block_per_grid; - PADDLE_ENFORCE_LE_UINT32_MAX(numel_64, "numel"); - const uint32_t numel = static_cast(numel_64); const int64_t block_len = static_cast(VecSize) * threads; - const int64_t main_offset_64 = (numel_64 / block_len) * block_len; - const int64_t tail_tid_64 = numel_64 % block_len; - uint32_t main_offset = static_cast(main_offset_64); - uint32_t tail_tid = static_cast(tail_tid_64); + const int64_t main_offset = (numel / block_len) * block_len; + const int64_t tail_tid = numel % block_len; if (classifier.all_elementwise) { VectorizedBroadcastKernel outs, ConditionalT src[VecSize], - kps::IndexType block_offset, + int64_t block_offset, int num, int read_lens) { OutT dst[NumOuts][VecSize]; @@ -701,7 +701,7 @@ template struct ElementwiseWriteDataCallerBc { __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( diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h index 4b51fd9002b3f9..596579acd75c2a 100644 --- a/paddle/phi/kernels/primitive/datamover_primitives.h +++ b/paddle/phi/kernels/primitive/datamover_primitives.h @@ -413,12 +413,12 @@ template __device__ __forceinline__ void ReadDataBc( T* dst, const T* __restrict__ src, - uint32_t block_offset, + uint64_t block_offset, const details::BroadcastConfig& config, uint64_t total_num_output, int stride_nx, int stride_ny) { - uint32_t thread_offset = block_offset + threadIdx.x; + uint64_t thread_offset = block_offset + threadIdx.x; uint64_t index_src = 0; #pragma unroll @@ -751,11 +751,11 @@ template __device__ __forceinline__ void ReadDataBc( T* dst, const T* __restrict__ src, - uint32_t block_offset, + uint64_t block_offset, const details::BroadcastConfig& config, uint64_t total_num_output, int read_lens = NX) { - uint32_t thread_offset = block_offset + threadIdx.x * NX; + uint64_t thread_offset = block_offset + threadIdx.x * NX; uint64_t index_src = 0; #pragma unroll @@ -812,11 +812,11 @@ template Date: Mon, 13 Jul 2026 08:26:25 +0000 Subject: [PATCH 07/17] fix bugs --- paddle/phi/kernels/funcs/broadcast_function.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h index 63005e3c85bbde..60da3782639cc6 100644 --- a/paddle/phi/kernels/funcs/broadcast_function.h +++ b/paddle/phi/kernels/funcs/broadcast_function.h @@ -380,8 +380,8 @@ __global__ void VectorizedBroadcastKernel( read_lens, func); } - uint64_t num = numel - block_offset; - if (num > 0) { + if (block_offset < numel) { + uint64_t num = numel - block_offset; VectorizedBroadcastKernelImpl Date: Mon, 13 Jul 2026 09:25:11 +0000 Subject: [PATCH 08/17] fix bugs --- paddle/phi/kernels/funcs/broadcast_function.h | 11 ++++++++--- paddle/phi/kernels/primitive/datamover_primitives.h | 8 +++++--- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h index 60da3782639cc6..aad10270209d2d 100644 --- a/paddle/phi/kernels/funcs/broadcast_function.h +++ b/paddle/phi/kernels/funcs/broadcast_function.h @@ -46,7 +46,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) { @@ -312,7 +317,8 @@ __device__ void VectorizedBroadcastKernelImpl( if (LoadType == kBroadcast) { uint64_t index_bc[Arity][VecSize] = {0}; Unroller::step(args); - uint64_t thread_offset = block_offset + threadIdx.x * VecSize; + uint64_t thread_offset = + block_offset + static_cast(threadIdx.x) * VecSize; #pragma unroll for (int k = 0; k < VecSize; ++k) { uint64_t idx = thread_offset + k; @@ -444,7 +450,6 @@ void LaunchBroadcastKernel( Functor func) { #ifdef PADDLE_WITH_XPU_KP const int64_t numel = classifier.numel; - PADDLE_ENFORCE_LE_INT_MAX(numel, "BroadcastKernel numel (XPU)"); const int threads = 64; const int blocks = 8; int read_lens = configs[0].buf_len; diff --git a/paddle/phi/kernels/primitive/datamover_primitives.h b/paddle/phi/kernels/primitive/datamover_primitives.h index 596579acd75c2a..d2fb686b5ef942 100644 --- a/paddle/phi/kernels/primitive/datamover_primitives.h +++ b/paddle/phi/kernels/primitive/datamover_primitives.h @@ -418,7 +418,7 @@ __device__ __forceinline__ void ReadDataBc( uint64_t total_num_output, int stride_nx, int stride_ny) { - uint64_t thread_offset = block_offset + threadIdx.x; + uint64_t thread_offset = block_offset + static_cast(threadIdx.x); uint64_t index_src = 0; #pragma unroll @@ -755,7 +755,8 @@ __device__ __forceinline__ void ReadDataBc( const details::BroadcastConfig& config, uint64_t total_num_output, int read_lens = NX) { - uint64_t thread_offset = block_offset + threadIdx.x * NX; + uint64_t thread_offset = + block_offset + static_cast(threadIdx.x) * NX; uint64_t index_src = 0; #pragma unroll @@ -816,7 +817,8 @@ __device__ __forceinline__ void ReadDataBc( const details::BroadcastConfig& config, uint64_t total_num_output, int read_lens = NX) { - uint64_t thread_offset = block_offset + threadIdx.x * NX; + uint64_t thread_offset = + block_offset + static_cast(threadIdx.x) * NX; uint64_t index_src = 0; #pragma unroll From a3cd99c67861fe7485a1b1b43f0219a06f9e1678 Mon Sep 17 00:00:00 2001 From: feixi139 <18210811610@163.com> Date: Tue, 14 Jul 2026 07:54:50 +0000 Subject: [PATCH 09/17] fix bugs --- paddle/phi/kernels/funcs/broadcast_function.h | 206 +++++++++++++----- .../kernels/primitive/datamover_primitives.h | 35 ++- 2 files changed, 164 insertions(+), 77 deletions(-) diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h index aad10270209d2d..26d5940975bc4f 100644 --- a/paddle/phi/kernels/funcs/broadcast_function.h +++ b/paddle/phi/kernels/funcs/broadcast_function.h @@ -14,6 +14,7 @@ limitations under the License. */ #pragma once +#include #include #include "paddle/common/enforce.h" #include "paddle/phi/kernels/funcs/elementwise_base.h" @@ -117,16 +118,20 @@ struct BroadcastTypeClassifier { }; // Common broadcast/elementwise Loader. -template +template struct BroadcastDataLoader { template static __device__ __forceinline__ void Apply(const Array1 &ins, ArgsT *args, const Array2 &configs, const Array3 &use_broadcast, - uint64_t block_offset, + IndexT block_offset, const int num, - uint64_t numel, + IndexT numel, int read_lens) { using Type = std::tuple_element_t; #ifdef PADDLE_WITH_XPU_KP @@ -150,7 +155,7 @@ struct BroadcastDataLoader { #else kps::Init(args, static_cast(1.0f)); if (use_broadcast[Index]) { - kps::ReadDataBc( + kps::ReadDataBc( args, reinterpret_cast(ins[Index]), block_offset, @@ -175,24 +180,24 @@ struct BroadcastDataLoader { /* BroadcastDataLoaders Partial specialization */ #ifndef PADDLE_WITH_XPU_KP // Scalar elementwise Loader with consideration of IsBoundary. -template -struct BroadcastDataLoader { +template +struct BroadcastDataLoader { template static __device__ __forceinline__ void Apply(const Array1 &ins, ArgsT *args, const Array2 &configs, const Array3 &use_broadcast, - uint64_t block_offset, + IndexT block_offset, const int num, - uint64_t numel, + IndexT numel, int read_lens) { using Type = std::tuple_element_t; - uint64_t thread_offset = - static_cast(threadIdx.x) * VecSize + block_offset; + IndexT thread_offset = + static_cast(threadIdx.x) * VecSize + block_offset; #pragma unroll for (int idx = 0; idx < VecSize; ++idx) { std::get(args[idx]) = static_cast(1); - uint64_t index = thread_offset + idx; + IndexT index = thread_offset + idx; if (index < numel) { std::get(args[idx]) = reinterpret_cast(ins[Index])[index]; @@ -202,24 +207,23 @@ struct BroadcastDataLoader { }; // Vectorized elementwise Loader without consideration of IsBoundary. -template -struct BroadcastDataLoader { +template +struct BroadcastDataLoader { template static __device__ __forceinline__ void Apply(const Array1 &ins, ArgsT *args, const Array2 &configs, const Array3 &use_broadcast, - uint64_t block_offset, + IndexT block_offset, const int num, - uint64_t numel, + IndexT numel, int read_lens) { using Type = std::tuple_element_t; using VecType = phi::kps::details::VectorType; VecType vec_temp; - int64_t thread_offset = - static_cast(threadIdx.x) + - static_cast(blockIdx.x) * static_cast(blockDim.x); + IndexT thread_offset = static_cast(threadIdx.x) + + static_cast(blockIdx.x) * blockDim.x; const VecType *__restrict__ vec_input = reinterpret_cast(ins[Index]); vec_temp = vec_input[thread_offset]; @@ -260,8 +264,13 @@ struct BroadcastDataSetter { #endif // static broadcast unroller -template