From cdae2d3ea6085669b92bd58aadd4828387179cb6 Mon Sep 17 00:00:00 2001
From: Alexis Montoison <35051714+amontoison@users.noreply.github.com>
Date: Tue, 17 Sep 2024 12:37:53 -0500
Subject: [PATCH] Add conversions between CuSparseVector and CuSparseMatrices
 (#2489)

---
 lib/cusparse/conversions.jl            | 45 ++++++++++++++++++++++++--
 test/libraries/cusparse/conversions.jl | 18 +++++++++++
 2 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/lib/cusparse/conversions.jl b/lib/cusparse/conversions.jl
index a82b1bfc95..69dbaf3f9b 100644
--- a/lib/cusparse/conversions.jl
+++ b/lib/cusparse/conversions.jl
@@ -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
@@ -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
diff --git a/test/libraries/cusparse/conversions.jl b/test/libraries/cusparse/conversions.jl
index c1edc56a75..1a5119f289 100644
--- a/test/libraries/cusparse/conversions.jl
+++ b/test/libraries/cusparse/conversions.jl
@@ -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