fix: 修复整数溢出风险

修复了多个文件中潜在的整数溢出问题：

1. src/kernels/tensor_core_benchmark.cuh
   - 添加矩阵元素数量溢出检查
   - 使用 size_t 进行乘法运算后再转换为 int

2. src/kernels/tensor_core_sgemm.cuh
   - 添加矩阵元素数量溢出检查
   - 确保传递给 kernel 的 size 参数安全

3. src/utils/benchmark.cuh
   - 使用 size_t 进行矩阵大小计算
   - 添加溢出检查防止内存分配错误

4. tests/test_sgemm.cu
   - 删除多余的命名空间闭合标记

这些问题在大矩阵场景下（如 M,K,N > 46340）可能导致：
- 数据损坏
- 内存访问越界
- 程序崩溃

Author: LessUp

src/kernels/tensor_core_benchmark.cuh

@@ -79,8 +79,19 @@ runTensorCoreComputeOnlyBenchmark(cublasHandle_t cublas_handle, int M, int K, in
     int gridSizeA = safeGridSize(static_cast<size_t>(M) * K, blockSize);
     int gridSizeB = safeGridSize(static_cast<size_t>(K) * N, blockSize);
 
-    float_to_half_kernel<<<gridSizeA, blockSize>>>(d_A.get(), d_A_fp16.get(), M * K);
-    float_to_half_kernel<<<gridSizeB, blockSize>>>(d_B.get(), d_B_fp16.get(), K * N);
+    size_t num_A = static_cast<size_t>(M) * K;
+    size_t num_B = static_cast<size_t>(K) * N;
+
+    // 检查矩阵元素数量是否超过 int 最大值
+    if (num_A > static_cast<size_t>(INT_MAX)) {
+        throw CudaError("Matrix A size overflow: too many elements for int parameter");
+    }
+    if (num_B > static_cast<size_t>(INT_MAX)) {
+        throw CudaError("Matrix B size overflow: too many elements for int parameter");
+    }
+
+    float_to_half_kernel<<<gridSizeA, blockSize>>>(d_A.get(), d_A_fp16.get(), static_cast<int>(num_A));
+    float_to_half_kernel<<<gridSizeB, blockSize>>>(d_B.get(), d_B_fp16.get(), static_cast<int>(num_B));
     CUDA_CHECK(cudaGetLastError());
     CUDA_CHECK(cudaDeviceSynchronize());
 

src/kernels/tensor_core_sgemm.cuh

@@ -258,6 +258,14 @@ inline void launch_tensor_core_sgemm_with_fallback(const float *A, const float *
     int gridSizeA = safeGridSize(num_A, blockSize);
     int gridSizeB = safeGridSize(num_B, blockSize);
 
+    // 检查矩阵元素数量是否超过 int 最大值
+    if (num_A > static_cast<size_t>(INT_MAX)) {
+        throw CudaError("Matrix A size overflow: too many elements for int parameter");
+    }
+    if (num_B > static_cast<size_t>(INT_MAX)) {
+        throw CudaError("Matrix B size overflow: too many elements for int parameter");
+    }
+
     float_to_half_kernel<<<gridSizeA, blockSize, 0, stream>>>(A, d_A_fp16.get(),
                                                               static_cast<int>(num_A));
     float_to_half_kernel<<<gridSizeB, blockSize, 0, stream>>>(B, d_B_fp16.get(),

src/utils/benchmark.cuh

@@ -65,17 +65,28 @@ class SGEMMBenchmark {
         result.N = N;
 
         // 初始化数据
-        std::vector<float> h_A(M * K), h_B(K * N), h_C(M * N), h_C_ref(M * N);
-        DeviceMemory<float> d_A(M * K);
-        DeviceMemory<float> d_B(K * N);
-        DeviceMemory<float> d_C(M * N);
-        DeviceMemory<float> d_C_ref(M * N);
+        // 安全计算矩阵大小，避免整数溢出
+        size_t size_A = static_cast<size_t>(M) * K;
+        size_t size_B = static_cast<size_t>(K) * N;
+        size_t size_C = static_cast<size_t>(M) * N;
+
+        // 检查是否超过 size_t 范围（实际上是检查是否合理）
+        if (size_A > static_cast<size_t>(INT_MAX) || size_B > static_cast<size_t>(INT_MAX) ||
+            size_C > static_cast<size_t>(INT_MAX)) {
+            throw CudaError("Matrix dimensions too large for benchmark");
+        }
+
+        std::vector<float> h_A(size_A), h_B(size_B), h_C(size_C), h_C_ref(size_C);
+        DeviceMemory<float> d_A(size_A);
+        DeviceMemory<float> d_B(size_B);
+        DeviceMemory<float> d_C(size_C);
+        DeviceMemory<float> d_C_ref(size_C);
 
         initRandomMatrix(h_A.data(), M, K, -1.0f, 1.0f, 42);
         initRandomMatrix(h_B.data(), K, N, -1.0f, 1.0f, 123);
 
-        d_A.copyFromHost(h_A.data(), M * K);
-        d_B.copyFromHost(h_B.data(), K * N);
+        d_A.copyFromHost(h_A.data(), size_A);
+        d_B.copyFromHost(h_B.data(), size_B);
 
         // 计算参考结果
         float alpha = 1.0f, beta = 0.0f;
@@ -94,8 +105,8 @@ class SGEMMBenchmark {
         result.efficiency = calculateEfficiency(result.gflops, getTheoreticalPeakGflops());
 
         // 验证正确性
-        d_C.copyToHost(h_C.data(), M * N);
-        d_C_ref.copyToHost(h_C_ref.data(), M * N);
+        d_C.copyToHost(h_C.data(), size_C);
+        d_C_ref.copyToHost(h_C_ref.data(), size_C);
 
         VerifyResult verify_result = compareMatrices(h_C.data(), h_C_ref.data(), M, N, tolerance);
         result.correct = verify_result.passed;
@@ -115,16 +126,27 @@ class SGEMMBenchmark {
         result.K = K;
         result.N = N;
 
-        std::vector<float> h_A(M * K), h_B(K * N);
-        DeviceMemory<float> d_A(M * K);
-        DeviceMemory<float> d_B(K * N);
-        DeviceMemory<float> d_C(M * N);
+        // 安全计算矩阵大小，避免整数溢出
+        size_t size_A = static_cast<size_t>(M) * K;
+        size_t size_B = static_cast<size_t>(K) * N;
+        size_t size_C = static_cast<size_t>(M) * N;
+
+        // 检查是否超过 size_t 范围
+        if (size_A > static_cast<size_t>(INT_MAX) || size_B > static_cast<size_t>(INT_MAX) ||
+            size_C > static_cast<size_t>(INT_MAX)) {
+            throw CudaError("Matrix dimensions too large for benchmark");
+        }
+
+        std::vector<float> h_A(size_A), h_B(size_B);
+        DeviceMemory<float> d_A(size_A);
+        DeviceMemory<float> d_B(size_B);
+        DeviceMemory<float> d_C(size_C);
 
         initRandomMatrix(h_A.data(), M, K, -1.0f, 1.0f, 42);
         initRandomMatrix(h_B.data(), K, N, -1.0f, 1.0f, 123);
 
-        d_A.copyFromHost(h_A.data(), M * K);
-        d_B.copyFromHost(h_B.data(), K * N);
+        d_A.copyFromHost(h_A.data(), size_A);
+        d_B.copyFromHost(h_B.data(), size_B);
 
         float alpha = 1.0f, beta = 0.0f;
 

tests/test_sgemm.cu

@@ -60,7 +60,6 @@ void computeReference(cublasHandle_t handle, const float *d_A, const float *d_B,
     CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, d_B, N, d_A, K,
                              &beta, d_C, N));
 }
-} // namespace
 
 class ErrorDetectionTest : public ::testing::Test {
   protected: