Use cublasHgemm "back" for fp16 computation with Volta GPU (microsoft#3765)

* Use cublasHgemm for fp16 computation with Volta GPU

Author: pengwa

onnxruntime/contrib_ops/cuda/bert/attention.cc

@@ -71,18 +71,19 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
 
   // Bias shape is (N), broadcast using B(N, M) = 1 * bias(N, 1) x ones(1, M) + 0 * B.
   // TODO: use custom kernel of expand to improve the performance.
+  auto& device_prop = GetDeviceProp();
   CUBLAS_RETURN_IF_ERROR(cublasGemmHelper(
       cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, 1, &one,
       reinterpret_cast<const CudaT*>(bias->template Data<T>()), n,
       GetConstOnes<CudaT>(m), 1,
-      &zero, reinterpret_cast<CudaT*>(gemm_buffer.get()), n));
+      &zero, reinterpret_cast<CudaT*>(gemm_buffer.get()), n, device_prop));
 
   // Gemm, note that CUDA assumes col-major, so result(N, M) = 1 * weights x input + 1 x B.
   CUBLAS_RETURN_IF_ERROR(cublasGemmHelper(
       cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &one,
       reinterpret_cast<const CudaT*>(weights->template Data<T>()), n,
       reinterpret_cast<const CudaT*>(input->template Data<T>()), k,
-      &one, reinterpret_cast<CudaT*>(gemm_buffer.get()), n));
+      &one, reinterpret_cast<CudaT*>(gemm_buffer.get()), n, device_prop));
 
   size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, num_heads_, head_size, sequence_length);
   auto temp_buffer = GetScratchBuffer<void>(workSpaceSize);

onnxruntime/core/providers/cuda/math/gemm.cc

@@ -65,12 +65,11 @@ Status Gemm<T>::ComputeInternal(OpKernelContext* ctx) const {
 
   CudaT one = ToCudaType<T>::FromFloat(1.0f);
   CudaT zero = ToCudaType<T>::FromFloat(0.0f);
-
+  auto& device_prop = GetDeviceProp();
   // broadcast bias if needed and is present
   if (beta_ != 0 && B != nullptr) {
     auto& b_shape = B->Shape();
     const CudaT* b_data = reinterpret_cast<const CudaT*>(B->template Data<T>());
-
     if (b_shape.Size() == 1) {
       // if B is (), (1,) or (1, 1), broadcast the scalar
       CUBLAS_RETURN_IF_ERROR(cublasCopyHelper(
@@ -91,7 +90,7 @@ Status Gemm<T>::ComputeInternal(OpKernelContext* ctx) const {
           b_data, N,
           GetConstOnes<CudaT>(M), 1,
           /*beta*/ &zero,
-          out_data, N));
+          out_data, N, device_prop));
     } else if (b_shape.NumDimensions() == 2 && b_shape[1] == 1) {
       // B is (M, 1), broadcast using Y(N,M) = 1 * ones(N,1) x B(1,M) + 0 * Y
       CUBLAS_RETURN_IF_ERROR(cublasGemmHelper(
@@ -103,7 +102,7 @@ Status Gemm<T>::ComputeInternal(OpKernelContext* ctx) const {
           GetConstOnes<CudaT>(N), N,
           b_data, 1,
           /*beta*/ &zero,
-          out_data, N));
+          out_data, N, device_prop));
     } else {
       // B is (M, N), no broadcast needed.
       CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(out_data, b_data, M * N * sizeof(float), cudaMemcpyDeviceToDevice));
@@ -126,7 +125,7 @@ Status Gemm<T>::ComputeInternal(OpKernelContext* ctx) const {
       // ideally we need to set the output buffer contents to 0 if bias is missing,
       // but passing 0 for beta is cheaper and it will ignore any junk in the output buffer
       B != nullptr ? &beta : &zero,
-      out_data, N));
+      out_data, N, device_prop));
 
   return Status::OK();
 }

onnxruntime/core/providers/cuda/math/matmul.cc

@@ -49,7 +49,7 @@ static bool CanUseStridedBatchedGemm(const TensorShape& left_shape, const Tensor
   int64_t left_k = transa ? left_shape[left_num_dims - 2] : left_shape[left_num_dims - 1];
 
   if (right_num_dims >= 3) {
-    int64_t right_p = right_shape.SizeToDimension(right_num_dims-2);
+    int64_t right_p = right_shape.SizeToDimension(right_num_dims - 2);
     if (left_p != right_p) {
       return false;
     }
@@ -102,7 +102,7 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
   const int ldb = transb ? static_cast<int>(helper.K()) : static_cast<int>(helper.N());
   const int ldc = static_cast<int>(helper.N());
   int64_t stride_A, stride_B, stride_C, batch_count;
-
+  auto& device_prop = GetDeviceProp();
   if (helper.OutputOffsets().size() == 1) {
     CUBLAS_RETURN_IF_ERROR(cublasGemmHelper(
         Base::CublasHandle(),
@@ -118,10 +118,11 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
         lda,
         &zero,
         reinterpret_cast<CudaT*>(Y->template MutableData<T>()),
-        ldc));
+        ldc,
+        device_prop));
     return Status::OK();
   } else if (CanUseStridedBatchedGemm(left_X->Shape(), right_X->Shape(),
-             transa, transb, stride_A, stride_B, stride_C, batch_count)) {
+                                      transa, transb, stride_A, stride_B, stride_C, batch_count)) {
     CUBLAS_RETURN_IF_ERROR(cublasGemmStridedBatchedHelper(Base::CublasHandle(),
                                                           transB,
                                                           transA,

onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h

@@ -18,20 +18,29 @@
 // Generalize library calls to be use in template functions
 
 // gemm
-inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, const float* alpha, const float* A, int lda, const float* B, int ldb, const float* beta, float* C, int ldc) {
+inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
+                                       int m, int n, int k, const float* alpha, const float* A, int lda,
+                                       const float* B, int ldb, const float* beta, float* C, int ldc,
+                                       const cudaDeviceProp& /*prop*/) {
   return cublasSgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
 }
-inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, const double* alpha, const double* A, int lda, const double* B, int ldb, const double* beta, double* C, int ldc) {
+inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
+                                       int m, int n, int k, const double* alpha, const double* A, int lda,
+                                       const double* B, int ldb, const double* beta, double* C, int ldc,
+                                       const cudaDeviceProp& /*prop*/) {
   return cublasDgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
 }
-inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, int k, const half* alpha, const half* A, int lda, const half* B, int ldb, const half* beta, half* C, int ldc) {
-  // Disable below to make sure merged result is on par with before-merge.
+inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
+                                       int m, int n, int k, const half* alpha, const half* A, int lda,
+                                       const half* B, int ldb, const half* beta, half* C, int ldc,
+                                       const cudaDeviceProp& prop) {
   // This does true FP16 computation which is slow for non-Volta GPUs
-  //if (onnxruntime::cuda::DeviceProp().GetDeviceProps().major >= 7) {
-  //   onnxruntime::cuda::CublasMathModeSetter math_mode_setter( handle, CUBLAS_TENSOR_OP_MATH );
-  //  return cublasHgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
-  //}
-  // This does pseudo FP16 computation (input/output in fp16, computation in fp32)
+  if (prop.major >= 7) {
+    onnxruntime::cuda::CublasMathModeSetter math_mode_setter(handle, CUBLAS_TENSOR_OP_MATH);
+    return cublasHgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
+  }
+
+  //This does pseudo FP16 computation (input/output in fp16, computation in fp32)
   float h_a = onnxruntime::math::halfToFloat(*reinterpret_cast<const uint16_t*>(alpha));
   float h_b = onnxruntime::math::halfToFloat(*reinterpret_cast<const uint16_t*>(beta));
   cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH);
@@ -79,7 +88,7 @@ inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle,
                                                      const double* beta,
                                                      double* C, int ldc,
                                                      long long int strideC,
-                                                     int batch_count){
+                                                     int batch_count) {
   return cublasDgemmStridedBatched(handle, transa, transb, m, n, k, alpha, A, lda, strideA, B, ldb, strideB, beta, C, ldc, strideC, batch_count);
 }