merge mirror master

apex/contrib/csrc/multihead_attn/strided_batched_gemm.cuh

@@ -7,8 +7,6 @@
 //#include <cuda_profiler_api.h>
 #include <cuda_runtime.h>
 
-#include <rocblas/rocblas.h>
-
 //#include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/cuda/Exceptions.h>
@@ -47,52 +45,6 @@ cublasOperation_t convertTransToCublasOperation(char trans) {
   }
 }
 
-// needed to work around calling rocblas API instead of hipblas API
-static rocblas_operation hipOperationToRocOperation(hipblasOperation_t op)
-{
-    switch(op)
-    {
-    case HIPBLAS_OP_N:
-        return rocblas_operation_none;
-    case HIPBLAS_OP_T:
-        return rocblas_operation_transpose;
-    case HIPBLAS_OP_C:
-        return rocblas_operation_conjugate_transpose;
-    }
-    AT_ERROR("HIPBLAS_STATUS_INVALID_ENUM");
-}
-
-static hipblasStatus_t rocBLASStatusToHIPStatus(rocblas_status error)
-{
-    switch(error)
-    {
-    case rocblas_status_size_unchanged:
-    case rocblas_status_size_increased:
-    case rocblas_status_success:
-    case rocblas_status_continue:
-        return HIPBLAS_STATUS_SUCCESS;
-    case rocblas_status_invalid_handle:
-        return HIPBLAS_STATUS_NOT_INITIALIZED;
-    case rocblas_status_not_implemented:
-    case rocblas_status_excluded_from_build:
-        return HIPBLAS_STATUS_NOT_SUPPORTED;
-    case rocblas_status_invalid_pointer:
-    case rocblas_status_invalid_size:
-    case rocblas_status_invalid_value:
-    case rocblas_status_size_query_mismatch:
-        return HIPBLAS_STATUS_INVALID_VALUE;
-    case rocblas_status_memory_error:
-        return HIPBLAS_STATUS_ALLOC_FAILED;
-    case rocblas_status_internal_error:
-    case rocblas_status_perf_degraded:
-    case rocblas_status_check_numerics_fail:
-        return HIPBLAS_STATUS_INTERNAL_ERROR;
-    case rocblas_status_arch_mismatch:
-        return HIPBLAS_STATUS_ARCH_MISMATCH;
-    }
-    AT_ERROR("HIPBLAS_STATUS_INVALID_ENUM");
-}
-
 void RocblasStridedBatchedGemm(char transa, char transb, long m, long n, long k,
                     float alpha, const half *a, long lda, long strideA, const half *b, long ldb, long strideB,
                     float beta, half *c, long ldc, long strideC, half *d, long ldd, long strideD, long batchCount, rocblas_gemm_algo algo, rocblas_int flags) {
@@ -105,13 +57,13 @@ void RocblasStridedBatchedGemm(char transa, char transb, long m, long n, long k,
     float fAlpha = alpha;
     float fBeta = beta;
     //THCublasCheck(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
-    TORCH_CUDABLAS_CHECK(rocBLASStatusToHIPStatus(rocblas_gemm_strided_batched_ex((rocblas_handle)handle,
-                                     hipOperationToRocOperation(opa), hipOperationToRocOperation(opb), (int)m, (int)n, (int)k,
+    TORCH_CUDABLAS_CHECK(rocblas_gemm_strided_batched_ex(handle,
+                                     opa, opb, (int)m, (int)n, (int)k,
                                      (void*)&fAlpha, a, rocblas_datatype_f16_r /*a_type*/, (int)lda, strideA,
                                      b, rocblas_datatype_f16_r /*b_type*/, (int)ldb, strideB,
                                      (void*)&fBeta, c, rocblas_datatype_f16_r /*c_type*/, (int)ldc, strideC,
                                      d, rocblas_datatype_f16_r /*d_type*/, int(ldd), strideD,
-                                     (int)batchCount, rocblas_datatype_f32_r /*compute_type*/, algo, 0 /*solution_index*/, flags)));
+                                     (int)batchCount, rocblas_datatype_f32_r /*compute_type*/, algo, 0 /*solution_index*/, flags));
 }
 
 void gemm_switch_fp32accum(char transa, char transb, long m, long n, long k,

csrc/fused_dense_cuda.cu

@@ -10,22 +10,10 @@
 #include <cublas_v2.h>
 #include <cuda_runtime.h>
 
-#include "utils.h"
-#include <rocblas/rocblas.h>
-
 #if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000
 // includes cublaslt
 #include <cublasLt.h>
 #endif
-
-// until we use hipblas v2
-// hipify correctly maps things like CUDA_R_16F to HIP_R_16F,
-// however hipblas v1 is still using its custom type
-#define HIP_R_64F  HIPBLAS_R_64F
-#define HIP_R_32F  HIPBLAS_R_32F
-#define HIP_R_16F  HIPBLAS_R_16F
-
-
 // FP64 Wrapper around cublas GEMMEx
 cublasStatus_t gemm_bias(
     cublasHandle_t handle,
@@ -42,6 +30,33 @@ cublasStatus_t gemm_bias(
     const float* beta,
     double* C,
     int ldc) {
+#ifdef __HIP_PLATFORM_HCC__
+  return rocblas_gemm_ex(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      rocblas_datatype_f64_r,
+      lda,
+      B,
+      rocblas_datatype_f64_r,
+      ldb,
+      beta,
+      C,
+      rocblas_datatype_f64_r,
+      ldc,
+      C,
+      rocblas_datatype_f64_r,
+      ldc,
+      rocblas_datatype_f64_r,
+      rocblas_gemm_algo_standard,
+      0,
+      0);  
+#else
   return cublasGemmEx(
       handle,
       transa,
@@ -62,6 +77,7 @@ cublasStatus_t gemm_bias(
       ldc,
       CUDA_R_64F,
       CUBLAS_GEMM_DEFAULT);
+#endif
 }
 
 // FP32 Wrapper around cublas GEMMEx
@@ -80,6 +96,34 @@ cublasStatus_t gemm_bias(
     const float* beta,
     float* C,
     int ldc) {
+#ifdef __HIP_PLATFORM_HCC__
+  return rocblas_gemm_ex(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      rocblas_datatype_f32_r,
+      lda,
+      B,
+      rocblas_datatype_f32_r,
+      ldb,
+      beta,
+      C,
+      rocblas_datatype_f32_r,
+      ldc,
+      C,
+      rocblas_datatype_f32_r,
+      ldc,
+      rocblas_datatype_f32_r,
+      rocblas_gemm_algo_standard,
+      0,
+      0);
+
+#else
   return cublasGemmEx(
       handle,
       transa,
@@ -100,6 +144,7 @@ cublasStatus_t gemm_bias(
       ldc,
       CUDA_R_32F,
       CUBLAS_GEMM_DEFAULT);
+#endif
 }
 
 // FP16 Tensor core wrapper around cublas GEMMEx
@@ -118,6 +163,7 @@ cublasStatus_t gemm_bias(
     const float* beta,
     at::Half* C,
     int ldc) {
+<<<<<<< HEAD
   if (parseEnvVarFlag("APEX_ROCBLAS_GEMM_ALLOW_HALF")) {
     half h_alpha = __float2half(*alpha);
     half h_beta = __float2half(*beta);
@@ -163,6 +209,56 @@ cublasStatus_t gemm_bias(
         CUDA_R_32F,
         CUBLAS_GEMM_DEFAULT_TENSOR_OP);
   }
+=======
+#ifdef __HIP_PLATFORM_HCC__
+  return rocblas_gemm_ex(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      rocblas_datatype_f16_r,
+      lda,
+      B,
+      rocblas_datatype_f16_r,
+      ldb,
+      beta,
+      C,
+      rocblas_datatype_f16_r,
+      ldc,
+      C,
+      rocblas_datatype_f16_r,
+      ldc,
+      rocblas_datatype_f32_r,
+      rocblas_gemm_algo_standard,
+      0,
+      0);
+#else
+  return cublasGemmEx(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      CUDA_R_16F,
+      lda,
+      B,
+      CUDA_R_16F,
+      ldb,
+      beta,
+      C,
+      CUDA_R_16F,
+      ldc,
+      CUDA_R_32F,
+      CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+#endif
+>>>>>>> mirror/master
 }
 
 

csrc/mlp_cuda.cu

@@ -13,8 +13,6 @@
 #include <cuda_runtime.h>
 #include "utils.h"
 
-#include <rocblas/rocblas.h>
-
 #if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000
 // includes cublaslt
 #include <cublasLt.h>
@@ -62,52 +60,6 @@ __device__ __inline__ float sigmoid(float a) {
   return (retf);
 }
 
-// needed to work around calling rocblas API instead of hipblas API
-static rocblas_operation hipOperationToRocOperation(hipblasOperation_t op)
-{
-    switch(op)
-    {
-    case HIPBLAS_OP_N:
-        return rocblas_operation_none;
-    case HIPBLAS_OP_T:
-        return rocblas_operation_transpose;
-    case HIPBLAS_OP_C:
-        return rocblas_operation_conjugate_transpose;
-    }
-    AT_ERROR("HIPBLAS_STATUS_INVALID_ENUM");
-}
-
-static hipblasStatus_t rocBLASStatusToHIPStatus(rocblas_status error)
-{
-    switch(error)
-    {
-    case rocblas_status_size_unchanged:
-    case rocblas_status_size_increased:
-    case rocblas_status_success:
-    case rocblas_status_continue:
-        return HIPBLAS_STATUS_SUCCESS;
-    case rocblas_status_invalid_handle:
-        return HIPBLAS_STATUS_NOT_INITIALIZED;
-    case rocblas_status_not_implemented:
-    case rocblas_status_excluded_from_build:
-        return HIPBLAS_STATUS_NOT_SUPPORTED;
-    case rocblas_status_invalid_pointer:
-    case rocblas_status_invalid_size:
-    case rocblas_status_invalid_value:
-    case rocblas_status_size_query_mismatch:
-        return HIPBLAS_STATUS_INVALID_VALUE;
-    case rocblas_status_memory_error:
-        return HIPBLAS_STATUS_ALLOC_FAILED;
-    case rocblas_status_internal_error:
-    case rocblas_status_perf_degraded:
-    case rocblas_status_check_numerics_fail:
-        return HIPBLAS_STATUS_INTERNAL_ERROR;
-    case rocblas_status_arch_mismatch:
-	return HIPBLAS_STATUS_ARCH_MISMATCH;
-    }
-    AT_ERROR("HIPBLAS_STATUS_INVALID_ENUM");
-}
-
 // FP64 Wrapper around cublas GEMMEx
 cublasStatus_t mlp_gemm(
     cublasHandle_t handle,
@@ -126,10 +78,10 @@ cublasStatus_t mlp_gemm(
     int ldc,
     int flag) {
 #ifdef __HIP_PLATFORM_HCC__
-  return rocBLASStatusToHIPStatus(rocblas_gemm_ex(
-      (rocblas_handle) handle,
-      hipOperationToRocOperation(transa),
-      hipOperationToRocOperation(transb),
+  return rocblas_gemm_ex(
+      handle,
+      transa,
+      transb,
       m,
       n,
       k,
@@ -150,7 +102,7 @@ cublasStatus_t mlp_gemm(
       rocblas_datatype_f64_r,
       rocblas_gemm_algo_standard,
       0,
-      flag));  
+      flag);  
 #else
   return cublasGemmEx(
       handle,
@@ -193,10 +145,10 @@ cublasStatus_t mlp_gemm(
     int ldc,
     int flag) {
 #ifdef __HIP_PLATFORM_HCC__
-  return rocBLASStatusToHIPStatus(rocblas_gemm_ex(
-      (rocblas_handle) handle,
-      hipOperationToRocOperation(transa),
-      hipOperationToRocOperation(transb),
+  return rocblas_gemm_ex(
+      handle,
+      transa,
+      transb,
       m,
       n,
       k,
@@ -217,7 +169,7 @@ cublasStatus_t mlp_gemm(
       rocblas_datatype_f32_r,
       rocblas_gemm_algo_standard,
       0,
-      flag));
+      flag);
 
 #else
   return cublasGemmEx(
@@ -261,61 +213,31 @@ cublasStatus_t mlp_gemm(
     int ldc,
     int flag) {
 #ifdef __HIP_PLATFORM_HCC__
-  if (parseEnvVarFlag("APEX_ROCBLAS_GEMM_ALLOW_HALF")) {
-    half h_alpha = __float2half(*alpha);
-    half h_beta = __float2half(*beta);
-    return rocBLASStatusToHIPStatus(rocblas_gemm_ex(
-        (rocblas_handle) handle,
-        hipOperationToRocOperation(transa),
-        hipOperationToRocOperation(transb),
-        m,
-        n,
-        k,
-        /* alpha */ &h_alpha,
-        A,
-        rocblas_datatype_f16_r,
-        lda,
-        B,
-        rocblas_datatype_f16_r,
-        ldb,
-        /* beta */ &h_beta,
-        C,
-        rocblas_datatype_f16_r,
-        ldc,
-        C,
-        rocblas_datatype_f16_r,
-        ldc,
-        /* rocblas_datatype_f32_r */ rocblas_datatype_f16_r,
-        rocblas_gemm_algo_standard,
-        0,
-        flag);
-  } else {
-    return rocBLASStatusToHIPStatus(rocblas_gemm_ex(
-        (rocblas_handle) handle,
-        hipOperationToRocOperation(transa),
-        hipOperationToRocOperation(transb),
-        m,
-        n,
-        k,
-        alpha,
-        A,
-        rocblas_datatype_f16_r,
-        lda,
-        B,
-        rocblas_datatype_f16_r,
-        ldb,
-        beta,
-        C,
-        rocblas_datatype_f16_r,
-        ldc,
-        C,
-        rocblas_datatype_f16_r,
-        ldc,
-        rocblas_datatype_f32_r,
-        rocblas_gemm_algo_standard,
-        0,
-        flag);
-  }
+  return rocblas_gemm_ex(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      rocblas_datatype_f16_r,
+      lda,
+      B,
+      rocblas_datatype_f16_r,
+      ldb,
+      beta,
+      C,
+      rocblas_datatype_f16_r,
+      ldc,
+      C,
+      rocblas_datatype_f16_r,
+      ldc,
+      rocblas_datatype_f32_r,
+      rocblas_gemm_algo_standard,
+      0,
+      flag);
 #else
   return cublasGemmEx(
       handle,