XsyevBatched! interface accepting 3D StridedCuArray (#2951)

Signed-off-by: Steven Hahn <hahnse@ornl.gov>

Author: quantumsteve

lib/cusolver/dense_generic.jl

@@ -505,8 +505,53 @@ function Xgeev!(jobvl::Char, jobvr::Char, A::StridedCuMatrix{T}) where {T <: Bla
 end
 
 # XsyevBatched
+function XsyevBatched!(jobz::Char, uplo::Char, A::StridedCuArray{T, 3}) where {T <: BlasFloat}
+    minimum_version = v"11.7.1"
+    CUSOLVER.version() < minimum_version && throw(ErrorException("This operation requires cuSOLVER
+        $(minimum_version) or later. Current cuSOLVER version: $(CUSOLVER.version())."))
+    chkuplo(uplo)
+    n = checksquare(A)
+    batch_size = size(A, 3)
+    R = real(T)
+    lda = max(1, stride(A, 2))
+    W = CuMatrix{R}(undef, n, batch_size)
+    params = CuSolverParameters()
+    dh = dense_handle()
+    resize!(dh.info, batch_size)
+
+    function bufferSize()
+        out_cpu = Ref{Csize_t}(0)
+        out_gpu = Ref{Csize_t}(0)
+        cusolverDnXsyevBatched_bufferSize(
+            dh, params, jobz, uplo, n,
+            T, A, lda, R, W, T, out_gpu, out_cpu, batch_size
+        )
+        return out_gpu[], out_cpu[]
+    end
+    with_workspaces(dh.workspace_gpu, dh.workspace_cpu, bufferSize()...) do buffer_gpu, buffer_cpu
+        cusolverDnXsyevBatched(
+            dh, params, jobz, uplo, n, T, A,
+            lda, R, W, T, buffer_gpu, sizeof(buffer_gpu),
+            buffer_cpu, sizeof(buffer_cpu), dh.info, batch_size
+        )
+    end
+
+    info = @allowscalar collect(dh.info)
+    for i in 1:batch_size
+        chkargsok(info[i] |> BlasInt)
+    end
+
+    if jobz == 'N'
+        return W
+    elseif jobz == 'V'
+        return W, A
+    end
+end
+
 function XsyevBatched!(jobz::Char, uplo::Char, A::StridedCuMatrix{T}) where {T <: BlasFloat}
-    CUSOLVER.version() < v"11.7.1" && throw(ErrorException("This operation is not supported by the current CUDA version."))
+    minimum_version = v"11.7.1"
+    CUSOLVER.version() < minimum_version && throw(ErrorException("This operation requires cuSOLVER
+        $(minimum_version) or later. Current cuSOLVER version: $(CUSOLVER.version())."))
     chkuplo(uplo)
     n, num_matrices = size(A)
     batch_size = num_matrices ÷ n

test/libraries/cusolver/dense_generic.jl

@@ -33,6 +33,36 @@ p = 5
         end
 
         @testset "syevBatched!" begin
+            batch_size = 5
+            for uplo in ('L', 'U')
+                (CUSOLVER.version() < v"11.7.2") && (uplo == 'L') && (elty == ComplexF32) && continue
+
+                A = rand(elty, n, n, batch_size)
+                B = rand(elty, n, n, batch_size)
+                for i in 1:batch_size
+                    S = rand(elty, n, n)
+                    S = S * S' + I
+                    B[:, :, i] .= S
+                    S = uplo == 'L' ? tril(S) : triu(S)
+                    A[:, :, i] .= S
+                end
+                d_A = CuArray(A)
+                d_W, d_V = CUSOLVER.XsyevBatched!('V', uplo, d_A)
+                W = collect(d_W)
+                V = collect(d_V)
+                for i in 1:batch_size
+                    Bᵢ = B[:, :, i]
+                    Wᵢ = Diagonal(W[:, i])
+                    Vᵢ = V[:, :, i]
+                    @test Bᵢ * Vᵢ ≈ Vᵢ * Diagonal(Wᵢ)
+                end
+
+                d_A = CuArray(A)
+                d_W = CUSOLVER.XsyevBatched!('N', uplo, d_A)
+            end
+        end
+
+        @testset "syevBatched! updated" begin
             batch_size = 5
             for uplo in ('L', 'U')
                 (CUSOLVER.version() < v"11.7.2") && (uplo == 'L') && (elty == ComplexF32) && continue