From 3a4482d09b8f152bb503c127cf5a0a5d351a2a29 Mon Sep 17 00:00:00 2001
From: Aasta Lin <aastal@nvidia.com>
Date: Mon, 7 Nov 2022 05:47:18 +0000
Subject: [PATCH] half matrixMulCUBLAS

---
 Makefile            |   5 +-
 matrixMulCUBLAS.cpp | 132 ++++++++++++++++++++++++++++++++------------
 2 files changed, 97 insertions(+), 40 deletions(-)

diff --git a/Makefile b/Makefile
index 9f5db72..c2a5190 100644
--- a/Makefile
+++ b/Makefile
@@ -274,7 +274,7 @@ ALL_LDFLAGS += $(addprefix -Xlinker ,$(LDFLAGS))
 ALL_LDFLAGS += $(addprefix -Xlinker ,$(EXTRA_LDFLAGS))
 
 # Common includes and paths for CUDA
-INCLUDES  := -I../../common/inc
+INCLUDES  := -I$(CUDA_PATH)/samples/common/inc
 LIBRARIES :=
 
 ################################################################################
@@ -354,14 +354,11 @@ matrixMulCUBLAS.o:matrixMulCUBLAS.cpp
 
 matrixMulCUBLAS: matrixMulCUBLAS.o
 	$(EXEC) $(NVCC) $(ALL_LDFLAGS) $(GENCODE_FLAGS) -o $@ $+ $(LIBRARIES)
-	$(EXEC) mkdir -p ../../bin/$(TARGET_ARCH)/$(TARGET_OS)/$(BUILD_TYPE)
-	$(EXEC) cp $@ ../../bin/$(TARGET_ARCH)/$(TARGET_OS)/$(BUILD_TYPE)
 
 run: build
 	$(EXEC) ./matrixMulCUBLAS
 
 clean:
 	rm -f matrixMulCUBLAS matrixMulCUBLAS.o
-	rm -rf ../../bin/$(TARGET_ARCH)/$(TARGET_OS)/$(BUILD_TYPE)/matrixMulCUBLAS
 
 clobber: clean
diff --git a/matrixMulCUBLAS.cpp b/matrixMulCUBLAS.cpp
index 132afd4..8e96a43 100644
--- a/matrixMulCUBLAS.cpp
+++ b/matrixMulCUBLAS.cpp
@@ -72,6 +72,52 @@ typedef struct _matrixSize      // Optional Command-line multiplier for matrix s
     unsigned int uiWA, uiHA, uiWB, uiHB, uiWC, uiHC;
 } sMatrixSize;
 
+// Modified from helper_image.h to support half precision arguments
+inline bool sdkCompareL2fe(const __half *reference, const __half *data,
+                           const unsigned int len, const float epsilon) {
+  assert(epsilon >= 0);
+
+  float error = 0;
+  float ref = 0;
+
+  for (unsigned int i = 0; i < len; ++i) {
+    float diff = (float)reference[i] - (float)data[i];
+    error += diff * diff;
+    ref += (float)reference[i] * (float)reference[i];
+  }
+
+  float normRef = sqrtf(ref);
+
+  if (fabs(ref) < 1e-7) {
+#ifdef _DEBUG
+    std::cerr << "ERROR, reference l2-norm is 0\n";
+#endif
+    return false;
+  }
+
+  float normError = sqrtf(error);
+  error = normError / normRef;
+  bool result = error < epsilon;
+#ifdef _DEBUG
+
+  if (!result) {
+    std::cerr << "ERROR, l2-norm error " << error << " is greater than epsilon "
+              << epsilon << "\n";
+  }
+
+#endif
+
+  return result;
+}
+
+inline int idx_n(int row, int col, int height, int width) {
+    return row * width + col;
+}
+
+inline int idx_t(int row, int col, int height, int width) {
+    return col * height + row;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 //! Compute reference data set matrix multiply on CPU
 //! C = A * B
@@ -82,32 +128,34 @@ typedef struct _matrixSize      // Optional Command-line multiplier for matrix s
 //! @param wB         width of matrix B
 ////////////////////////////////////////////////////////////////////////////////
 void
-matrixMulCPU(float *C, const float *A, const float *B, unsigned int hA, unsigned int wA, unsigned int wB)
+matrixMulCPU(__half *C, const __half *A, const __half *B, unsigned int hA, unsigned int wA, unsigned int wB, const bool ta, const bool tb)
 {
+    auto idx_a = (ta) ? idx_t : idx_n;
+    auto idx_b = (tb) ? idx_t : idx_n;
     for (unsigned int i = 0; i < hA; ++i)
         for (unsigned int j = 0; j < wB; ++j)
         {
-            double sum = 0;
+            float sum = 0;
 
             for (unsigned int k = 0; k < wA; ++k)
             {
-                double a = A[i * wA + k];
-                double b = B[k * wB + j];
+                __half a = A[idx_a(i, k, hA, wA)];
+                __half b = B[idx_b(k, j, wA, wB)];
                 sum += a * b;
             }
 
-            C[i * wB + j] = (float)sum;
+            C[i * wB + j] = sum;
         }
 }
 
-// Allocates a matrix with random float entries.
-void randomInit(float *data, int size)
+// Allocates a matrix with random __half entries.
+void randomInit(__half *data, int size)
 {
     for (int i = 0; i < size; ++i)
-        data[i] = rand() / (float)RAND_MAX;
+        data[i] = (__half)(rand() / (float)(RAND_MAX));
 }
 
-void printDiff(float *data1, float *data2, int width, int height, int iListLength, float fListTol)
+void printDiff(__half *data1, __half *data2, int width, int height, int iListLength, float fListTol)
 {
     printf("Listing first %d Differences > %.6f...\n", iListLength, fListTol);
     int i,j,k;
@@ -129,7 +177,7 @@ void printDiff(float *data1, float *data2, int width, int height, int iListLengt
             {
                 if (error_count < iListLength)
                 {
-                    printf("    Loc(%d,%d)\tCPU=%.5f\tGPU=%.5f\tDiff=%.6f\n", i, j, data1[k], data2[k], fDiff);
+                    printf("    Loc(%d,%d)\tCPU=%.5f\tGPU=%.5f\tDiff=%.6f\n", i, j, (float)data1[k], (float)data2[k], fDiff);
                 }
 
                 error_count++;
@@ -170,12 +218,13 @@ void initializeCUDA(int argc, char **argv, int &devID, int &iSizeMultiple, sMatr
 
     int block_size = 32;
 
-    matrix_size.uiWA = 3 * block_size * iSizeMultiple;
-    matrix_size.uiHA = 4 * block_size * iSizeMultiple;
-    matrix_size.uiWB = 2 * block_size * iSizeMultiple;
-    matrix_size.uiHB = 3 * block_size * iSizeMultiple;
-    matrix_size.uiWC = 2 * block_size * iSizeMultiple;
-    matrix_size.uiHC = 4 * block_size * iSizeMultiple;
+    const int N = 2048;
+    matrix_size.uiWA = N;//3 * block_size * iSizeMultiple;
+    matrix_size.uiHA = N;//4 * block_size * iSizeMultiple;
+    matrix_size.uiWB = N;//2 * block_size * iSizeMultiple;
+    matrix_size.uiHB = N;//3 * block_size * iSizeMultiple;
+    matrix_size.uiWC = N;//2 * block_size * iSizeMultiple;
+    matrix_size.uiHC = N;//4 * block_size * iSizeMultiple;
 
     printf("MatrixA(%u,%u), MatrixB(%u,%u), MatrixC(%u,%u)\n",
            matrix_size.uiHA, matrix_size.uiWA,
@@ -194,7 +243,7 @@ void initializeCUDA(int argc, char **argv, int &devID, int &iSizeMultiple, sMatr
 ////////////////////////////////////////////////////////////////////////////////
 //! Run a simple test matrix multiply using CUBLAS
 ////////////////////////////////////////////////////////////////////////////////
-int matrixMultiply(int argc, char **argv, int devID, sMatrixSize &matrix_size)
+int matrixMultiply(int argc, char **argv, int devID, sMatrixSize &matrix_size, bool ta, bool tb)
 {
     cudaDeviceProp deviceProp;
 
@@ -207,11 +256,11 @@ int matrixMultiply(int argc, char **argv, int devID, sMatrixSize &matrix_size)
 
     // allocate host memory for matrices A and B
     unsigned int size_A = matrix_size.uiWA * matrix_size.uiHA;
-    unsigned int mem_size_A = sizeof(float) * size_A;
-    float *h_A = (float *)malloc(mem_size_A);
+    unsigned int mem_size_A = sizeof(__half) * size_A;
+    __half *h_A = (__half *)malloc(mem_size_A);
     unsigned int size_B = matrix_size.uiWB * matrix_size.uiHB;
-    unsigned int mem_size_B = sizeof(float) * size_B;
-    float *h_B = (float *)malloc(mem_size_B);
+    unsigned int mem_size_B = sizeof(__half) * size_B;
+    __half *h_B = (__half *)malloc(mem_size_B);
 
     // set seed for rand()
     srand(2006);
@@ -221,13 +270,13 @@ int matrixMultiply(int argc, char **argv, int devID, sMatrixSize &matrix_size)
     randomInit(h_B, size_B);
 
     // allocate device memory
-    float *d_A, *d_B, *d_C;
+    __half *d_A, *d_B, *d_C;
     unsigned int size_C = matrix_size.uiWC * matrix_size.uiHC;
-    unsigned int mem_size_C = sizeof(float) * size_C;
+    unsigned int mem_size_C = sizeof(__half) * size_C;
 
     // allocate host memory for the result
-    float *h_C      = (float *) malloc(mem_size_C);
-    float *h_CUBLAS = (float *) malloc(mem_size_C);
+    __half *h_C      = (__half *) malloc(mem_size_C);
+    __half *h_CUBLAS = (__half *) malloc(mem_size_C);
 
     checkCudaErrors(cudaMalloc((void **) &d_A, mem_size_A));
     checkCudaErrors(cudaMalloc((void **) &d_B, mem_size_B));
@@ -243,19 +292,30 @@ int matrixMultiply(int argc, char **argv, int devID, sMatrixSize &matrix_size)
     printf("Computing result using CUBLAS...");
 
     // execute the kernel
-    int nIter = 30;
+    int nIter = 9999999;
 
     // CUBLAS version 2.0
     {
-        const float alpha = 1.0f;
-        const float beta  = 0.0f;
+        cublasOperation_t trans_A = (ta) ? CUBLAS_OP_T : CUBLAS_OP_N;
+        cublasOperation_t trans_B = (tb) ? CUBLAS_OP_T : CUBLAS_OP_N;
+        int m = matrix_size.uiWC;
+        int n = matrix_size.uiHC;
+        int k = matrix_size.uiWA;
+        int lda = (trans_A == CUBLAS_OP_N) ? k : n;
+        int ldb = (trans_B == CUBLAS_OP_N) ? m : k;
+        int ldc = m;
+        const __half alpha = 1.0f;
+        const __half beta  = 0.0f;
         cublasHandle_t handle;
         cudaEvent_t start, stop;
 
         checkCudaErrors(cublasCreate(&handle));
 
         //Perform warmup operation with cublas
-        checkCudaErrors(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, matrix_size.uiWB, matrix_size.uiHA, matrix_size.uiWA, &alpha, d_B, matrix_size.uiWB, d_A, matrix_size.uiWA, &beta, d_C, matrix_size.uiWB));
+        checkCudaErrors(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+        cublasGemmAlgo_t algo = cublasGemmAlgo_t::CUBLAS_GEMM_ALGO4_TENSOR_OP;
+        checkCudaErrors(cublasGemmEx(handle, trans_B, trans_A, m, n, k, &alpha, d_B, CUDA_R_16F, 
+                    ldb, d_A, CUDA_R_16F, lda, &beta, d_C, CUDA_R_16F, ldc, CUBLAS_COMPUTE_16F, algo));
 
         // Allocate CUDA events that we'll use for timing
         checkCudaErrors(cudaEventCreate(&start));
@@ -268,8 +328,8 @@ int matrixMultiply(int argc, char **argv, int devID, sMatrixSize &matrix_size)
         {
             //note cublas is column primary!
             //need to transpose the order
-            checkCudaErrors(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, matrix_size.uiWB, matrix_size.uiHA, matrix_size.uiWA, &alpha, d_B, matrix_size.uiWB, d_A, matrix_size.uiWA, &beta, d_C, matrix_size.uiWB));
-
+            checkCudaErrors(cublasGemmEx(handle, trans_B, trans_A, m, n, k, &alpha, d_B, CUDA_R_16F, 
+                        ldb, d_A, CUDA_R_16F, lda, &beta, d_C, CUDA_R_16F, ldc, CUBLAS_COMPUTE_16F, algo));
         }
 
         printf("done.\n");
@@ -302,16 +362,16 @@ int matrixMultiply(int argc, char **argv, int devID, sMatrixSize &matrix_size)
 
     // compute reference solution
     printf("Computing result using host CPU...");
-    float *reference = (float *)malloc(mem_size_C);
-    matrixMulCPU(reference, h_A, h_B, matrix_size.uiHA, matrix_size.uiWA, matrix_size.uiWB);
+    __half *reference = (__half *)malloc(mem_size_C);
+    matrixMulCPU(reference, h_A, h_B, matrix_size.uiHA, matrix_size.uiWA, matrix_size.uiWB, ta, tb);
     printf("done.\n");
 
     // check result (CUBLAS)
-    bool resCUBLAS = sdkCompareL2fe(reference, h_CUBLAS, size_C, 1.0e-6f);
+    bool resCUBLAS = sdkCompareL2fe(reference, h_CUBLAS, size_C, 1.0e-2f);
 
     if (resCUBLAS != true)
     {
-        printDiff(reference, h_CUBLAS, matrix_size.uiWC, matrix_size.uiHC, 100, 1.0e-5f);
+        printDiff(reference, h_CUBLAS, matrix_size.uiWC, matrix_size.uiHC, 100, 1.0e-2f);
     }
 
     printf("Comparing CUBLAS Matrix Multiply with CPU results: %s\n", (true == resCUBLAS) ? "PASS" : "FAIL");
@@ -349,7 +409,7 @@ int main(int argc, char **argv)
 
     initializeCUDA(argc, argv, devID, sizeMult, matrix_size);
 
-    int matrix_result = matrixMultiply(argc, argv, devID, matrix_size);
+    int matrix_result = matrixMultiply(argc, argv, devID, matrix_size, true, false);
 
     return matrix_result;
 }
-- 
2.25.1

