Interface the batch routines

exanauts · Dec 11, 2024 · 1132a51 · 1132a51
1 parent 5efca42
commit 1132a51
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 1 deletion.
diff --git a/src/helpers.jl b/src/helpers.jl
@@ -7,6 +7,9 @@ export CudssMatrix, CudssData, CudssConfig
     matrix = CudssMatrix(v::CuVector{T})
     matrix = CudssMatrix(A::CuMatrix{T})
     matrix = CudssMatrix(A::CuSparseMatrixCSR{T,Cint}, struture::String, view::Char; index::Char='O')
+    matrix = CudssMatrix(v::Vector{CuVector{T}})
+    matrix = CudssMatrix(A::Vector{CuMatrix{T}})
+    matrix = CudssMatrix(A::Vector{CuSparseMatrixCSR{T,Cint}}, struture::String, view::Char; index::Char='O')
 
 The type `T` can be `Float32`, `Float64`, `ComplexF32` or `ComplexF64`.
 
@@ -79,12 +82,58 @@ mutable struct CudssMatrix{T}
         m,n = size(A)
         matrix_ref = Ref{cudssMatrix_t}()
         cudssMatrixCreateCsr(matrix_ref, m, n, nnz(A), A.rowPtr, CU_NULL,
-                             A.colVal, A.nzVal, eltype(A.rowPtr), T, structure,
+                             A.colVal, A.nzVal, Cint, T, structure,
                              view, index)
         obj = new{T}(T, matrix_ref[])
         finalizer(cudssMatrixDestroy, obj)
         obj
     end
+
+    function CudssMatrix(v::Vector{CuVector{T}}) where T <: BlasFloat
+        matrix_ref = Ref{cudssMatrix_t}()
+        nbatch = length(v)
+        nrows = [length(vᵢ) for vᵢ in v]
+        ncols = [1 for i = 1:nbatch]
+        ld = nrows
+        cudssMatrixCreateBatchDn(matrix, nbatch, nrows, ncols, ld, v, T, 'C')
+        obj = new{T}(T, matrix_ref[])
+        finalizer(cudssMatrixDestroy, obj)
+        obj
+    end
+
+    function CudssMatrix(A::Vector{CuMatrix{T}}; transposed::Bool=false) where T <: BlasFloat
+        matrix_ref = Ref{cudssMatrix_t}()
+        nbatch = length(A)
+        nrows = [size(Aᵢ,1) for Aᵢ in A]
+        ncols = [size(Aᵢ,2) for Aᵢ in A]
+        ld = nrows
+        if transposed
+            cudssMatrixCreateBatchDn(matrix_ref, nbatch, ncols, nrows, ld, A, T, 'R')
+        else
+            cudssMatrixCreateBatchDn(matrix_ref, nbatch, nrows, ncols, ld, A, T, 'C')
+        end
+        obj = new{T}(T, matrix_ref[])
+        finalizer(cudssMatrixDestroy, obj)
+        obj
+    end
+
+    function CudssMatrix(A::Vector{CuSparseMatrixCSR{T,Cint}}, structure::String, view::Char; index::Char='O') where T <: BlasFloat
+        matrix_ref = Ref{cudssMatrix_t}()
+        nbatch = length(A)
+        nrows = [size(Aᵢ,1) for Aᵢ in A]
+        ncols = [size(Aᵢ,2) for Aᵢ in A]
+        nnzA = [nnz(Aᵢ) for Aᵢ in A]
+        rowsPtrs = [A.rowPtr for Aᵢ in A]
+        colVals = [A.colVal for Aᵢ in A]
+        nzVals = [A.nzVal for Aᵢ in A]
+        PTR_CU_NULL = Ptr{CuPtr{Cvoid}}()
+        cudssMatrixCreateBatchCsr(matrix_ref, nbatch, nrows, ncols, nnzA, rowPtrs,
+                                  PTR_CU_NULL, colVals, nzVals, Cint, T, structure,
+                                  view, index)
+        obj = new{T}(T, matrix_ref[])
+        finalizer(cudssMatrixDestroy, obj)
+        obj
+    end
 end
 
 Base.unsafe_convert(::Type{cudssMatrix_t}, matrix::CudssMatrix) = matrix.matrix

diff --git a/src/interfaces.jl b/src/interfaces.jl
@@ -49,6 +49,10 @@ end
     cudss_set(matrix::CudssMatrix{T}, A::CuMatrix{T})
     cudss_set(matrix::CudssMatrix{T}, A::CuSparseMatrixCSR{T,Cint})
     cudss_set(solver::CudssSolver{T}, A::CuSparseMatrixCSR{T,Cint})
+    cudss_set(matrix::CudssMatrix{T}, v::Vector{CuVector{T}})
+    cudss_set(matrix::CudssMatrix{T}, A::Vector{CuMatrix{T}})
+    cudss_set(matrix::CudssMatrix{T}, A::Vector{CuSparseMatrixCSR{T,Cint}})
+    cudss_set(solver::CudssSolver{T}, A::Vector{CuSparseMatrixCSR{T,Cint}})
     cudss_set(solver::CudssSolver, parameter::String, value)
     cudss_set(config::CudssConfig, parameter::String, value)
     cudss_set(data::CudssData, parameter::String, value)
@@ -93,6 +97,26 @@ function cudss_set(solver::CudssSolver{T}, A::CuSparseMatrixCSR{T,Cint}) where T
   cudss_set(solver.matrix, A)
 end
 
+function cudss_set(matrix::CudssMatrix{T}, v::Vector{CuVector{T}}) where T <: BlasFloat
+  cudssMatrixSetBatchValues(matrix, v)
+end
+
+function cudss_set(matrix::CudssMatrix{T}, v::Vector{CuMatrix{T}}) where T <: BlasFloat
+  cudssMatrixSetBatchValues(matrix, v)
+end
+
+function cudss_set(matrix::CudssMatrix{T}, A::Vector{CuSparseMatrixCSR{T,Cint}}) where T <: BlasFloat
+  rowsPtrs = [A.rowPtr for Aᵢ in A]
+  colVals = [A.colVal for Aᵢ in A]
+  nzVals = [A.nzVal for Aᵢ in A]
+  PTR_CU_NULL = Ptr{CuPtr{Cvoid}}()
+  cudssMatrixSetBatchCsrPointers(matrix, rowsPtrs, PTR_CU_NULL, colVals, nzVals)
+end
+
+function cudss_set(solver::CudssSolver{T}, A::Vector{CuSparseMatrixCSR{T,Cint}}) where T <: BlasFloat
+  cudss_set(solver.matrix, A)
+end
+
 function cudss_set(solver::CudssSolver, parameter::String, value)
   if parameter ∈ CUDSS_CONFIG_PARAMETERS
     cudss_set(solver.config, parameter, value)