Harden CUDA fp16 transpose index math against int overflow

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 <limits>
+
 #define TRANS_TILE_DIM 32
 #define BLOCK_ROWS 8
 #define COPY_TILE_DIM 1024
@@ -31,7 +33,8 @@
   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<int64_t>(y + j) * m + x;
+      tile[threadIdx.y + j][threadIdx.x] = idata[input_offset];
     }
   }
 
@@ -44,7 +47,8 @@
 
   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<int64_t>(y + j) * n + x;
+    odata[output_offset] = tile[threadIdx.x][threadIdx.y + j];
   }
 }
 
@@ -64,11 +68,23 @@
 }  // 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<unsigned int>((static_cast<int64_t>(n) + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM);
+  const auto grid_y = static_cast<unsigned int>((static_cast<int64_t>(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 row * stride + col addressing in device code.
+  // Keep fallback disabled when total element count would overflow int32 indexing.
+  if (static_cast<int64_t>(m) * static_cast<int64_t>(n) > std::numeric_limits<int>::max()) {
+    return false;
+  }
+
   dim3 dimGrid = cublasTransposeHelperDimGrid(m, n);
   return dimGrid.y < 65536;
 }
@@ -78,6 +94,7 @@
     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);
   } else {

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,16 @@ TEST(CudaUtilsTest, FillCorrectness) {
   TestFillCorrectness<double>(1 << 20, 5.0);
 }
 
+TEST(CudaUtilsTest, CanUseTransposeHelperRejectsOverflowingElementCount) {
+  EXPECT_TRUE(CanUse_cublasTransposeHelper_MLFloat16(100, 100));
+  EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(100, 25000000));
+}
+
+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