Add more BLAS tests (#601)

pxl-th · web-flow · commit 44e581885f1e · 2024-02-22T13:57:14.000+02:00
diff --git a/.codecov.yml b/.codecov.yml
@@ -1,6 +1,8 @@
 coverage:
   ignore:
+    - "src/*/lib*.jl"
     - "src/device"
+    - "docs/"
   status:
     patch:
       default:
diff --git a/src/blas/highlevel.jl b/src/blas/highlevel.jl
@@ -2,11 +2,9 @@ rocblas_size(t::Char, M::ROCVecOrMat) = (size(M, t=='N' ? 1 : 2), size(M, t=='N'
 
 const ROCBLASArray{T<:ROCBLASFloat} = ROCArray{T}
 
-###########
 #
 # BLAS 1
 #
-###########
 
 LinearAlgebra.rmul!(x::ROCArray{<:ROCBLASFloat}, k::Number) =
   scal!(length(x), convert(eltype(x), k), x, 1)
@@ -84,11 +82,9 @@ function LinearAlgebra.reflect!(
     x, y
 end
 
-############
 #
 # BLAS 2
 #
-############
 
 if VERSION ≥ v"1.10-"
     # multiplication
@@ -145,36 +141,56 @@ else
     end
 end
 
-#########
 # GEMV
-##########
 
-function gemv_wrapper!(
-    y::ROCVector{T}, tA::Char, A::ROCMatrix{T}, x::ROCVector{T},
-    alpha = one(T), beta = zero(T),
-) where T <: ROCBLASFloat
-    mA, nA = rocblas_size(tA, A)
-    if nA != length(x)
-        throw(DimensionMismatch("second dimension of A, $nA, does not match length of x, $(length(x))"))
-    end
-    if mA != length(y)
-        throw(DimensionMismatch("first dimension of A, $mA, does not match length of y, $(length(y))"))
+function LinearAlgebra.generic_matvecmul!(
+    Y::ROCVector, tA::AbstractChar, A::StridedROCMatrix, B::StridedROCVector,
+    _add::MulAddMul,
+)
+    mA, nA = tA == 'N' ? size(A) : reverse(size(A))
+
+    nA != length(B) && throw(DimensionMismatch(
+        "second dimension of A, $nA, does not match length of B, $(length(B))"))
+    mA != length(Y) && throw(DimensionMismatch(
+        "first dimension of A, $mA, does not match length of Y, $(length(Y))"))
+
+    mA == 0 && return Y
+    nA == 0 && return rmul!(Y, 0)
+
+    T = eltype(Y)
+    alpha, beta = _add.alpha, _add.beta
+    if alpha isa Union{Bool,T} && beta isa Union{Bool,T}
+        α, β = T(alpha), T(beta)
+        if T <: ROCBLASFloat && eltype(A) == eltype(B) == T
+            if tA in ('N', 'T', 'C')
+                return gemv!(tA, α, A, B, β, Y)
+            elseif tA in ('S', 's')
+                return symv!(tA == 'S' ? 'U' : 'L', α, A, B, β, Y)
+            elseif tA in ('H', 'h')
+                return hemv!(tA == 'H' ? 'U' : 'L', α, A, B, β, Y)
+            end
+        end
     end
-    mA == 0 && return y
-    nA == 0 && return rmul!(y, 0)
-    gemv!(tA, alpha, A, x, beta, y)
+    LinearAlgebra.generic_matmatmul!(Y, tA, 'N', A, B, MulAddMul(alpha, beta))
 end
 
-LinearAlgebra.mul!(Y::ROCVector{T}, A::ROCMatrix{T}, B::ROCVector{T}) where T<:ROCBLASFloat = gemv_wrapper!(Y, 'N', A,  B)
-LinearAlgebra.lmul!(Y::ROCVector{T}, A::LinearAlgebra.Transpose{<:Any, ROCMatrix{T}}, B::ROCVector{T}) where T<:ROCBLASFloat = gemv_wrapper!(Y, 'T', A.parent, B)
-LinearAlgebra.lmul!(Y::ROCVector{T}, A::LinearAlgebra.Adjoint{<:Any, ROCMatrix{T}}, B::ROCVector{T}) where T<:ROCBLASFloat = gemv_wrapper!(Y, 'T', A.parent, B)
-LinearAlgebra.lmul!(Y::ROCVector{T}, A::LinearAlgebra.Adjoint{<:Any, ROCMatrix{T}}, B::ROCVector{T}) where T<:ROCBLASComplex = gemv_wrapper!(Y, 'C', A.parent, B)
+if VERSION < v"1.10.0-DEV.1365"
+    @inline LinearAlgebra.gemv!(
+        Y::ROCVector, tA::AbstractChar, A::StridedROCMatrix,
+        B::StridedROCVector, a::Number, b::Number,
+    ) = LinearAlgebra.generic_matvecmul!(Y, tA, A, B, MulAddMul(a, b))
+
+    # disambiguation with LinearAlgebra.jl
+    @inline LinearAlgebra.gemv!(
+        Y::ROCVector{T}, tA::AbstractChar, A::StridedROCMatrix{T},
+        B::StridedROCVector{T}, a::Number, b::Number,
+    ) where T <: ROCBLASFloat =
+        LinearAlgebra.generic_matvecmul!(Y, tA, A, B, MulAddMul(a, b))
+end
 
-############
 #
 # BLAS 3
 #
-############
 
 ########
 # GEMM
diff --git a/src/blas/rocBLAS.jl b/src/blas/rocBLAS.jl
@@ -1,12 +1,12 @@
 module rocBLAS
 
 using ..AMDGPU
-import AMDGPU: librocblas, AnyROCArray
+import AMDGPU: librocblas, AnyROCArray, StridedROCVector, StridedROCMatrix
 import AMDGPU: HandleCache, HIP, library_state
 import .HIP: HIPContext, HIPStream, hipStream_t, hipEvent_t
 
 using LinearAlgebra
-using LinearAlgebra: AdjOrTrans
+using LinearAlgebra: AdjOrTrans, MulAddMul
 if VERSION ≥ v"1.10-"
     using LinearAlgebra: wrap, UpperOrLowerTriangular
 end
diff --git a/src/blas/util.jl b/src/blas/util.jl
@@ -3,16 +3,6 @@ const ROCBLASComplex = Union{ComplexF32, ComplexF64}
 const ROCBLASFloat = Union{ROCBLASReal, ROCBLASComplex}
 const ROCBLASFloatWithHalf = Union{Float16, ROCBLASFloat}
 
-# Utility functions
-
-# convert Char {N,T,C} to rocblas_operation
-function rocblasop(trans::Char)
-    trans == 'N' && return rocblas_operation_none
-    trans == 'T' && return rocblas_operation_transpose
-    trans == 'C' && return rocblas_operation_conjugate_transpose
-    throw(ArgumentError("unknown rocblas operation $trans"))
-end
-
 function Base.convert(::Type{rocblas_operation}, trans::Char)
     if trans == 'N'
         return rocblas_operation_none
@@ -25,13 +15,6 @@ function Base.convert(::Type{rocblas_operation}, trans::Char)
     end
 end
 
-# convert Char {U,L} to rocblas_fill
-function rocblasfill(uplo::Char)
-    uplo == 'U' && return rocblas_fill_upper
-    uplo == 'L' && return rocblas_fill_lower
-    throw(ArgumentError("unknown rocblas fill mode $uplo"))
-end
-
 function Base.convert(::Type{rocblas_fill}, uplo::Char)
     if uplo == 'U'
         return rocblas_fill_upper
@@ -42,13 +25,6 @@ function Base.convert(::Type{rocblas_fill}, uplo::Char)
     end
 end
 
-# convert Char {U,N} to rocblas_diagonal
-function rocblasdiag(diag::Char)
-    diag == 'U' && return rocblas_diagonal_unit
-    diag == 'N' && return rocblas_diagonal_non_unit
-    throw(ArgumentError("unknown rocblas diag mode $diag"))
-end
-
 function Base.convert(::Type{rocblas_diagonal}, diag::Char)
     if diag == 'U'
         return rocblas_diagonal_unit
@@ -59,13 +35,6 @@ function Base.convert(::Type{rocblas_diagonal}, diag::Char)
     end
 end
 
-# convert Char {L,R} to rocblas_side
-function rocblasside(side::Char)
-    side == 'L' && return rocblas_side_left
-    side == 'R' && return rocblas_side_right
-    throw(ArgumentError("unknown rocblas side mode $side"))
-end
-
 function Base.convert(::Type{rocblas_side}, side::Char)
     if side == 'L'
         return rocblas_side_left
diff --git a/src/dnn/MIOpen.jl b/src/dnn/MIOpen.jl
@@ -8,7 +8,7 @@ import AMDGPU: libMIOpen_path
 import AMDGPU.Runtime.Mem # TODO remove?
 import .HIP: hipStream_t
 
-include("low_level.jl")
+include("libMIOpen.jl")
 
 const STATUS_DESCRIPTORS = Dict(
     miopenStatusSuccess => "Success",
diff --git a/src/dnn/libMIOpen.jl b/src/dnn/libMIOpen.jl
diff --git a/src/hsa/HSA.jl b/src/hsa/HSA.jl
@@ -1,6 +1,6 @@
 module HSA
 
-include("LibHSARuntime.jl")
+include("libHSA.jl")
 
 # Forward prefixed names
 hsa_names = map(string, names(LibHSARuntime))
diff --git a/src/hsa/libHSA.jl b/src/hsa/libHSA.jl
diff --git a/test/core_tests.jl b/test/core_tests.jl
@@ -86,10 +86,8 @@ include("rocarray/base.jl")
 include("rocarray/broadcast.jl")
 
 const IS_NAVI3 = AMDGPU.device().gcn_arch in ("gfx1100", "gfx1101", "gfx1102", "gfx1103")
-
-# TODO rework, hangs on Navi 3
 if !IS_NAVI3
-    include("tls.jl")
+    include("tls.jl") # TODO hangs on Navi 3
 end
 
 end
diff --git a/test/hip_core_tests.jl b/test/hip_core_tests.jl
@@ -1,4 +1,4 @@
-@testitem "hip - core" begin
+@testitem "hip - core" setup=[TSCore] begin
 
 using Test
 using LinearAlgebra
@@ -32,14 +32,4 @@ if length(AMDGPU.devices()) > 1
     end
 end
 
-if AMDGPU.functional(:rocblas)
-    include("rocarray/blas.jl")
-end
-if AMDGPU.functional(:MIOpen)
-    include("dnn/miopen.jl")
-end
-if AMDGPU.functional(:rocrand)
-    include("rocarray/random.jl")
-end
-
 end
diff --git a/test/hip_miopen_tests.jl b/test/hip_miopen_tests.jl
@@ -0,0 +1,19 @@
+@testitem "hip - MIOpen" setup=[TSCore] begin
+
+using Test
+using LinearAlgebra
+using Random
+
+using AMDGPU: HIP, Runtime, Device, Mem
+import AMDGPU: @allowscalar
+
+Random.seed!(1)
+AMDGPU.allowscalar(false)
+
+if AMDGPU.functional(:MIOpen)
+    include("dnn/miopen.jl")
+else
+    @test_skip "MIOpen"
+end
+
+end
diff --git a/test/hip_rocblas_tests.jl b/test/hip_rocblas_tests.jl
@@ -0,0 +1,19 @@
+@testitem "hip - rocBLAS" setup=[TSCore] begin
+
+using Test
+using LinearAlgebra
+using Random
+
+using AMDGPU: HIP, Runtime, Device, Mem
+import AMDGPU: @allowscalar
+
+Random.seed!(1)
+AMDGPU.allowscalar(false)
+
+if AMDGPU.functional(:rocblas)
+    include("rocarray/blas.jl")
+else
+    @test_skip "rocBLAS"
+end
+
+end
diff --git a/test/hip_rocfft_tests.jl b/test/hip_rocfft_tests.jl
@@ -10,11 +10,10 @@ import AMDGPU: @allowscalar
 Random.seed!(1)
 AMDGPU.allowscalar(false)
 
-# TODO rocFFT needs an update to work with ROCm 6.0+.
-if HIP.runtime_version() ≥ v"6-"
-    @test_skip "rocFFT"
-else
+if AMDGPU.functional(:rocfft)
     include("rocarray/fft.jl")
+else
+    @test_skip "rocFFT"
 end
 
 end
diff --git a/test/hip_rocrand_tests.jl b/test/hip_rocrand_tests.jl
@@ -0,0 +1,19 @@
+@testitem "hip - rocRAND" setup=[TSCore] begin
+
+using Test
+using LinearAlgebra
+using Random
+
+using AMDGPU: HIP, Runtime, Device, Mem
+import AMDGPU: @allowscalar
+
+Random.seed!(1)
+AMDGPU.allowscalar(false)
+
+if AMDGPU.functional(:rocrand)
+    include("rocarray/random.jl")
+else
+    @test_skip "rocRAND"
+end
+
+end
diff --git a/test/hip_rocsparse_tests.jl b/test/hip_rocsparse_tests.jl
@@ -10,13 +10,10 @@ import AMDGPU: @allowscalar
 Random.seed!(1)
 AMDGPU.allowscalar(false)
 
-@test_skip "rocSPARSE"
-
-# TODO hangs Navi 3 (and thus CI)
-# if AMDGPU.functional(:rocsparse)
-#     include("rocsparse/rocsparse.jl")
-# else
-#     @test_skip "rocSPARSE"
-# end
+if AMDGPU.functional(:rocsparse)
+    include("rocsparse/rocsparse.jl")
+else
+    @test_skip "rocSPARSE"
+end
 
 end
diff --git a/test/rocarray/blas.jl b/test/rocarray/blas.jl
diff --git a/test/runtests.jl b/test/runtests.jl
diff --git a/test/tls.jl b/test/tls.jl
