More tests for CUBLAS and a bugfix (#2659)

Author: kshyatt

lib/cublas/wrappers.jl

@@ -70,8 +70,6 @@ function juliaStorageType(T::Type{<:Complex}, ct::cublasComputeType_t)
         return Complex{Float32}
     elseif ct == CUBLAS_COMPUTE_64F || ct == CUBLAS_COMPUTE_64F_PEDANTIC
         return Complex{Float64}
-    elseif ct == CUBLAS_COMPUTE_32I || ct == CUBLAS_COMPUTE_32I_PEDANTIC
-        return Complex{Int32}
     else
         throw(ArgumentError("Julia type equivalent for compute type $ct does not exist!"))
     end
@@ -1174,14 +1172,10 @@ function gemmExComputeType(TA, TB, TC, m, k, n)
 
     # gemmEx requires sm_50 or higher
     cap = capability(device())
-    if cap < v"5"
-        return nothing
-    end
+    cap < v"5" && return nothing
 
     # source: CUBLAS Features and Technical Specifications
-    if Float16 in sig && cap < v"5.3"
-        return nothing
-    end
+    Float16 in sig && cap < v"5.3" && return nothing
 
     math_mode = CUDA.math_mode()
     reduced_precision = CUDA.math_precision()
@@ -1192,15 +1186,10 @@ function gemmExComputeType(TA, TB, TC, m, k, n)
     end
 
     if sig === (Int8, Int32)
-        # starting with CUDA 11.2, this is unsupported (NVIDIA bug #3221266)
-        # TODO: might be fixed in a later version?
-        version() >= v"11.3.1" && return nothing
-
         # Int32=Int8*Int8 requires m,n,k to be multiples of 4
         # https://forums.developer.nvidia.com/t/cublasgemmex-cant-use-cuda-r-8i-compute-type-on-gtx1080/58100/2
-        if m%4 == 0 && n%4 == 0 && k%4 == 0
-            return math_mode==CUDA.PEDANTIC_MATH ? CUBLAS_COMPUTE_32I_PEDANTIC : CUBLAS_COMPUTE_32I
-        end
+        all_mod_4 = (m%4 == 0 && n%4 == 0 && k%4 == 0)
+        all_mod_4 && return math_mode==CUDA.PEDANTIC_MATH ? CUBLAS_COMPUTE_32I_PEDANTIC : CUBLAS_COMPUTE_32I
     end
 
     if math_mode == CUDA.FAST_MATH
@@ -1231,13 +1220,8 @@ function gemmExComputeType(TA, TB, TC, m, k, n)
        sig === (Complex{Float64}, Complex{Float64})
         return math_mode==CUDA.PEDANTIC_MATH ? CUBLAS_COMPUTE_64F_PEDANTIC : CUBLAS_COMPUTE_64F
     end
-
-    # BFloat16 support was added in CUDA 11
-    if version() >= v"11"
-        if sig === (BFloat16, BFloat16) ||
-           sig === (BFloat16, Float32)
-            return math_mode==CUDA.PEDANTIC_MATH ? CUBLAS_COMPUTE_32F_PEDANTIC : CUBLAS_COMPUTE_32F
-        end
+    if sig === (BFloat16, BFloat16) || sig === (BFloat16, Float32)
+        return math_mode==CUDA.PEDANTIC_MATH ? CUBLAS_COMPUTE_32F_PEDANTIC : CUBLAS_COMPUTE_32F
     end
 
     return nothing
@@ -1263,20 +1247,11 @@ function gemmEx!(transA::Char, transB::Char,
     isnothing(computeType) &&
         throw(ArgumentError("gemmEx does not support $(eltype(C))=$(eltype(A))*$(eltype(B))"))
     computeT = juliaStorageType(eltype(C), computeType)
-    if version() >= v"11.0"
-        # with CUDA 11, the compute type encodes the math mode.
-        cublasGemmEx(
-            handle(), transA, transB, m, n, k, CuRef{computeT}(alpha), A, eltype(A), lda, B,
-            eltype(B), ldb, CuRef{computeT}(beta), C, eltype(C), ldc, computeType, algo
-        )
-    else
-        # before CUDA 11, it was a plain cudaDataType.
-        computeType = convert(cudaDataType, computeT)
-        cublasGemmEx_old(
-            handle(), transA, transB, m, n, k, CuRef{computeT}(alpha), A, eltype(A), lda, B,
-            eltype(B), ldb, CuRef{computeT}(beta), C, eltype(C), ldc, computeType, algo
-        )
-    end
+    
+    cublasGemmEx(
+        handle(), transA, transB, m, n, k, CuRef{computeT}(alpha), A, eltype(A), lda, B,
+        eltype(B), ldb, CuRef{computeT}(beta), C, eltype(C), ldc, computeType, algo
+    )
     C
 end
 
@@ -1311,15 +1286,11 @@ function gemmBatchedEx!(transA::Char, transB::Char,
     Aptrs = unsafe_batch(A)
     Bptrs = unsafe_batch(B)
     Cptrs = unsafe_batch(C)
-    if version() >= v"11.0"
-        # with CUDA 11, the compute type encodes the math mode.
-        cublasGemmBatchedEx(
-            handle(), transA, transB, m, n, k, CuRef{computeT}(alpha), Aptrs, eltype(A[1]), lda, Bptrs,
-            eltype(B[1]), ldb, CuRef{computeT}(beta), Cptrs, eltype(C[1]), ldc, length(A), computeType, algo
-        )
-    else
-        error("Not implemented for CUDA 11 and below.")
-    end
+
+    cublasGemmBatchedEx(
+        handle(), transA, transB, m, n, k, CuRef{computeT}(alpha), Aptrs, eltype(A[1]), lda, Bptrs,
+        eltype(B[1]), ldb, CuRef{computeT}(beta), Cptrs, eltype(C[1]), ldc, length(A), computeType, algo
+    )
     unsafe_free!(Cptrs)
     unsafe_free!(Bptrs)
     unsafe_free!(Aptrs)
@@ -1357,15 +1328,10 @@ function gemmStridedBatchedEx!(
     isnothing(computeType) &&
     throw(ArgumentError("gemmEx does not support $(eltype(C))=$(eltype(A))*$(eltype(B))"))
     computeT = juliaStorageType(eltype(C), computeType)
-    if version() >= v"11.0"
-        # with CUDA 11, the compute type encodes the math mode.
-        cublasGemmStridedBatchedEx(
-            handle(), transA, transB, m, n, k, CuRef{computeT}(alpha), A, eltype(A), lda, strideA,
-            B, eltype(B), ldb, strideB, CuRef{computeT}(beta), C, eltype(C), ldc, strideC,
-                                   batchCount, computeType, algo)
-    else
-        error("Not implemented for CUDA 11 and below.")
-    end
+    cublasGemmStridedBatchedEx(
+        handle(), transA, transB, m, n, k, CuRef{computeT}(alpha), A, eltype(A), lda, strideA,
+        B, eltype(B), ldb, strideB, CuRef{computeT}(beta), C, eltype(C), ldc, strideC,
+        batchCount, computeType, algo)
     C
 end
 
@@ -1382,6 +1348,8 @@ end
     #ptrs = [pointer(strided, (i-1)*batch_stride + 1) for i in 1:batch_size]
     # fill the array on the GPU to avoid synchronous copies and support larger batch sizes
     ptrs = CuArray{CuPtr{T}}(undef, batch_size)
+    # device-side code
+    ## COV_EXCL_START
     function compute_pointers()
         i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x
         grid_stride = gridDim().x * blockDim().x
@@ -1392,6 +1360,7 @@ end
         end
         return
     end
+    ## COV_EXCL_STOP
     kernel = @cuda launch = false compute_pointers()
     config = launch_configuration(kernel.fun)
     threads = min(config.threads, batch_size)
@@ -2337,7 +2306,7 @@ for (fname, elty) in ((:cublasDgetriBatched, :Float64),
             ldc = max(1, stride(C[1], 2))
             Aptrs = unsafe_batch(A)
             Cptrs = unsafe_batch(C)
-            info = CuArrays.zeros(Cint, length(A))
+            info = CUDA.zeros(Cint, length(A))
             $fname(handle(), n, Aptrs, lda, pivotArray, Cptrs, ldc, info, length(A))
             unsafe_free!(Cptrs)
             unsafe_free!(Aptrs)

test/libraries/cublas/extensions.jl

@@ -54,6 +54,14 @@ k = 13
                 end
                 @test inv(P)*dL*dU ≈ inv(C.P) * C.L * C.U
             end
+            # generate bad matrices
+            A_bad = vcat([rand(elty,m,m) for i in 1:9], [rand(elty, m, m+1)])
+            # move to device
+            d_A_bad = CuArray{elty, 2}[]
+            for i in 1:length(A)
+                push!(d_A_bad,CuArray(A_bad[i]))
+            end
+            @test_throws DimensionMismatch CUBLAS.getrf_batched!(d_A_bad, true)
         end
 
         @testset "getrf_batched" begin
@@ -128,6 +136,11 @@ k = 13
                 end
                 @test inv(P)*dL*dU ≈ inv(C.P) * C.L * C.U
             end
+            # generate bad strided matrix
+            A = rand(elty,m,m+1,10)
+            # move to device
+            d_A = CuArray(A)
+            @test_throws DimensionMismatch CUBLAS.getrf_strided_batched!(d_A, true)
         end
 
         @testset "getrf_strided_batched" begin
@@ -168,24 +181,41 @@ k = 13
         for (opchar,opfun) in (('N',identity), ('T',transpose), ('C',adjoint))
 
             @testset "getrs_batched!" begin
-                A                   = [rand(elty,n,n) for _ in 1:k];
-                d_A                 = [CuArray(a) for a in A];
-                d_A2                = deepcopy(d_A);
-                d_pivot, info, d_LU = CUDA.CUBLAS.getrf_batched!(d_A, true);
+                A                   = [rand(elty,n,n) for _ in 1:k]
+                d_A                 = [CuArray(a) for a in A]
+                d_A2                = deepcopy(d_A)
+                d_pivot, info, d_LU = CUDA.CUBLAS.getrf_batched!(d_A, true)
                 @test d_LU == d_A
-                d_pivot2            = similar(d_pivot);
-                info2               = similar(info);
-                CUDA.CUBLAS.getrf_batched!(d_A2, d_pivot2, info2);
+                d_pivot2            = similar(d_pivot)
+                info2               = similar(info)
+                CUDA.CUBLAS.getrf_batched!(d_A2, d_pivot2, info2)
                 @test isapprox(d_pivot, d_pivot2)
                 @test isapprox(info, info2)
-                B                   = [rand(elty,n,m) for _ in 1:k];
-                d_B                 = [CuArray(b) for b in B];
-                info2, d_Bhat       = CUDA.CUBLAS.getrs_batched!(opchar, d_LU, d_B, d_pivot);
+                B                   = [rand(elty,n,m) for _ in 1:k]
+                d_B                 = [CuArray(b) for b in B]
+                info2, d_Bhat       = CUDA.CUBLAS.getrs_batched!(opchar, d_LU, d_B, d_pivot)
                 @test d_Bhat == d_B
-                h_Bhat              = [collect(bh) for bh in d_Bhat];
+                h_Bhat              = [collect(bh) for bh in d_Bhat]
                 for i in 1:k
                     @test h_Bhat[i] ≈ opfun(A[i]) \ B[i]
                 end
+                
+                # generate bad matrices
+                A_bad = vcat([rand(elty,m,m) for i in 1:9], [rand(elty, m, m+1)])
+                # move to device
+                d_A_bad = CuArray{elty, 2}[]
+                for i in 1:length(A_bad)
+                    push!(d_A_bad,CuArray(A_bad[i]))
+                end
+                @test_throws DimensionMismatch CUBLAS.getrs_batched!(opchar, d_A_bad, d_B, d_pivot)
+                # generate bad matrices
+                A_bad = [rand(elty,m+1,m+1) for i in 1:10]
+                # move to device
+                d_A_bad = CuArray{elty, 2}[]
+                for i in 1:length(A_bad)
+                    push!(d_A_bad,CuArray(A_bad[i]))
+                end
+                @test_throws DimensionMismatch CUBLAS.getrs_batched!(opchar, d_A_bad, d_B, d_pivot)
             end
 
             @testset "getrs_batched" begin
@@ -210,24 +240,31 @@ k = 13
             end
 
             @testset "getrs_strided_batched!" begin
-                A                   = rand(elty,n,n,k);
-                d_A                 = CuArray(A);
-                d_A2                = copy(d_A);
-                d_pivot, info, d_LU = CUDA.CUBLAS.getrf_strided_batched!(d_A, true);
+                A                   = rand(elty,n,n,k)
+                d_A                 = CuArray(A)
+                d_A2                = copy(d_A)
+                d_pivot, info, d_LU = CUDA.CUBLAS.getrf_strided_batched!(d_A, true)
                 @test d_LU == d_A
-                d_pivot2            = similar(d_pivot);
-                info2               = similar(info);
-                CUDA.CUBLAS.getrf_strided_batched!(d_A2, d_pivot2, info2);
+                d_pivot2            = similar(d_pivot)
+                info2               = similar(info)
+                CUDA.CUBLAS.getrf_strided_batched!(d_A2, d_pivot2, info2)
                 @test isapprox(d_pivot, d_pivot2)
                 @test isapprox(info, info2)
-                B                   = rand(elty,n,m,k);
-                d_B                 = CuArray(B);
-                info2, d_Bhat       = CUDA.CUBLAS.getrs_strided_batched!(opchar, d_LU, d_B, d_pivot);
+                B                   = rand(elty,n,m,k)
+                d_B                 = CuArray(B)
+                info2, d_Bhat       = CUDA.CUBLAS.getrs_strided_batched!(opchar, d_LU, d_B, d_pivot)
                 @test d_Bhat == d_B
-                h_Bhat              = collect(d_Bhat);
+                h_Bhat              = collect(d_Bhat)
                 for i in 1:k
                     @test h_Bhat[:,:,i] ≈ opfun(A[:,:,i]) \ B[:,:,i]
                 end
+
+                A_bad               = rand(elty,n+1,n,k)
+                d_A_bad             = CuArray(A_bad)
+                @test_throws DimensionMismatch CUDA.CUBLAS.getrs_strided_batched!(opchar, d_A_bad, d_B, d_pivot)
+                A_bad               = rand(elty,n+1,n+1,k)
+                d_A_bad             = CuArray(A_bad)
+                @test_throws DimensionMismatch CUDA.CUBLAS.getrs_strided_batched!(opchar, d_A_bad, d_B, d_pivot)
             end
 
             @testset "getrs_strided_batched" begin
@@ -267,6 +304,12 @@ k = 13
                 @test h_info[Cs] == 0
                 @test B ≈ Array(d_B[:,:,Cs]) rtol=1e-3
             end
+
+            A_bad = rand(elty,m+1,m,10)
+            d_A_bad = CuArray(A_bad)
+            d_B = similar(d_A)
+            pivot, info = CUBLAS.getrf_strided_batched!(d_A, true)
+            @test_throws DimensionMismatch CUBLAS.getri_strided_batched!(d_A_bad, d_B, pivot)
         end
 
         @testset "getri_batched" begin
@@ -290,6 +333,32 @@ k = 13
                 @test h_info[Cs] == 0
                 @test C ≈ h_C rtol=1e-2
             end
+
+            d_A = CuArray{elty, 2}[]
+            for i in 1:length(A)
+                push!(d_A,CuArray(A[i]))
+            end
+            pivot, info = CUBLAS.getrf_batched!(d_A, true)
+            h_info = Array(info)
+            for Cs in 1:length(h_info)
+                @test h_info[Cs] == 0
+            end
+            d_C = CuMatrix{elty}[similar(d_A[1]) for i in 1:length(d_A)]
+            info = CUBLAS.getri_batched!(d_A, d_C, pivot)
+            h_info = Array(info)
+            for Cs in 1:length(d_C)
+                C   = inv(A[Cs])
+                h_C = Array(d_C[Cs])
+                @test h_info[Cs] == 0
+                @test C ≈ h_C rtol=1e-2
+            end
+
+            A_bad = [rand(elty,m+1,m) for i in 1:10]
+            d_A_bad = CuArray{elty, 2}[]
+            for i in 1:length(A)
+                push!(d_A_bad,CuArray(A_bad[i]))
+            end
+            @test_throws DimensionMismatch CUBLAS.getri_batched(d_A_bad, pivot)
         end
 
         @testset "matinv_batched" begin
@@ -308,6 +377,15 @@ k = 13
             end
             push!(d_A, CUDA.rand(elty, m, m+1))
             @test_throws DimensionMismatch CUBLAS.matinv_batched(d_A)
+
+            # matinv_batched only supports matrices smaller than 32x32
+            A = [rand(elty,64,64) for i in 1:10]
+            # move to device
+            d_A_too_big = CuArray{elty, 2}[]
+            for i in 1:length(A)
+                push!(d_A_too_big,CuArray(A[i]))
+            end
+            @test_throws ArgumentError("matinv requires all matrices be smaller than 32 x 32") CUBLAS.matinv_batched(d_A_too_big)
         end
 
         @testset "geqrf_batched!" begin
@@ -343,7 +421,7 @@ k = 13
             for i in 1:length(A)
                 push!(d_A,CuArray(A[i]))
             end
-            tau, d_B = CUBLAS.geqrf_batched!(d_A)
+            tau, d_B = CUBLAS.geqrf_batched(d_A)
             for Bs in 1:length(d_B)
                 C   = qr(A[Bs])
                 h_B = Array(d_B[Bs])
@@ -392,6 +470,18 @@ k = 13
             end
             # system is now not overdetermined
             @test_throws ArgumentError CUBLAS.gels_batched!('N',d_A, d_C)
+
+            # generate bad matrices
+            A = [rand(elty,n,k) for i in 1:10]
+            C = [rand(elty,n+1,k) for i in 1:10]
+            # move to device
+            d_A = CuArray{elty, 2}[]
+            d_C = CuArray{elty, 2}[]
+            for i in 1:length(A)
+                push!(d_A,CuArray(A[i]))
+                push!(d_C,CuArray(C[i]))
+            end
+            @test_throws DimensionMismatch CUBLAS.gels_batched!('N',d_A, d_C)
         end
 
         @testset "gels_batched" begin

test/libraries/cublas/level1.jl

@@ -166,6 +166,17 @@ k = 13
         @test testf(axpy!, rand(), rand(T, m), rand(T, m))
         @test testf(LinearAlgebra.axpby!, rand(), rand(T, m), rand(), rand(T, m))
 
+        @testset "scal!" begin
+            x = rand(T, m)
+            d_x = CuArray(x)
+            α = rand(Float32)
+            d_α = CuArray([α])
+            y = α * x
+            d_x = CUBLAS.scal!(m, d_α, d_x)
+            h_y = Array(d_x)
+            @test h_y ≈ y
+        end
+
         if T <: Complex
             @test testf(dot, rand(T, m), rand(T, m))
             x = rand(T, m)