Skip to content

Commit

Permalink
Add conversions between CuSparseVector and CuSparseMatrices (#2489)
Browse files Browse the repository at this point in the history
  • Loading branch information
amontoison authored Sep 17, 2024
1 parent 1e1f4ef commit cdae2d3
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 2 deletions.
45 changes: 43 additions & 2 deletions lib/cusparse/conversions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -614,12 +614,12 @@ function CuSparseMatrixCOO{Tv}(csc::CuSparseMatrixCSC{Tv}; index::SparseChar='O'
coo = sort_coo(coo, 'R')
end

### BSR to COO and viceversa
### BSR to COO and vice-versa

CuSparseMatrixBSR(coo::CuSparseMatrixCOO, blockdim) = CuSparseMatrixBSR(CuSparseMatrixCSR(coo), blockdim) # no direct conversion
CuSparseMatrixCOO(bsr::CuSparseMatrixBSR) = CuSparseMatrixCOO(CuSparseMatrixCSR(bsr)) # no direct conversion

### BSR to CSC and viceversa
### BSR to CSC and vice-versa

CuSparseMatrixBSR(csc::CuSparseMatrixCSC, blockdim) = CuSparseMatrixBSR(CuSparseMatrixCSR(csc), blockdim) # no direct conversion
CuSparseMatrixCSC(bsr::CuSparseMatrixBSR) = CuSparseMatrixCSC(CuSparseMatrixCSR(bsr)) # no direct conversion
Expand Down Expand Up @@ -668,3 +668,44 @@ end
function CuSparseMatrixCOO(A::CuMatrix{T}; index::SparseChar='O') where {T}
densetosparse(A, :coo, index)
end

## CuSparseVector to CuSparseMatrices and vice-versa
function CuSparseVector(A::CuSparseMatrixCSC{T}) where T
m, n = size(A)
(n == 1) || error("A doesn't have one column and can't be converted to a CuSparseVector.")
CuSparseVector{T}(A.rowVal, A.nzVal, m)
end

# no direct conversion
function CuSparseVector(A::CuSparseMatrixCSR{T}) where T
m, n = size(A)
(n == 1) || error("A doesn't have one column and can't be converted to a CuSparseVector.")
B = CuSparseMatrixCSC{T}(A)
CuSparseVector(B)
end

function CuSparseVector(A::CuSparseMatrixCOO{T}) where T
m, n = size(A)
(n == 1) || error("A doesn't have one column and can't be converted to a CuSparseVector.")
CuSparseVector{T}(A.rowInd, A.nzVal, m)
end

function CuSparseMatrixCSC(x::CuSparseVector{T}) where T
n = length(x)
colPtr = CuVector{Int32}([1; nnz(x)+1])
CuSparseMatrixCSC{T}(colPtr, x.iPtr, nonzeros(x), (n,1))
end

# no direct conversion
function CuSparseMatrixCSR(x::CuSparseVector{T}) where T
A = CuSparseMatrixCSC(x)
CuSparseMatrixCSR{T}(A)
end

function CuSparseMatrixCOO(x::CuSparseVector{T}) where T
n = length(x)
nnzx = nnz(x)
colInd = CuVector{Int32}(undef, nnzx)
fill!(colInd, one(Int32))
CuSparseMatrixCOO{T}(x.iPtr, colInd, nonzeros(x), (n,1), nnzx)
end
18 changes: 18 additions & 0 deletions test/libraries/cusparse/conversions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,24 @@ end
end
end

if !(v"12.0" <= CUSPARSE.version() < v"12.1")
x = [0.0; 1.0; 2.0; 0.0; 3.0] |> SparseVector |> CuSparseVector
A = Matrix{Float64}(undef, 5, 1)
A[:, 1] .= [0.0; 1.0; 2.0; 0.0; 3.0]
A = SparseMatrixCSC(A)
for CuSparseMatrixType in (CuSparseMatrixCSC, CuSparseMatrixCSR, CuSparseMatrixCOO)
@testset "conversion CuSparseVector --> $CuSparseMatrixType" begin
B = CuSparseMatrixType(x)
@test collect(B)[:] collect(x)
end
@testset "conversion $CuSparseMatrixType --> CuSparseVector" begin
B = CuSparseMatrixType(A)
y = CuSparseVector(B)
@test collect(B)[:] collect(y)
end
end
end

for (n, bd, p) in [(100, 5, 0.02), (5, 1, 0.8), (4, 2, 0.5)]
v"12.0" <= CUSPARSE.version() < v"12.1" && n == 4 && continue
@testset "conversions between CuSparseMatrices (n, bd, p) = ($n, $bd, $p)" begin
Expand Down

0 comments on commit cdae2d3

Please sign in to comment.