More tests and bugfixes for CUSOLVER (#2707)

* More tests and bugfixes for CUSOLVER

* Add a test for sytrs

Author: kshyatt

lib/cusolver/dense_generic.jl

@@ -122,7 +122,7 @@ function Xgeqrf!(A::StridedCuMatrix{T}) where {T <: BlasFloat}
 end
 
 # Xsytrs
-function sytrs!(uplo::Char, A::StridedCuMatrix{T}, p::CuVector{Int64}, B::StridedCuMatrix{T}) where {T <: BlasFloat}
+function sytrs!(uplo::Char, A::StridedCuMatrix{T}, p::CuVector{Int64}, B::StridedCuVecOrMat{T}) where {T <: BlasFloat}
     chkuplo(uplo)
     n = checksquare(A)
     nrhs = size(B, 2)
@@ -149,7 +149,7 @@ function sytrs!(uplo::Char, A::StridedCuMatrix{T}, p::CuVector{Int64}, B::Stride
     B
 end
 
-function sytrs!(uplo::Char, A::StridedCuMatrix{T}, B::StridedCuMatrix{T}) where {T <: BlasFloat}
+function sytrs!(uplo::Char, A::StridedCuMatrix{T}, B::StridedCuVecOrMat{T}) where {T <: BlasFloat}
     chkuplo(uplo)
     n = checksquare(A)
     nrhs = size(B, 2)

lib/cusolver/linalg.jl

@@ -66,7 +66,7 @@ function Base.:\(_A::CuMatOrAdj, _B::CuOrAdj)
 end
 
 function Base.:\(_A::Symmetric{<:Any,<:CuMatOrAdj}, _B::CuOrAdj)
-    uplo = A.uplo
+    uplo = _A.uplo
     A, B = copy_cublasfloat(_A.data, _B)
 
     # LDLᴴ decomposition with partial pivoting
@@ -371,9 +371,9 @@ for (triangle, uplo, diag) in ((:LowerTriangular, 'L', 'N'),
     @eval begin
         function LinearAlgebra.inv(A::$triangle{T,<:StridedCuMatrix{T}}) where T <: BlasFloat
             n = checksquare(A)
-            B = copy(A.data)
-            trtri!(uplo, diag, B)
-            return B
+            B = copy(parent(A))
+            trtri!($uplo, $diag, B)
+            return $triangle(B)
         end
     end
 end

test/libraries/cusolver/dense.jl

@@ -89,9 +89,9 @@ k = 1
             A = rand(elty,n,n)
             A = uplo == 'L' ? tril(A) : triu(A)
             A = diag == 'N' ? A : A - Diagonal(A) + I
-            dA = triangle(CuArray(A))
+            dA = triangle(view(CuArray(A), 1:2:n, 1:2:n)) # without this view, we are hitting the CUBLAS method!
             dA⁻¹ = inv(dA)
-            dI = dA.data * dA⁻¹
+            dI = CuArray(dA) * CuArray(dA⁻¹)
             @test Array(dI) ≈ I
         end
     end
@@ -265,6 +265,8 @@ k = 1
         A             += A'
         d_A            = CuArray(A)
         Eig            = eigen(LinearAlgebra.Hermitian(A))
+        d_eig          = eigen(d_A)
+        @test Eig.values ≈ collect(d_eig.values)
         d_eig          = eigen(LinearAlgebra.Hermitian(d_A))
         @test Eig.values ≈ collect(d_eig.values)
         h_V            = collect(d_eig.vectors)
@@ -533,6 +535,7 @@ k = 1
         B              = rand(elty, n)
         d_B            = CuArray(B)
         @test collect(d_M \ d_B) ≈ M \ B
+        @test_throws DimensionMismatch("arguments must have the same number of rows") d_M \ CUDA.ones(elty, n+1)
         A              = rand(elty, m, n)  # A is a matrix and B,C is a vector
         d_A            = CuArray(A)
         M              = qr(A)
@@ -782,7 +785,9 @@ end
         Bf = cublasfloat.(B)
         bf = cublasfloat.(b)
         @test Array(d_A \ d_B) ≈ (Af \ Bf)
+        @test Array(Symmetric(d_A) \ d_B) ≈ (Af \ Bf)
         @test Array(d_A \ d_b) ≈ (Af \ bf)
+        @test Array(Symmetric(d_A) \ d_b) ≈ (Af \ bf)
         @inferred d_A \ d_B
         @inferred d_A \ d_b
     end

test/libraries/cusolver/dense_generic.jl

@@ -96,21 +96,27 @@ p = 5
             @testset "pivoting = $pivoting" for pivoting in (false, true)
                 A = rand(elty,n,n)
                 B = rand(elty,n,p)
+                C = rand(elty,n)
                 A = A + transpose(A)
                 d_A = CuMatrix(A)
                 d_B = CuMatrix(B)
+                d_C = CuVector(C)
                 !pivoting && (CUSOLVER.version() < v"11.7.2") && continue
                 if pivoting
                     d_A, d_ipiv, _ = CUSOLVER.sytrf!(uplo, d_A; pivoting)
                     d_ipiv = CuVector{Int64}(d_ipiv)
                     CUSOLVER.sytrs!(uplo, d_A, d_ipiv, d_B)
+                    CUSOLVER.sytrs!(uplo, d_A, d_ipiv, d_C)
                 else
                     d_A, _ = CUSOLVER.sytrf!(uplo, d_A; pivoting)
                     CUSOLVER.sytrs!(uplo, d_A, d_B)
+                    CUSOLVER.sytrs!(uplo, d_A, d_C)
                 end
                 A, ipiv, _ = LAPACK.sytrf!(uplo, A)
                 LAPACK.sytrs!(uplo, A, ipiv, B)
+                LAPACK.sytrs!(uplo, A, ipiv, C)
                 @test B ≈ collect(d_B)
+                @test C ≈ collect(d_C)
             end
         end
     end