Add support for tensor cores in float16.

Author: abergeron

src/gpuarray_blas_cuda_cublas.c

@@ -443,7 +443,7 @@ static int hgemm(cb_order order, cb_transpose transA, cb_transpose transB,
   ASSERT_BUF(B);
   ASSERT_BUF(C);
 
-  if (cublasSgemmEx == NULL)
+  if (cublasGemmEx == NULL && cublasSgemmEx == NULL)
     return error_set(ctx->err, GA_DEVSUP_ERROR, "cublasSgemmEx unavailable");
 
   if (LARGE_VAL(M) || LARGE_VAL(N) || LARGE_VAL(K) ||
@@ -476,16 +476,30 @@ static int hgemm(cb_order order, cb_transpose transA, cb_transpose transB,
   GA_CUDA_EXIT_ON_ERROR(ctx, cuda_wait(B, CUDA_WAIT_READ));
   GA_CUDA_EXIT_ON_ERROR(ctx, cuda_wait(C, CUDA_WAIT_ALL));
 
-  CUBLAS_EXIT_ON_ERROR(ctx, cublasSgemmEx(h->h, convT(transA), convT(transB),
-                                          M, N, K,
-                                          &alpha, ((uint16_t *)A->ptr) + offA,
-                                          CUDA_R_16F,
-                                          lda, ((uint16_t *)B->ptr) + offB,
-                                          CUDA_R_16F,
-                                          ldb, &beta, ((uint16_t *)C->ptr) + offC,
-                                          CUDA_R_16F,
-                                          ldc));
-
+  if (cublasGemmEx) {
+    CUBLAS_EXIT_ON_ERROR(ctx, cublasGemmEx(h->h, convT(transA), convT(transB),
+					   M, N, K,
+					   &alpha, ((uint16_t *)A->ptr) + offA,
+					   CUDA_R_16F,
+					   lda, ((uint16_t *)B->ptr) + offB,
+					   CUDA_R_16F,
+					   ldb, &beta, ((uint16_t *)C->ptr) + offC,
+					   CUDA_R_16F,
+					   ldc,
+					   CUDA_R_32F,
+					   CUBLAS_GEMM_DFALT_TENSOR_OP));
+  } else {
+    CUBLAS_EXIT_ON_ERROR(ctx, cublasSgemmEx(h->h, convT(transA), convT(transB),
+					    M, N, K,
+					    &alpha, ((uint16_t *)A->ptr) + offA,
+					    CUDA_R_16F,
+					    lda, ((uint16_t *)B->ptr) + offB,
+					    CUDA_R_16F,
+					    ldb, &beta, ((uint16_t *)C->ptr) + offC,
+					    CUDA_R_16F,
+					    ldc));
+  }
+    
   GA_CUDA_EXIT_ON_ERROR(ctx, cuda_record(A, CUDA_WAIT_READ));
   GA_CUDA_EXIT_ON_ERROR(ctx, cuda_record(B, CUDA_WAIT_READ));
   GA_CUDA_EXIT_ON_ERROR(ctx, cuda_record(C, CUDA_WAIT_ALL));

src/loaders/libcublas.fn

@@ -21,6 +21,8 @@ DEF_PROC_V2(cublasDger, (cublasHandle_t handle, int m, int n, const double *alph
 
 DEF_PROC_OPT(cublasSgemmEx, (cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, const float *alpha, const void *A, cudaDataType Atype, int lda, const void *B, cudaDataType Btype, int ldb, const float *beta, void *C, cudaDataType Ctype, int ldc));
 
+DEF_PROC_OPT(cublasGemmEx, (cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, const void *alpha, const void *A, cudaDataType_t Atype, int lda, const void *B, cudaDataType_t Btype, int ldb, const void *beta, void *C, cudaDataType_t Ctype, int ldc, cudaDataType_t computeType, cublasGemmAlgo_t algo));
+
 DEF_PROC(cublasSgemmBatched, (cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, const float *alpha, const float *Aarray[], int lda, const float *Barray[], int ldb, const float *beta, float *Carray[], int ldc, int batchCount));
 DEF_PROC(cublasDgemmBatched, (cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, const double *alpha, const double *Aarray[], int lda, const double *Barray[], int ldb, const double *beta, double *Carray[], int ldc, int batchCount));
 

src/loaders/libcublas.h

@@ -34,6 +34,34 @@ typedef enum cudaDataType_t
   CUDA_C_32U= 13   // complex as a pair of unsigned int numbers
 } cudaDataType;
 
+typedef cudaDataType cudaDataType_t;
+
+typedef enum {
+    CUBLAS_GEMM_DFALT               = -1,
+    CUBLAS_GEMM_ALGO0               =  0,
+    CUBLAS_GEMM_ALGO1               =  1,
+    CUBLAS_GEMM_ALGO2               =  2,
+    CUBLAS_GEMM_ALGO3               =  3,
+    CUBLAS_GEMM_ALGO4               =  4,
+    CUBLAS_GEMM_ALGO5               =  5,
+    CUBLAS_GEMM_ALGO6               =  6,
+    CUBLAS_GEMM_ALGO7               =  7,
+    CUBLAS_GEMM_ALGO8               =  8,
+    CUBLAS_GEMM_ALGO9               =  9,
+    CUBLAS_GEMM_ALGO10              =  10,
+    CUBLAS_GEMM_ALGO11              =  11,
+    CUBLAS_GEMM_ALGO12              =  12,
+    CUBLAS_GEMM_ALGO13              =  13,
+    CUBLAS_GEMM_ALGO14              =  14,
+    CUBLAS_GEMM_ALGO15              =  15,
+    CUBLAS_GEMM_ALGO16              =  16,
+    CUBLAS_GEMM_ALGO17              =  17,
+    CUBLAS_GEMM_DFALT_TENSOR_OP     =  99,
+    CUBLAS_GEMM_ALGO0_TENSOR_OP     =  100,
+    CUBLAS_GEMM_ALGO1_TENSOR_OP     =  101,
+    CUBLAS_GEMM_ALGO2_TENSOR_OP     =  102
+} cublasGemmAlgo_t;
+
 typedef struct CUstream_st *cudaStream_t;
 
 typedef enum {