diff --git a/onnxruntime/core/providers/cuda/fpgeneric.cu b/onnxruntime/core/providers/cuda/fpgeneric.cu index 78426c1294883..21097d1080f34 100644 --- a/onnxruntime/core/providers/cuda/fpgeneric.cu +++ b/onnxruntime/core/providers/cuda/fpgeneric.cu @@ -13,6 +13,8 @@ #include "core/providers/cuda/curand_wrapper.h" #include "core/providers/cuda/cu_inc/common.cuh" +#include + #define TRANS_TILE_DIM 32 #define BLOCK_ROWS 8 #define COPY_TILE_DIM 1024 @@ -31,7 +33,8 @@ __global__ void transposeNoOverlap(half* odata, const half* idata, const int m, if (x < m) { for (int j = 0; j < TRANS_TILE_DIM; j += BLOCK_ROWS) { if (j >= (n - y)) continue; - tile[threadIdx.y + j][threadIdx.x] = idata[(y + j) * m + x]; + const int64_t input_offset = static_cast(y + j) * m + x; + tile[threadIdx.y + j][threadIdx.x] = idata[input_offset]; } } @@ -44,7 +47,8 @@ __global__ void transposeNoOverlap(half* odata, const half* idata, const int m, for (int j = 0; j < TRANS_TILE_DIM; j += BLOCK_ROWS) { if ((y + j) >= m) return; - odata[(y + j) * n + x] = tile[threadIdx.x][threadIdx.y + j]; + const int64_t output_offset = static_cast(y + j) * n + x; + odata[output_offset] = tile[threadIdx.x][threadIdx.y + j]; } } @@ -64,20 +68,34 @@ __global__ void CopyVectorBFloat16(const onnxruntime::BFloat16* x, int incx, onn } // namespace dim3 cublasTransposeHelperDimGrid(int m, int n) { - return dim3((n + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM, (m + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM, 1); + const auto grid_x = static_cast((static_cast(n) + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM); + const auto grid_y = static_cast((static_cast(m) + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM); + return dim3(grid_x, grid_y, 1); } // cublasTransposeHelper can only be used if it won't overflow the 65536 grid y dimension size __host__ bool CanUse_cublasTransposeHelper_MLFloat16(int m, int n) { + if (m <= 0 || n <= 0) { + return false; + } + + // transposeNoOverlap uses int64_t row * stride + col addressing in device code. + // Keep fallback disabled when total element count would overflow 32-bit launch/indexing assumptions. + if (static_cast(m) * static_cast(n) > std::numeric_limits::max()) { + return false; + } + dim3 dimGrid = cublasTransposeHelperDimGrid(m, n); return dimGrid.y < 65536; } cublasStatus_t cublasTransposeHelper(cudaStream_t stream, cublasHandle_t, cublasOperation_t, cublasOperation_t, int m, int n, const half*, const half* A, int, const half*, const half*, int, half* C, int) { + ORT_ENFORCE(m > 0 && n > 0); if (C != A) { dim3 dimGrid = cublasTransposeHelperDimGrid(m, n); dim3 dimBlock(TRANS_TILE_DIM, BLOCK_ROWS, 1); + ORT_ENFORCE(static_cast(m) * static_cast(n) <= std::numeric_limits::max()); ORT_ENFORCE(dimGrid.y < 65536); // To prevent this, call CanUse_cublasTransposeHelper_MLFloat16 first transposeNoOverlap<<>>(C, A, n, m); } else { diff --git a/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc b/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc index 7468a5718425e..5cf06d02ca804 100644 --- a/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc @@ -8,6 +8,7 @@ #include "core/common/common.h" #include "core/providers/cuda/shared_inc/cuda_call.h" +#include "core/providers/cuda/shared_inc/fpgeneric.h" #include "core/providers/cuda/shared_inc/cuda_utils.h" namespace onnxruntime { @@ -48,6 +49,23 @@ TEST(CudaUtilsTest, FillCorrectness) { TestFillCorrectness(1 << 20, 5.0); } +TEST(CudaUtilsTest, CanUseTransposeHelperRejectsOverflowingElementCount) { + EXPECT_TRUE(CanUse_cublasTransposeHelper_MLFloat16(100, 100)); + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(100, 25000000)); +} + +TEST(CudaUtilsTest, CanUseTransposeHelperRejectsNonPositiveDimensions) { + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(0, 100)); + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(100, 0)); + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(-1, 100)); + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(100, -1)); +} + +TEST(CudaUtilsTest, CanUseTransposeHelperRejectsGridYOverflow) { + // For TRANS_TILE_DIM=32, m=2097152 yields grid_y=65536, which is out of range. + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(2097152, 1)); +} + } // namespace test } // namespace cuda } // namespace onnxruntime