long m n k, mnk < LONG_MAX

Author: Dong Jun Woun

src/components/amd_smi/tests/amdsmi_gemm.c

@@ -58,9 +58,9 @@ static const int kTestEventCount =
     (int)(sizeof(kTestEventTemplates) / sizeof(kTestEventTemplates[0]));
 
 struct run_config {
-    int m_dim;
-    int k_dim;
-    int n_dim;
+    long m_dim;
+    long k_dim;
+    long n_dim;
     int event_count;
     int iterations;
     useconds_t iteration_delay_us;
@@ -211,33 +211,58 @@ static int parse_test_override(int *argc, char **argv) {
 
         long dims[3] = {0};
         for (int d = 0; d < 3; ++d) {
-            char *endptr = NULL;
+            const char *arg = argv[i + 1 + d];
             errno = 0;
-            long value = strtol(argv[i + 1 + d], &endptr, 10);
-            if (errno != 0 || !endptr || *endptr != '\0') {
+            char *endptr = NULL;
+            long value = strtol(arg, &endptr, 10);
+            if (endptr == arg || !endptr || *endptr != '\0') {
                 fprintf(stderr, "Invalid integer value for %s argument: %s\n",
-                        request_rocblas ? "--testrocblas" : "--test", argv[i + 1 + d]);
+                        request_rocblas ? "--testrocblas" : "--test", arg);
+                return -1;
+            }
+            if (errno == ERANGE) {
+                fprintf(stderr, "Dimension for %s is out of range: %s\n",
+                        request_rocblas ? "--testrocblas" : "--test", arg);
                 return -1;
             }
-            if (value <= 0 || value > INT_MAX) {
-                fprintf(stderr, "Dimension for %s must be between 1 and %d: %s\n",
-                        request_rocblas ? "--testrocblas" : "--test", INT_MAX, argv[i + 1 + d]);
+            if (value <= 0) {
+                fprintf(stderr, "Dimension for %s must be a positive integer: %s\n",
+                        request_rocblas ? "--testrocblas" : "--test", arg);
                 return -1;
             }
             dims[d] = value;
         }
 
-        g_run_config.m_dim = (int)dims[0];
-        g_run_config.k_dim = (int)dims[1];
-        g_run_config.n_dim = (int)dims[2];
+        if (request_rocblas) {
+            for (int d = 0; d < 3; ++d) {
+                if (dims[d] > INT_MAX) {
+                    fprintf(stderr, "rocBLAS test requires dimensions to be <= %d: %s\n",
+                            INT_MAX, argv[i + 1 + d]);
+                    return -1;
+                }
+            }
+        }
+
+        long m = dims[0];
+        long k = dims[1];
+        long n = dims[2];
+        if ((m > LONG_MAX / k) || (n > LONG_MAX / k) || (m > LONG_MAX / n)) {
+            fprintf(stderr, "Input dimensions are too large; ensure M*K, N*K, and M*N are <= %ld.\n",
+                    LONG_MAX);
+            return -1;
+        }
+
+        g_run_config.m_dim = m;
+        g_run_config.k_dim = k;
+        g_run_config.n_dim = n;
         g_run_config.event_count = kTestEventCount;
         g_run_config.iterations = ITERATIONS_PER_STREAM;
         g_run_config.iteration_delay_us = 5000000; /* 5 second pause between iterations */
         g_run_config.test_mode = true;
         g_run_config.use_rocblas = request_rocblas;
         g_run_config.device_index = 0;
         snprintf(g_run_config.csv_path, sizeof(g_run_config.csv_path),
-                 "amdsmi_gemm_%d_%d_%d%s.csv",
+                 "amdsmi_gemm_%ld_%ld_%ld%s.csv",
                  g_run_config.m_dim,
                  g_run_config.k_dim,
                  g_run_config.n_dim,
@@ -352,35 +377,42 @@ static void *monitor_events(void *args) {
  *        A: MxK, B: KxN, C: MxN (row-major)
  */
 __global__ void dgemm_kernel(const double *A, const double *B, double *C,
-                             int M, int N, int K, double alpha, double beta) {
-    int row = blockIdx.y * blockDim.y + threadIdx.y;
-    int col = blockIdx.x * blockDim.x + threadIdx.x;
+                             long M, long N, long K, double alpha, double beta) {
+    long row = (long)blockIdx.y * blockDim.y + threadIdx.y;
+    long col = (long)blockIdx.x * blockDim.x + threadIdx.x;
 
     if (row < M && col < N) {
         double sum = 0.0;
-        for (int k = 0; k < K; k++) {
-            sum += A[row * K + k] * B[k * N + col];
+        long row_offset = row * K;
+        for (long k = 0; k < K; ++k) {
+            sum += A[row_offset + k] * B[k * N + col];
         }
-        C[row * N + col] = alpha * sum + beta * C[row * N + col];
+        long c_index = row * N + col;
+        C[c_index] = alpha * sum + beta * C[c_index];
     }
 }
 
 static int rocblas_dgemm_row_major(rocblas_handle handle,
                                    const double *dA,
                                    const double *dB,
                                    double *dC,
-                                   int M,
-                                   int N,
-                                   int K,
+                                   long M,
+                                   long N,
+                                   long K,
                                    double alpha,
                                    double beta) {
     if (!handle) {
         return -1;
     }
 
-    const int m_prime = N;
-    const int n_prime = M;
-    const int k_prime = K;
+    if (M > INT_MAX || N > INT_MAX || K > INT_MAX) {
+        fprintf(stderr, "rocBLAS dgemm requires dimensions <= %d.\n", INT_MAX);
+        return -1;
+    }
+
+    const int m_prime = (int)N;
+    const int n_prime = (int)M;
+    const int k_prime = (int)K;
 
     const int lda = m_prime;
     const int ldb = k_prime;
@@ -458,9 +490,9 @@ static int real_main(const HarnessOpts *opts) {
         SKIP("Unable to locate the amd_smi component (PAPI built without ROCm?)");
     }
 
-    const int m_dim = g_run_config.m_dim;
-    const int k_dim = g_run_config.k_dim;
-    const int n_dim = g_run_config.n_dim;
+    const long m_dim = g_run_config.m_dim;
+    const long k_dim = g_run_config.k_dim;
+    const long n_dim = g_run_config.n_dim;
     const int event_count = g_run_config.event_count;
     const int iterations = g_run_config.iterations;
     const useconds_t pause_us = g_run_config.iteration_delay_us;
@@ -716,8 +748,8 @@ static int real_main(const HarnessOpts *opts) {
     double beta  = 0.5;
 
     dim3 blockDim(32, 32);
-    dim3 gridDim((n_dim + blockDim.x - 1) / blockDim.x,
-                 (m_dim + blockDim.y - 1) / blockDim.y);
+    dim3 gridDim((unsigned int)((n_dim + blockDim.x - 1) / blockDim.x),
+                 (unsigned int)((m_dim + blockDim.y - 1) / blockDim.y));
 
     for (int iter = 0; iter < iterations; ++iter) {
         for (int s = 0; s < NUM_STREAMS; s++) {