diff --git a/faiss/gpu/utils/MatrixMult-inl.cuh b/faiss/gpu/utils/MatrixMult-inl.cuh index b28ea56da..a0ef20f6c 100644 --- a/faiss/gpu/utils/MatrixMult-inl.cuh +++ b/faiss/gpu/utils/MatrixMult-inl.cuh @@ -22,19 +22,26 @@ struct GetCudaType; #ifdef USE_AMD_ROCM +#if ((hipblasVersionMajor == 2 && defined(HIPBLAS_V2)) || \ + hipblasVersionMajor >= 3) +using hipDataTypeCompat = hipDataType; +#else +using hipDataTypeCompat = hipblasDatatype_t; +#endif + template <> struct GetCudaType { - static constexpr hipblasDatatype_t Type = HIPBLAS_R_32F; + static constexpr hipDataTypeCompat Type = HIPBLAS_R_32F; }; template <> struct GetCudaType { - static constexpr hipblasDatatype_t Type = HIPBLAS_R_16F; + static constexpr hipDataTypeCompat Type = HIPBLAS_R_16F; }; template <> struct GetCudaType<__hip_bfloat16> { - static constexpr hipblasDatatype_t Type = HIPBLAS_R_16B; + static constexpr hipDataTypeCompat Type = HIPBLAS_R_16B; }; #else @@ -76,6 +83,13 @@ cublasStatus_t rawGemm( auto cBT = GetCudaType::Type; #ifdef USE_AMD_ROCM +#if ((hipblasVersionMajor == 2 && defined(HIPBLAS_V2)) || \ + hipblasVersionMajor >= 3) + auto computeType = HIPBLAS_COMPUTE_32F; +#else + auto computeType = HIPBLAS_R_32F; +#endif + return hipblasGemmEx( handle, transa, @@ -94,7 +108,7 @@ cublasStatus_t rawGemm( C, HIPBLAS_R_32F, ldc, - HIPBLAS_R_32F, + computeType, HIPBLAS_GEMM_DEFAULT); #else @@ -177,6 +191,13 @@ cublasStatus_t rawBatchGemm( // Always accumulate in f32 #ifdef USE_AMD_ROCM +#if ((hipblasVersionMajor == 2 && defined(HIPBLAS_V2)) || \ + hipblasVersionMajor >= 3) + auto computeType = HIPBLAS_COMPUTE_32F; +#else + auto computeType = HIPBLAS_R_32F; +#endif + return hipblasGemmStridedBatchedEx( handle, transa, @@ -199,7 +220,7 @@ cublasStatus_t rawBatchGemm( ldc, strideC, batchCount, - HIPBLAS_R_32F, + computeType, HIPBLAS_GEMM_DEFAULT); #else return cublasGemmStridedBatchedEx(