[CUSOLVER] Support symmetric factorization without pivoting (#2640)

Author: amontoison

lib/cusolver/dense.jl

@@ -206,10 +206,34 @@ for (bname, fname,elty) in ((:cusolverDnSsytrf_bufferSize, :cusolverDnSsytrf, :F
             A, ipiv, info
         end
 
-        function sytrf!(uplo::Char, A::StridedCuMatrix{$elty})
+        function sytrf!(uplo::Char, A::StridedCuMatrix{$elty}; pivoting::Bool=true)
             n = checksquare(A)
-            ipiv = CuArray{Cint}(undef, n)
-            sytrf!(uplo, A, ipiv)
+            if pivoting
+                ipiv = CuArray{Cint}(undef, n)
+                return sytrf!(uplo, A, ipiv)
+            else
+                CUSOLVER.version() < v"11.7.2" && throw(ErrorException("This operation is not supported by the current CUDA version."))
+                chkuplo(uplo)
+                n = checksquare(A)
+                lda = max(1, stride(A, 2))
+                dh = dense_handle()
+
+                function bufferSize()
+                    out = Ref{Cint}(0)
+                    $bname(dh, n, A, lda, out)
+                    return out[] * sizeof($elty)
+                end
+
+                with_workspace(dh.workspace_gpu, bufferSize) do buffer
+                    $fname(dh, uplo, n, A, lda, CU_NULL,
+                           buffer, sizeof(buffer) ÷ sizeof($elty), dh.info)
+                end
+
+                info = @allowscalar dh.info[1]
+                chkargsok(info |> BlasInt)
+
+                return A, info
+            end
         end
     end
 end

lib/cusolver/dense_generic.jl

@@ -149,6 +149,33 @@ 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}
+    chkuplo(uplo)
+    n = checksquare(A)
+    nrhs = size(B, 2)
+    lda = max(1, stride(A, 2))
+    ldb = max(1, stride(B, 2))
+    dh = dense_handle()
+
+    function bufferSize()
+        out_cpu = Ref{Csize_t}(0)
+        out_gpu = Ref{Csize_t}(0)
+        cusolverDnXsytrs_bufferSize(dh, uplo, n, nrhs, T, A,
+                                    lda, CU_NULL, T, B, ldb, out_gpu, out_cpu)
+        out_gpu[], out_cpu[]
+    end
+    with_workspaces(dh.workspace_gpu, dh.workspace_cpu,
+                    bufferSize()...) do buffer_gpu, buffer_cpu
+        cusolverDnXsytrs(dh, uplo, n, nrhs, T, A, lda, CU_NULL,
+                         T, B, ldb, buffer_gpu, sizeof(buffer_gpu),
+                         buffer_cpu, sizeof(buffer_cpu), dh.info)
+    end
+
+    flag = @allowscalar dh.info[1]
+    chkargsok(flag |> BlasInt)
+    B
+end
+
 # Xtrtri
 function trtri!(uplo::Char, diag::Char, A::StridedCuMatrix{T}) where {T <: BlasFloat}
     chkuplo(uplo)

test/libraries/cusolver/dense_generic.jl

@@ -92,18 +92,26 @@ p = 5
     end
 
     @testset "sytrs!" begin
-        for uplo in ('L', 'U')
-            A = rand(elty,n,n)
-            B = rand(elty,n,p)
-            A = A + transpose(A)
-            d_A = CuMatrix(A)
-            d_B = CuMatrix(B)
-            d_A, d_ipiv, _ = CUSOLVER.sytrf!(uplo, d_A)
-            d_ipiv = CuVector{Int64}(d_ipiv)
-            A, ipiv, _ = LAPACK.sytrf!(uplo, A)
-            CUSOLVER.sytrs!(uplo, d_A, d_ipiv, d_B)
-            LAPACK.sytrs!(uplo, A, ipiv, B)
-            @test B ≈ collect(d_B)
+        @testset "uplo = $uplo" for uplo in ('L', 'U')
+            @testset "pivoting = $pivoting" for pivoting in (false, true)
+                A = rand(elty,n,n)
+                B = rand(elty,n,p)
+                A = A + transpose(A)
+                d_A = CuMatrix(A)
+                d_B = CuMatrix(B)
+                !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)
+                else
+                    d_A, _ = CUSOLVER.sytrf!(uplo, d_A; pivoting)
+                    CUSOLVER.sytrs!(uplo, d_A, d_B)
+                end
+                A, ipiv, _ = LAPACK.sytrf!(uplo, A)
+                LAPACK.sytrs!(uplo, A, ipiv, B)
+                @test B ≈ collect(d_B)
+            end
         end
     end