Skip to content

Fix int32 overflow in CUDA Cast and UnaryElementWise kernels for tensors with >2^31 elements #28386

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tianleiwu merged 6 commits into from
May 19, 2026
Merged

Fix int32 overflow in CUDA Cast and UnaryElementWise kernels for tensors with >2^31 elements #28386

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
3ace5f3
Initial plan
Copilot May 6, 2026
5f5b236
Fix int32 overflow in CUDA Cast and UnaryElementWise kernels for tens…
Copilot May 6, 2026
92ac4fa
Fix placeholder issue number in test comment
Copilot May 6, 2026
f28a737
Address review feedback: use size_t for pair_count, rename test, add …
Copilot May 18, 2026
081aa2f
Merge remote-tracking branch 'origin/main' into copilot/fix-cuda-cast…
Copilot May 18, 2026
ad13c0e
check blocksPerGrid
tianleiwu May 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions onnxruntime/core/providers/cuda/cu_inc/unary_elementwise_impl.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,11 @@ __global__ void _UnaryElementWise(
const InT* input_data,
OutT* output_data,
const FuncT functor,
CUDA_LONG N) {
CUDA_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
int64_t N) {
int64_t start = static_cast<int64_t>(NumElementsPerThread) * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
InT value[NumElementsPerThread];

CUDA_LONG id = start;
int64_t id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
Expand Down Expand Up @@ -47,8 +47,10 @@ void UnaryElementWiseImpl(
if (count == 0) // special case where there's a dim value of 0 in the shape
return;

int blocksPerGrid = static_cast<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
CUDA_LONG N = static_cast<CUDA_LONG>(count);
size_t blocksPerGridSize = CeilDiv(count, static_cast<size_t>(GridDim::maxThreadsPerBlock) * GridDim::maxElementsPerThread);
ORT_ENFORCE(blocksPerGridSize <= static_cast<size_t>(INT32_MAX), "Grid size exceeds CUDA limits");
int blocksPerGrid = static_cast<int>(blocksPerGridSize);
int64_t N = static_cast<int64_t>(count);
_UnaryElementWise<InT, OutT, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
input_data,
Expand Down
44 changes: 26 additions & 18 deletions onnxruntime/core/providers/cuda/tensor/cast_op.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -220,8 +220,8 @@ struct CastStd<Float4E2M1x2, float> {
#endif // DISABLE_FLOAT4_TYPES

template <int NumThreadsPerBlock, int NumElementsPerThread, typename OutT, typename InT>
__global__ void CastKernelStd(const InT* input, OutT* output, CUDA_LONG N, CastStd<OutT, InT> cast) {
CUDA_LONG id = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
__global__ void CastKernelStd(const InT* input, OutT* output, int64_t N, CastStd<OutT, InT> cast) {
int64_t id = static_cast<int64_t>(NumElementsPerThread) * NumThreadsPerBlock * blockIdx.x + threadIdx.x;

#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
Expand All @@ -237,11 +237,13 @@ Status CudaCastStd(cudaStream_t stream, const InT* input, OutT* output, size_t n
if (num_of_elements <= 0)
return Status::OK();

int blocksPerGrid = static_cast<int>(CeilDiv(num_of_elements, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
size_t blocksPerGridSize = CeilDiv(num_of_elements, static_cast<size_t>(GridDim::maxThreadsPerBlock) * GridDim::maxElementsPerThread);
ORT_RETURN_IF_NOT(blocksPerGridSize <= static_cast<size_t>(INT32_MAX), "Grid size exceeds CUDA limits");
int blocksPerGrid = static_cast<int>(blocksPerGridSize);
CastKernelStd<GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread, OutT, InT><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
input,
output,
static_cast<int>(num_of_elements),
static_cast<int64_t>(num_of_elements),
CastStd<OutT, InT>());
return Status::OK();
}
Expand All @@ -251,10 +253,10 @@ Status CudaCastStd(cudaStream_t stream, const InT* input, OutT* output, size_t n
template <int NumThreadsPerBlock, int NumElementsPerThread, bool is_odd, typename OutPairType, typename InPairType,
typename OutSingleType, typename InSingleType>
__global__ void CudaCastPairwiseKernel(const InPairType* input, OutPairType* output,
CUDA_LONG pair_count,
int64_t pair_count,
CastStd<OutPairType, InPairType> pair_caster,
CastStd<OutSingleType, InSingleType> singleton_caster) {
CUDA_LONG id = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
int64_t id = static_cast<int64_t>(NumElementsPerThread) * NumThreadsPerBlock * blockIdx.x + threadIdx.x;

#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
Expand Down Expand Up @@ -284,9 +286,11 @@ Status CudaCastPairwise(cudaStream_t stream, const Float4E2M1x2* input, float* o

bool is_odd = (num_of_elements & 0x01) != 0;

int pair_count = static_cast<int>(num_of_elements / 2);
size_t pair_count = num_of_elements / 2;

int blocksPerGrid = static_cast<int>(CeilDiv(pair_count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
size_t blocksPerGridSize = CeilDiv(pair_count, static_cast<size_t>(GridDim::maxThreadsPerBlock) * GridDim::maxElementsPerThread);
ORT_RETURN_IF_NOT(blocksPerGridSize <= static_cast<size_t>(INT32_MAX), "Grid size exceeds CUDA limits");
int blocksPerGrid = static_cast<int>(blocksPerGridSize);

if (pair_count == 0) {
blocksPerGrid = 1;
Expand All @@ -296,14 +300,14 @@ Status CudaCastPairwise(cudaStream_t stream, const Float4E2M1x2* input, float* o
CudaCastPairwiseKernel<GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread, true,
float2, Float4E2M1x2, float, Float4E2M1x2>
<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
input, reinterpret_cast<float2*>(output), pair_count,
input, reinterpret_cast<float2*>(output), static_cast<int64_t>(pair_count),
CastStd<float2, Float4E2M1x2>(),
CastStd<float, Float4E2M1x2>());
} else {
CudaCastPairwiseKernel<GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread, false,
float2, Float4E2M1x2, float, Float4E2M1x2>
<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
input, reinterpret_cast<float2*>(output), pair_count,
input, reinterpret_cast<float2*>(output), static_cast<int64_t>(pair_count),
CastStd<float2, Float4E2M1x2>(),
CastStd<float, Float4E2M1x2>());
}
Expand All @@ -318,9 +322,11 @@ Status CudaCastPairwise(cudaStream_t stream, const float* input, Float4E2M1x2* o

bool is_odd = (num_of_elements & 0x01) != 0;

int pair_count = static_cast<int>(num_of_elements / 2);
size_t pair_count = num_of_elements / 2;

int blocksPerGrid = static_cast<int>(CeilDiv(pair_count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
size_t blocksPerGridSize = CeilDiv(pair_count, static_cast<size_t>(GridDim::maxThreadsPerBlock) * GridDim::maxElementsPerThread);
ORT_RETURN_IF_NOT(blocksPerGridSize <= static_cast<size_t>(INT32_MAX), "Grid size exceeds CUDA limits");
int blocksPerGrid = static_cast<int>(blocksPerGridSize);

if (pair_count == 0) {
blocksPerGrid = 1;
Expand All @@ -330,14 +336,14 @@ Status CudaCastPairwise(cudaStream_t stream, const float* input, Float4E2M1x2* o
CudaCastPairwiseKernel<GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread, true,
Float4E2M1x2, float2, Float4E2M1x2, float>
<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
reinterpret_cast<const float2*>(input), output, pair_count,
reinterpret_cast<const float2*>(input), output, static_cast<int64_t>(pair_count),
CastStd<Float4E2M1x2, float2>(),
CastStd<Float4E2M1x2, float>());
} else {
CudaCastPairwiseKernel<GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread, false,
Float4E2M1x2, float2, Float4E2M1x2, float>
<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
reinterpret_cast<const float2*>(input), output, pair_count,
reinterpret_cast<const float2*>(input), output, static_cast<int64_t>(pair_count),
CastStd<Float4E2M1x2, float2>(),
CastStd<Float4E2M1x2, float>());
}
Expand All @@ -353,8 +359,8 @@ template Status CudaCastPairwise<Float4E2M1x2, float>(cudaStream_t stream, const
#if !defined(DISABLE_FLOAT8_TYPES)

template <int NumThreadsPerBlock, int NumElementsPerThread, typename OutT, typename InT>
__global__ void CastKernelSat(const InT* input, OutT* output, CUDA_LONG N, CastSat<OutT, InT> cast, bool saturate) {
CUDA_LONG id = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
__global__ void CastKernelSat(const InT* input, OutT* output, int64_t N, CastSat<OutT, InT> cast, bool saturate) {
int64_t id = static_cast<int64_t>(NumElementsPerThread) * NumThreadsPerBlock * blockIdx.x + threadIdx.x;

#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
Expand All @@ -370,11 +376,13 @@ Status CudaCastSat(cudaStream_t stream, const InT* input, OutT* output, size_t n
if (num_of_element <= 0)
return Status::OK();

int blocksPerGrid = static_cast<int>(CeilDiv(num_of_element, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
size_t blocksPerGridSize = CeilDiv(num_of_element, static_cast<size_t>(GridDim::maxThreadsPerBlock) * GridDim::maxElementsPerThread);
ORT_RETURN_IF_NOT(blocksPerGridSize <= static_cast<size_t>(INT32_MAX), "Grid size exceeds CUDA limits");
int blocksPerGrid = static_cast<int>(blocksPerGridSize);
CastKernelSat<GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread, OutT, InT><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
input,
output,
static_cast<int>(num_of_element),
static_cast<int64_t>(num_of_element),
CastSat<OutT, InT>(),
saturate);
return Status::OK();
Expand Down
57 changes: 57 additions & 0 deletions onnxruntime/test/providers/cpu/tensor/cast_op_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3127,6 +3127,63 @@ TEST(CastOpTest, CopyCpuTensor_SubByteTypes_DistinctBuffers) {
}
}

// Correctness test for Cast kernel with a moderately large tensor.
// Exercises the same kernel code path as tensors > 2^31 elements but stays within
// CI GPU memory limits. For the actual overflow scenario, see the host-side test below.
TEST(CastOpTest, CastKernelCorrectness_ModerateSize) {
constexpr int64_t num_elements = 1 << 24; // 16M elements
const std::vector<int64_t> shape = {num_elements};

std::vector<float> input(num_elements);
std::vector<int32_t> expected(num_elements);
for (int64_t i = 0; i < num_elements; ++i) {
input[i] = static_cast<float>(i % 1000);
expected[i] = static_cast<int32_t>(i % 1000);
}

TestCastOp<float, int32_t>(gsl::make_span(input), gsl::make_span(expected), shape);
}

// Host-side regression test that verifies the grid launch arithmetic uses 64-bit
// types for element counts exceeding INT32_MAX. This validates the fix without
// needing to allocate > 8 GB of GPU memory.
// The fix changed:
// CUDA_LONG N = static_cast<CUDA_LONG>(count) // was int32 truncation
// to:
// int64_t N = static_cast<int64_t>(count) // correct 64-bit
TEST(CastOpTest, CastKernel_Int64IndexArithmetic_NoOverflow) {
// Simulate the grid launch calculation from UnaryElementWiseImpl / CudaCastStd
// with a count that exceeds INT32_MAX.
constexpr size_t count = static_cast<size_t>(INT32_MAX) + 65536; // 2^31 + 65536
constexpr int maxThreadsPerBlock = 256;
constexpr int maxElementsPerThread = 4;

// Verify N is correctly represented (not truncated to int32)
int64_t N = static_cast<int64_t>(count);
ASSERT_GT(N, static_cast<int64_t>(INT32_MAX));
ASSERT_EQ(N, static_cast<int64_t>(count));

// Verify blocksPerGrid calculation doesn't overflow
// (uses size_t arithmetic for the divisor)
size_t elements_per_block = static_cast<size_t>(maxThreadsPerBlock) * maxElementsPerThread;
int blocksPerGrid = static_cast<int>((count + elements_per_block - 1) / elements_per_block);
ASSERT_GT(blocksPerGrid, 0);
// For count = 2^31 + 65536, elements_per_block = 1024, we expect ~2M blocks
ASSERT_EQ(blocksPerGrid, static_cast<int>((count + 1023) / 1024));

// Verify that the per-thread index calculation doesn't overflow in int64_t
// Simulate the last block's thread 0: id = NumElementsPerThread * NumThreadsPerBlock * (blocksPerGrid-1) + 0
int64_t last_block_start = static_cast<int64_t>(maxElementsPerThread) * maxThreadsPerBlock *
(blocksPerGrid - 1);
ASSERT_GT(last_block_start, 0); // Positive (no overflow)
ASSERT_LE(last_block_start, N); // Within bounds

// Verify the old int32 code would have failed:
// static_cast<int32_t>(count) would silently wrap
int32_t truncated_N = static_cast<int32_t>(count);
ASSERT_LT(truncated_N, 0); // Proves the old code was broken (wraps negative)
}

#if !defined(DISABLE_FLOAT8_TYPES)

float FloatFromBits(uint32_t bits) {
Expand Down
Loading