diff --git a/onnxruntime/core/providers/cuda/cu_inc/unary_elementwise_impl.cuh b/onnxruntime/core/providers/cuda/cu_inc/unary_elementwise_impl.cuh index c8ddbadb12fb2..5959482e5664e 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/unary_elementwise_impl.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/unary_elementwise_impl.cuh @@ -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(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) { @@ -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(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); - CUDA_LONG N = static_cast(count); + size_t blocksPerGridSize = CeilDiv(count, static_cast(GridDim::maxThreadsPerBlock) * GridDim::maxElementsPerThread); + ORT_ENFORCE(blocksPerGridSize <= static_cast(INT32_MAX), "Grid size exceeds CUDA limits"); + int blocksPerGrid = static_cast(blocksPerGridSize); + int64_t N = static_cast(count); _UnaryElementWise <<>>( input_data, diff --git a/onnxruntime/core/providers/cuda/tensor/cast_op.cu b/onnxruntime/core/providers/cuda/tensor/cast_op.cu index a8cd6caaa5d5f..c56d613e25241 100644 --- a/onnxruntime/core/providers/cuda/tensor/cast_op.cu +++ b/onnxruntime/core/providers/cuda/tensor/cast_op.cu @@ -220,8 +220,8 @@ struct CastStd { #endif // DISABLE_FLOAT4_TYPES template -__global__ void CastKernelStd(const InT* input, OutT* output, CUDA_LONG N, CastStd cast) { - CUDA_LONG id = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x; +__global__ void CastKernelStd(const InT* input, OutT* output, int64_t N, CastStd cast) { + int64_t id = static_cast(NumElementsPerThread) * NumThreadsPerBlock * blockIdx.x + threadIdx.x; #pragma unroll for (int i = 0; i < NumElementsPerThread; i++) { @@ -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(CeilDiv(num_of_elements, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); + size_t blocksPerGridSize = CeilDiv(num_of_elements, static_cast(GridDim::maxThreadsPerBlock) * GridDim::maxElementsPerThread); + ORT_RETURN_IF_NOT(blocksPerGridSize <= static_cast(INT32_MAX), "Grid size exceeds CUDA limits"); + int blocksPerGrid = static_cast(blocksPerGridSize); CastKernelStd<<>>( input, output, - static_cast(num_of_elements), + static_cast(num_of_elements), CastStd()); return Status::OK(); } @@ -251,10 +253,10 @@ Status CudaCastStd(cudaStream_t stream, const InT* input, OutT* output, size_t n template __global__ void CudaCastPairwiseKernel(const InPairType* input, OutPairType* output, - CUDA_LONG pair_count, + int64_t pair_count, CastStd pair_caster, CastStd singleton_caster) { - CUDA_LONG id = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x; + int64_t id = static_cast(NumElementsPerThread) * NumThreadsPerBlock * blockIdx.x + threadIdx.x; #pragma unroll for (int i = 0; i < NumElementsPerThread; i++) { @@ -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(num_of_elements / 2); + size_t pair_count = num_of_elements / 2; - int blocksPerGrid = static_cast(CeilDiv(pair_count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); + size_t blocksPerGridSize = CeilDiv(pair_count, static_cast(GridDim::maxThreadsPerBlock) * GridDim::maxElementsPerThread); + ORT_RETURN_IF_NOT(blocksPerGridSize <= static_cast(INT32_MAX), "Grid size exceeds CUDA limits"); + int blocksPerGrid = static_cast(blocksPerGridSize); if (pair_count == 0) { blocksPerGrid = 1; @@ -296,14 +300,14 @@ Status CudaCastPairwise(cudaStream_t stream, const Float4E2M1x2* input, float* o CudaCastPairwiseKernel <<>>( - input, reinterpret_cast(output), pair_count, + input, reinterpret_cast(output), static_cast(pair_count), CastStd(), CastStd()); } else { CudaCastPairwiseKernel <<>>( - input, reinterpret_cast(output), pair_count, + input, reinterpret_cast(output), static_cast(pair_count), CastStd(), CastStd()); } @@ -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(num_of_elements / 2); + size_t pair_count = num_of_elements / 2; - int blocksPerGrid = static_cast(CeilDiv(pair_count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); + size_t blocksPerGridSize = CeilDiv(pair_count, static_cast(GridDim::maxThreadsPerBlock) * GridDim::maxElementsPerThread); + ORT_RETURN_IF_NOT(blocksPerGridSize <= static_cast(INT32_MAX), "Grid size exceeds CUDA limits"); + int blocksPerGrid = static_cast(blocksPerGridSize); if (pair_count == 0) { blocksPerGrid = 1; @@ -330,14 +336,14 @@ Status CudaCastPairwise(cudaStream_t stream, const float* input, Float4E2M1x2* o CudaCastPairwiseKernel <<>>( - reinterpret_cast(input), output, pair_count, + reinterpret_cast(input), output, static_cast(pair_count), CastStd(), CastStd()); } else { CudaCastPairwiseKernel <<>>( - reinterpret_cast(input), output, pair_count, + reinterpret_cast(input), output, static_cast(pair_count), CastStd(), CastStd()); } @@ -353,8 +359,8 @@ template Status CudaCastPairwise(cudaStream_t stream, const #if !defined(DISABLE_FLOAT8_TYPES) template -__global__ void CastKernelSat(const InT* input, OutT* output, CUDA_LONG N, CastSat cast, bool saturate) { - CUDA_LONG id = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x; +__global__ void CastKernelSat(const InT* input, OutT* output, int64_t N, CastSat cast, bool saturate) { + int64_t id = static_cast(NumElementsPerThread) * NumThreadsPerBlock * blockIdx.x + threadIdx.x; #pragma unroll for (int i = 0; i < NumElementsPerThread; i++) { @@ -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(CeilDiv(num_of_element, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); + size_t blocksPerGridSize = CeilDiv(num_of_element, static_cast(GridDim::maxThreadsPerBlock) * GridDim::maxElementsPerThread); + ORT_RETURN_IF_NOT(blocksPerGridSize <= static_cast(INT32_MAX), "Grid size exceeds CUDA limits"); + int blocksPerGrid = static_cast(blocksPerGridSize); CastKernelSat<<>>( input, output, - static_cast(num_of_element), + static_cast(num_of_element), CastSat(), saturate); return Status::OK(); diff --git a/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc b/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc index 0e14bc59a09c9..038a8eaade116 100644 --- a/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc @@ -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 shape = {num_elements}; + + std::vector input(num_elements); + std::vector expected(num_elements); + for (int64_t i = 0; i < num_elements; ++i) { + input[i] = static_cast(i % 1000); + expected[i] = static_cast(i % 1000); + } + + TestCastOp(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(count) // was int32 truncation +// to: +// int64_t N = static_cast(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(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(count); + ASSERT_GT(N, static_cast(INT32_MAX)); + ASSERT_EQ(N, static_cast(count)); + + // Verify blocksPerGrid calculation doesn't overflow + // (uses size_t arithmetic for the divisor) + size_t elements_per_block = static_cast(maxThreadsPerBlock) * maxElementsPerThread; + int blocksPerGrid = static_cast((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((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(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(count) would silently wrap + int32_t truncated_N = static_cast(count); + ASSERT_LT(truncated_N, 0); // Proves the old code was broken (wraps negative) +} + #if !defined(DISABLE_FLOAT8_TYPES) float FloatFromBits(uint32_t bits) {