Harden CUDA fp16 transpose index math against int overflow

[Akshay Sonawane (apsonawane)]

This pull request improves the safety and robustness of CUDA-based matrix transposition in ONNX Runtime by adding overflow checks, using safer integer types for indexing, and expanding test coverage. The main focus is on preventing integer overflows and ensuring correct grid dimension calculations for large matrices.

Safety and overflow prevention:

• Updated all indexing calculations in the transposeNoOverlap CUDA kernel to use int64_t for offsets, preventing integer overflows when handling large matrices (onnxruntime/core/providers/cuda/fpgeneric.cu). [1] [2]
• Modified cublasTransposeHelperDimGrid to use int64_t for grid size calculations, then safely cast to unsigned int (onnxruntime/core/providers/cuda/fpgeneric.cu).
• Added explicit checks in CanUse_cublasTransposeHelper_MLFloat16 and cublasTransposeHelper to reject cases where the element count would overflow a 32-bit integer, or where grid dimensions would exceed CUDA limits (onnxruntime/core/providers/cuda/fpgeneric.cu). [1] [2]

Testing improvements:

• Added unit tests to verify that CanUse_cublasTransposeHelper_MLFloat16 correctly rejects overflowing element counts and grid dimensions (onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc).
• Included the appropriate header for the tested functions in the test file (onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc).

Other:

• Included <limits> to support overflow checks (onnxruntime/core/providers/cuda/fpgeneric.cu).

[Copilot]

Pull request overview

This PR hardens the CUDA MLFloat16 transpose fast-path by moving transpose kernel indexing to 64-bit math and adding host-side guards to prevent invalid/overflowing launch configurations for very large matrices.

Changes:

• Updated transposeNoOverlap indexing to use int64_t offsets and adjusted grid-dimension math to avoid intermediate overflow.
• Added stricter eligibility checks (CanUse_cublasTransposeHelper_MLFloat16) and runtime asserts (ORT_ENFORCE) to prevent out-of-range CUDA grid dimensions / unsafe sizes.
• Added CUDA unit tests to validate the new rejection behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
onnxruntime/core/providers/cuda/fpgeneric.cu	Hardened fp16 transpose kernel indexing and added size/grid validation/guards.
onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc	Added unit tests covering CanUse_cublasTransposeHelper_MLFloat16 rejection cases.

Suppressed comments (1)

onnxruntime/core/providers/cuda/fpgeneric.cu:99

• cublasTransposeHelper relies on callers to validate dimensions, but the current ORT_ENFORCE checks don’t reject non-positive m/n. If m==0 or n==0, cublasTransposeHelperDimGrid produces a 0 grid dimension and the kernel launch configuration is invalid. Adding an explicit precondition check makes the helper robust even if it’s called without CanUse_cublasTransposeHelper_MLFloat16.

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) {
  if (C != A) {
    dim3 dimGrid = cublasTransposeHelperDimGrid(m, n);
    dim3 dimBlock(TRANS_TILE_DIM, BLOCK_ROWS, 1);

    ORT_ENFORCE(static_cast<int64_t>(m) * static_cast<int64_t>(n) <= std::numeric_limits<int>::max());
    ORT_ENFORCE(dimGrid.y < 65536);  // To prevent this, call CanUse_cublasTransposeHelper_MLFloat16 first
    transposeNoOverlap<<<dimGrid, dimBlock, 0, stream>>>(C, A, n, m);

Duplicate approval posted after an equivalent submitted review was already present on the same head.

[Tianlei Wu (tianleiwu)]

The implementation looks correct: arithmetic is widened before overflow can occur, accepted inputs keep the remaining coordinate math representable, invalid launch dimensions are rejected, and oversized shapes retain the generic transpose fallback. I have one non-blocking suggestion to pin the exact limits in the new predicate tests.