[rocSPARSE] Interface additional sparse routines (#707)

* [rocSPARSE] Interface additional sparse routines

* Interface sparsetodense and densetosparse

* Add more constructors for ROCSparseMatrixDescriptor

* Fix the errors

src/sparse/conversions.jl

@@ -201,6 +201,15 @@ for (elty, felty) in ((:Int32, :Float32), (:Int64, :Float64), (:Int128, :Complex
     end
 end
 
+## ROCSparseVector to ROCVector
+ROCVector(x::ROCSparseVector{T}) where {T} = ROCVector{T}(x)
+
+function ROCVector{T}(sv::ROCSparseVector{T}) where {T}
+    n = length(sv)
+    dv = AMDGPU.zeros(T, n)
+    scatter!(dv, sv, 'O')
+end
+
 ## CSR to BSR and vice-versa
 
 for (fname,elty) in ((:rocsparse_scsr2bsr, :Float32),
@@ -400,7 +409,7 @@ for (elty, welty) in ((:Float16, :Float32), (:ComplexF16, :ComplexF32))
     end
 end
 
-function Base.copyto!(dest::Array{T, 2}, src::AbstractROCSparseMatrix{T}) where T
+function Base.copyto!(dest::Matrix{T}, src::AbstractROCSparseMatrix{T}) where T
     copyto!(dest, ROCMatrix{T}(src))
 end
 

src/sparse/generic.jl

@@ -2,13 +2,128 @@
 
 ## API functions
 
+function sparsetodense(A::Union{ROCSparseMatrixCSC{T},ROCSparseMatrixCSR{T},ROCSparseMatrixCOO{T}}, index::SparseChar,
+                       algo::rocsparse_sparse_to_dense_alg=rocsparse_sparse_to_dense_alg_default) where {T}
+    m,n = size(A)
+    B = ROCMatrix{T}(undef, m, n)
+    desc_sparse = ROCSparseMatrixDescriptor(A, index)
+    desc_dense = ROCDenseMatrixDescriptor(B)
+
+    function bufferSize()
+        out = Ref{Csize_t}()
+        rocsparse_sparse_to_dense(handle(), desc_sparse, desc_dense, algo, out, C_NULL)
+        return out[]
+    end
+
+    buffer_size = Ref{Csize_t}()
+    with_workspace(bufferSize) do buffer
+        buffer_size[] = sizeof(buffer)
+        rocsparse_sparse_to_dense(handle(), desc_sparse, desc_dense, algo, buffer_size, buffer)
+    end
+    return B
+end
+
+function densetosparse(A::ROCMatrix{T}, fmt::Symbol, index::SparseChar,
+                       algo::rocsparse_dense_to_sparse_alg=rocsparse_dense_to_sparse_alg_default) where {T}
+    m,n = size(A)
+    local rowPtr, colPtr, desc_sparse, B
+    if fmt == :coo
+        desc_sparse = ROCSparseMatrixDescriptor(ROCSparseMatrixCOO, T, Cint, m, n, index)
+    elseif fmt == :csr
+        rowPtr = ROCVector{Cint}(undef, m+1)
+        desc_sparse = ROCSparseMatrixDescriptor(ROCSparseMatrixCSR, rowPtr, T, Cint, m, n, index)
+    elseif fmt == :csc
+        colPtr = ROCVector{Cint}(undef, n+1)
+        desc_sparse = ROCSparseMatrixDescriptor(ROCSparseMatrixCSC, colPtr, T, Cint, m, n, index)
+    else
+        error("Format :$fmt not available, use :csc, :csr or :coo.")
+    end
+    desc_dense = ROCDenseMatrixDescriptor(A)
+
+    function bufferSize()
+        out = Ref{Csize_t}()
+        rocsparse_dense_to_sparse(handle(), desc_dense, desc_sparse, algo, out, C_NULL)
+        return out[]
+    end
+
+    buffer_size = Ref{Csize_t}()
+    with_workspace(bufferSize) do buffer
+        buffer_size[] = sizeof(buffer)
+        # Analysis
+        rocsparse_dense_to_sparse(handle(), desc_dense, desc_sparse, algo, C_NULL, buffer)
+        nnzB = Ref{Int64}()
+        rocsparse_spmat_get_size(desc_sparse, Ref{Int64}(), Ref{Int64}(), nnzB)
+        if fmt == :coo
+            rowInd = ROCVector{Cint}(undef, nnzB[])
+            colInd = ROCVector{Cint}(undef, nnzB[])
+            nzVal = ROCVector{T}(undef, nnzB[])
+            B = ROCSparseMatrixCOO{T, Cint}(rowInd, colInd, nzVal, (m,n))
+            rocsparse_coo_set_pointers(desc_sparse, B.rowInd, B.colInd, B.nzVal)
+        elseif fmt == :csr
+            colVal = ROCVector{Cint}(undef, nnzB[])
+            nzVal = ROCVector{T}(undef, nnzB[])
+            B = ROCSparseMatrixCSR{T, Cint}(rowPtr, colVal, nzVal, (m,n))
+            rocsparse_csr_set_pointers(desc_sparse, B.rowPtr, B.colVal, B.nzVal)
+        elseif fmt == :csc
+            rowVal = ROCVector{Cint}(undef, nnzB[])
+            nzVal = ROCVector{T}(undef, nnzB[])
+            B = ROCSparseMatrixCSC{T, Cint}(colPtr, rowVal, nzVal, (m,n))
+            rocsparse_csc_set_pointers(desc_sparse, B.colPtr, B.rowVal, B.nzVal)
+        else
+            error("Format :$fmt not available, use :csc, :csr or :coo.")
+        end
+        rocsparse_dense_to_sparse(handle(), desc_dense, desc_sparse, algo, buffer_size, buffer)
+    end
+    return B
+end
+
 function gather!(X::ROCSparseVector, Y::ROCVector, index::SparseChar)
     descX = ROCSparseVectorDescriptor(X, index)
     descY = ROCDenseVectorDescriptor(Y)
     rocsparse_gather(handle(), descY, descX)
     X
 end
 
+function scatter!(Y::ROCVector, X::ROCSparseVector, index::SparseChar)
+    descX = ROCSparseVectorDescriptor(X, index)
+    descY = ROCDenseVectorDescriptor(Y)
+    rocsparse_scatter(handle(), descX, descY)
+    return Y
+end
+
+function axpby!(alpha::Number, X::ROCSparseVector{T}, beta::Number, Y::ROCVector{T}, index::SparseChar) where {T}
+    descX = ROCSparseVectorDescriptor(X, index)
+    descY = ROCDenseVectorDescriptor(Y)
+    rocsparse_axpby(handle(), Ref{T}(alpha), descX, Ref{T}(beta), descY)
+    return Y
+end
+
+function rot!(X::ROCSparseVector{T}, Y::ROCVector{T}, c::Number, s::Number, index::SparseChar) where {T}
+    descX = ROCSparseVectorDescriptor(X, index)
+    descY = ROCDenseVectorDescriptor(Y)
+    rocsparse_rot(handle(), Ref{T}(c), Ref{T}(s), descX, descY)
+    return X, Y
+end
+
+function vv!(transx::SparseChar, X::ROCSparseVector{T}, Y::DenseROCVector{T}, index::SparseChar) where {T}
+    descX = ROCSparseVectorDescriptor(X, index)
+    descY = ROCDenseVectorDescriptor(Y)
+    result = Ref{T}()
+
+    function bufferSize()
+        out = Ref{Csize_t}()
+        rocsparse_spvv(handle(), transx, descX, descY, result, T, out, C_NULL)
+        return out[]
+    end
+
+    buffer_size = Ref{Csize_t}()
+    with_workspace(bufferSize) do buffer
+        buffer_size[] = sizeof(buffer)
+        rocsparse_spvv(handle(), transx, descX, descY, result, T, buffer_size, buffer)
+    end
+    return result[]
+end
+
 function mv!(
     transa::SparseChar, alpha::Number, A::Union{ROCSparseMatrixCSR{T}, ROCSparseMatrixCSC{T}, ROCSparseMatrixCOO{T}},
     X::DenseROCVector{T}, beta::Number, Y::DenseROCVector{T}, index::SparseChar,

src/sparse/helpers.jl

@@ -104,6 +104,33 @@ Base.unsafe_convert(::Type{rocsparse_dnmat_descr}, desc::ROCDenseMatrixDescripto
 mutable struct ROCSparseMatrixDescriptor
     handle::rocsparse_spmat_descr
 
+    function ROCSparseMatrixDescriptor(A::ROCSparseMatrixCOO, IndexBase::Char; transposed::Bool=false)
+        desc_ref = Ref{rocsparse_spmat_descr}()
+        if transposed
+            rocsparse_create_coo_descr(
+                desc_ref, reverse(size(A))..., nnz(A),
+                A.colInd, A.rowInd, nonzeros(A),
+                eltype(A.colInd), IndexBase, eltype(nonzeros(A))
+            )
+        else
+            rocsparse_create_coo_descr(
+                desc_ref, size(A)..., nnz(A),
+                A.rowInd, A.colInd, nonzeros(A),
+                eltype(A.rowInd), IndexBase, eltype(nonzeros(A))
+            )
+        end
+        obj = new(desc_ref[])
+        return finalizer(rocsparse_destroy_spmat_descr, obj)
+    end
+
+    function ROCSparseMatrixDescriptor(::Type{ROCSparseMatrixCOO}, Tv::DataType, Ti::DataType, m::Integer, n::Integer, IndexBase::Char)
+        desc_ref = Ref{rocsparse_spmat_descr}()
+        rocsparse_create_coo_descr(desc_ref, m, n, Ti(0), C_NULL, C_NULL, C_NULL, Ti, IndexBase, Tv)
+        obj = new(desc_ref[])
+        finalizer(rocsparse_destroy_spmat_descr, obj)
+        return obj
+    end
+
     function ROCSparseMatrixDescriptor(A::ROCSparseMatrixCSR, IndexBase::Char; transposed::Bool=false)
         desc_ref = Ref{rocsparse_spmat_descr}()
         if transposed
@@ -121,6 +148,14 @@ mutable struct ROCSparseMatrixDescriptor
         return finalizer(rocsparse_destroy_spmat_descr, obj)
     end
 
+    function ROCSparseMatrixDescriptor(::Type{ROCSparseMatrixCSR}, rowPtr::ROCVector, Tv::DataType, Ti::DataType, m::Integer, n::Integer, IndexBase::Char)
+        desc_ref = Ref{rocsparse_spmat_descr}()
+        rocsparse_create_csr_descr(desc_ref, m, n, Ti(0), rowPtr, C_NULL, C_NULL, Ti, Ti, IndexBase, Tv)
+        obj = new(desc_ref[])
+        finalizer(rocsparse_destroy_spmat_descr, obj)
+        return obj
+    end
+
     function ROCSparseMatrixDescriptor(A::ROCSparseMatrixCSC, IndexBase::Char; transposed::Bool=false)
         desc_ref = Ref{rocsparse_spmat_descr}()
         if transposed
@@ -138,23 +173,12 @@ mutable struct ROCSparseMatrixDescriptor
         return finalizer(rocsparse_destroy_spmat_descr, obj)
     end
 
-    function ROCSparseMatrixDescriptor(A::ROCSparseMatrixCOO, IndexBase::Char; transposed::Bool=false)
+    function ROCSparseMatrixDescriptor(::Type{ROCSparseMatrixCSC}, colPtr::ROCVector, Tv::DataType, Ti::DataType, m::Integer, n::Integer, IndexBase::Char)
         desc_ref = Ref{rocsparse_spmat_descr}()
-        if transposed
-            rocsparse_create_coo_descr(
-                desc_ref, reverse(size(A))..., nnz(A),
-                A.colInd, A.rowInd, nonzeros(A),
-                eltype(A.colInd), IndexBase, eltype(nonzeros(A))
-            )
-        else
-            rocsparse_create_coo_descr(
-                desc_ref, size(A)..., nnz(A),
-                A.rowInd, A.colInd, nonzeros(A),
-                eltype(A.rowInd), IndexBase, eltype(nonzeros(A))
-            )
-        end
+        rocsparse_create_csc_descr(desc_ref, m, n, Ti(0), colPtr, C_NULL, C_NULL, Ti, Ti, IndexBase, Tv)
         obj = new(desc_ref[])
-        return finalizer(rocsparse_destroy_spmat_descr, obj)
+        finalizer(rocsparse_destroy_spmat_descr, obj)
+        return obj
     end
 end
 

src/sparse/interfaces.jl

@@ -22,6 +22,12 @@ function mm_wrapper(
     mm!(transa, transb, alpha, A, B, beta, C, 'O')
 end
 
+LinearAlgebra.dot(x::ROCSparseVector{T}, y::DenseROCVector{T}) where {T <: BlasReal} = vv!('N', x, y, 'O')
+LinearAlgebra.dot(x::DenseROCVector{T}, y::ROCSparseVector{T}) where {T <: BlasReal} = dot(y, x)
+
+LinearAlgebra.dot(x::ROCSparseVector{T}, y::DenseROCVector{T}) where {T <: BlasComplex} = vv!('C', x, y, 'O')
+LinearAlgebra.dot(x::DenseROCVector{T}, y::ROCSparseVector{T}) where {T <: BlasComplex} = conj(dot(y,x))
+
 # legacy methods with final MulAddMul argument
 LinearAlgebra.generic_matvecmul!(C::ROCVector{T}, tA::AbstractChar, A::ROCSparseMatrix{T}, B::DenseROCVector{T}, _add::MulAddMul) where T <: BlasFloat =
     LinearAlgebra.generic_matvecmul!(C, tA, A, B, _add.alpha, _add.beta)

test/rocsparse/generic.jl

@@ -1,3 +1,95 @@
+fmt = Dict(ROCSparseMatrixCSC => :csc,
+           ROCSparseMatrixCSR => :csr,
+           ROCSparseMatrixCOO => :coo)
+
+for SparseMatrixType in [ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseMatrixCOO]
+    @testset "$SparseMatrixType -- densetosparse algo=$algo" for algo in [rocSPARSE.rocsparse_dense_to_sparse_alg_default]
+        @testset "densetosparse $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
+            A_sparse = sprand(T, 10, 20, 0.5)
+            A_dense = Matrix{T}(A_sparse)
+            dA_dense = ROCMatrix{T}(A_dense)
+            dA_sparse = rocSPARSE.densetosparse(dA_dense, fmt[SparseMatrixType], 'O', algo)
+            @test A_sparse ≈ collect(dA_sparse)
+        end
+    end
+    @testset "$SparseMatrixType -- sparsetodense algo=$algo" for algo in [rocSPARSE.rocsparse_sparse_to_dense_alg_default]
+        @testset "sparsetodense $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
+            A_dense = rand(T, 10, 20)
+            A_sparse = sparse(A_dense)
+            dA_sparse = SparseMatrixType(A_sparse)
+            dA_dense = rocSPARSE.sparsetodense(dA_sparse, 'O', algo)
+            @test A_dense ≈ collect(dA_dense)
+        end
+    end
+end
+
+@testset "gather! $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
+    X = sprand(T, 20, 0.5)
+    dX = ROCSparseVector{T}(X)
+    Y = rand(T, 20)
+    dY = ROCVector{T}(Y)
+    rocSPARSE.gather!(dX, dY, 'O')
+    Z = copy(X)
+    for i = 1:nnz(X)
+        Z[X.nzind[i]] = Y[X.nzind[i]]
+    end
+    @test Z ≈ sparse(collect(dX))
+end
+
+@testset "scatter! $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
+    X = sprand(T, 20, 0.5)
+    dX = ROCSparseVector{T}(X)
+    Y = rand(T, 20)
+    dY = ROCVector{T}(Y)
+    rocSPARSE.scatter!(dY, dX, 'O')
+    Z = copy(Y)
+    for i = 1:nnz(X)
+        Z[X.nzind[i]] = X.nzval[i]
+    end
+    @test Z ≈ collect(dY)
+end
+
+@testset "axpby! $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
+    X = sprand(T, 20, 0.5)
+    dX = ROCSparseVector{T}(X)
+    Y = rand(T, 20)
+    dY = ROCVector{T}(Y)
+    alpha = rand(T)
+    beta = rand(T)
+    rocSPARSE.axpby!(alpha, dX, beta, dY, 'O')
+    @test alpha * X + beta * Y ≈ collect(dY)
+end
+
+@testset "rot! $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
+    X = sprand(T, 20, 0.5)
+    dX = ROCSparseVector{T}(X)
+    Y = rand(T, 20)
+    dY = ROCVector{T}(Y)
+    c = rand(T)
+    s = rand(T)
+    rocSPARSE.rot!(dX, dY, c, s, 'O')
+    W = copy(X)
+    Z = copy(Y)
+    for i = 1:nnz(X)
+        W[X.nzind[i]] =  c * X.nzval[i] + s * Y[X.nzind[i]]
+        Z[X.nzind[i]] = -s * X.nzval[i] + c * Y[X.nzind[i]]
+    end
+    @test W ≈ collect(dX)
+    @test Z ≈ collect(dY)
+end
+
+@testset "vv! $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
+    for (transx, opx) in [('N', identity), ('C', conj)]
+        T <: Real && transx == 'C' && continue
+        X = sprand(T, 20, 0.5)
+        dX = ROCSparseVector{T}(X)
+        Y = rand(T, 20)
+        dY = ROCVector{T}(Y)
+        result = rocSPARSE.vv!(transx, dX, dY, 'O')
+        @test sum(opx(X[i]) * Y[i] for i=1:20) ≈ result
+    end
+end
+
 @testset "generic mv!" for T in (Float32, Float64, ComplexF32, ComplexF64)
     A = sprand(T, 10, 10, 0.1)
     x = rand(T, 10)