Introduce Abstract types for sparse arrays

.gitignore

@@ -12,3 +12,5 @@ Manifest.toml
 
 # MacOS generated files
 *.DS_Store
+
+/.vscode/

Project.toml

@@ -12,6 +12,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
@@ -24,5 +25,6 @@ Printf = "1"
 Random = "1"
 Reexport = "1"
 Serialization = "1"
+SparseArrays = "1"
 Statistics = "1"
 julia = "1.10"

lib/GPUArraysCore/Project.toml

@@ -5,7 +5,9 @@ version = "0.2.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 Adapt = "4.0"
+SparseArrays = "1"
 julia = "1.6"

lib/GPUArraysCore/src/GPUArraysCore.jl

@@ -1,14 +1,19 @@
 module GPUArraysCore
 
 using Adapt
-
+using SparseArrays
 
 ## essential types
 
 export AbstractGPUArray, AbstractGPUVector, AbstractGPUMatrix, AbstractGPUVecOrMat,
        WrappedGPUArray, AnyGPUArray, AbstractGPUArrayStyle,
        AnyGPUArray, AnyGPUVector, AnyGPUMatrix
 
+export AbstractGPUSparseArray, AbstractGPUSparseMatrix, AbstractGPUSparseVector, AbstractGPUSparseVecOrMat,
+       WrappedGPUSparseArray, AnyGPUSparseArray, AnyGPUSparseVector, AnyGPUSparseMatrix,
+       AbstractGPUSparseMatrixCSC, AbstractGPUSparseMatrixCSR, AbstractGPUSparseMatrixCOO,
+       WrappedGPUSparseMatrixCSC, AnyGPUSparseMatrixCSC, WrappedGPUSparseMatrixCSR, AnyGPUSparseMatrixCSR, WrappedGPUSparseMatrixCOO, AnyGPUSparseMatrixCOO
+
 """
     AbstractGPUArray{T, N} <: DenseArray{T, N}
 
@@ -28,6 +33,31 @@ const AnyGPUArray{T,N} = Union{AbstractGPUArray{T,N}, WrappedGPUArray{T,N}}
 const AnyGPUVector{T} = AnyGPUArray{T, 1}
 const AnyGPUMatrix{T} = AnyGPUArray{T, 2}
 
+## sparse arrays
+
+abstract type AbstractGPUSparseArray{Tv, Ti, N} <: AbstractSparseArray{Tv, Ti, N} end
+
+const AbstractGPUSparseMatrix{Tv, Ti} = AbstractGPUSparseArray{Tv, Ti, 2}
+const AbstractGPUSparseVector{Tv, Ti} = AbstractGPUSparseArray{Tv, Ti, 1}
+const AbstractGPUSparseVecOrMat{Tv, Ti} = Union{AbstractGPUSparseVector{Tv, Ti}, AbstractGPUSparseMatrix{Tv, Ti}}
+
+const WrappedGPUSparseArray{Tv, Ti, N} = WrappedArray{Tv,N,AbstractGPUSparseArray,AbstractGPUSparseArray{Tv, Ti, N}}
+const AnyGPUSparseArray{Tv, Ti, N} = Union{AbstractGPUSparseArray{Tv, Ti, N}, WrappedGPUSparseArray{Tv, Ti, N}}
+const AnyGPUSparseVector{Tv, Ti} = AnyGPUSparseArray{Tv, Ti, 1}
+const AnyGPUSparseMatrix{Tv, Ti} = AnyGPUSparseArray{Tv, Ti, 2}
+
+abstract type AbstractGPUSparseMatrixCSC{Tv,Ti<:Integer} <: AbstractGPUSparseMatrix{Tv,Ti} end
+abstract type AbstractGPUSparseMatrixCSR{Tv,Ti<:Integer} <: AbstractGPUSparseMatrix{Tv,Ti} end
+abstract type AbstractGPUSparseMatrixCOO{Tv,Ti<:Integer} <: AbstractGPUSparseMatrix{Tv,Ti} end
+
+const WrappedGPUSparseMatrixCSC{Tv,Ti} = WrappedArray{Tv,AbstractGPUSparseMatrixCSC,AbstractGPUSparseMatrixCSC{Tv,Ti}}
+const AnyGPUSparseMatrixCSC{Tv,Ti} = Union{AbstractGPUSparseMatrixCSC{Tv,Ti}, WrappedGPUSparseMatrixCSC{Tv,Ti}}
+
+const WrappedGPUSparseMatrixCSR{Tv,Ti} = WrappedArray{Tv,AbstractGPUSparseMatrixCSR,AbstractGPUSparseMatrixCSR{Tv,Ti}}
+const AnyGPUSparseMatrixCSR{Tv,Ti} = Union{AbstractGPUSparseMatrixCSR{Tv,Ti}, WrappedGPUSparseMatrixCSR{Tv,Ti}}
+
+const WrappedGPUSparseMatrixCOO{Tv,Ti} = WrappedArray{Tv,AbstractGPUSparseMatrixCOO,AbstractGPUSparseMatrixCOO{Tv,Ti}}
+const AnyGPUSparseMatrixCOO{Tv,Ti} = Union{AbstractGPUSparseMatrixCOO{Tv,Ti}, WrappedGPUSparseMatrixCOO{Tv,Ti}}
 
 ## broadcasting
 

lib/JLArrays/Project.toml

@@ -8,10 +8,12 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 Adapt = "2.0, 3.0, 4.0"
 GPUArrays = "11.1"
 KernelAbstractions = "0.9"
 Random = "1"
+SparseArrays = "1"
 julia = "1.8"

lib/JLArrays/src/JLArrays.jl

@@ -11,6 +11,7 @@ export JLArray, JLVector, JLMatrix, jl, JLBackend
 using GPUArrays
 
 using Adapt
+using SparseArrays
 
 import KernelAbstractions
 import KernelAbstractions: Adapt, StaticArrays, Backend, Kernel, StaticSize, DynamicSize, partition, blocks, workitems, launch_config
@@ -387,4 +388,6 @@ Adapt.adapt_storage(::JLBackend, a::Array) = Adapt.adapt(JLArrays.JLArray, a)
 Adapt.adapt_storage(::JLBackend, a::JLArrays.JLArray) = a
 Adapt.adapt_storage(::KernelAbstractions.CPU, a::JLArrays.JLArray) = convert(Array, a)
 
+include("sparse.jl")
+
 end

lib/JLArrays/src/sparse.jl

@@ -0,0 +1,33 @@
+export JLSparseMatrixCSC
+
+struct JLSparseMatrixCSC{Tv,Ti<:Integer} <: AbstractGPUSparseMatrixCSC{Tv,Ti}
+    m::Int                  # Number of rows
+    n::Int                  # Number of columns
+    colptr::JLVector{Ti}      # Column i is in colptr[i]:(colptr[i+1]-1)
+    rowval::JLVector{Ti}      # Row indices of stored values
+    nzval::JLVector{Tv}       # Stored values, typically nonzeros
+
+    function JLSparseMatrixCSC{Tv,Ti}(m::Integer, n::Integer, colptr::JLVector{Ti},
+                            rowval::JLVector{Ti}, nzval::JLVector{Tv}) where {Tv,Ti<:Integer}
+        SparseArrays.sparse_check_Ti(m, n, Ti)
+        GPUArrays._goodbuffers_csc(m, n, colptr, rowval, nzval) ||
+            throw(ArgumentError("Invalid buffers for JLSparseMatrixCSC construction n=$n, colptr=$(summary(colptr)), rowval=$(summary(rowval)), nzval=$(summary(nzval))"))
+        new(Int(m), Int(n), colptr, rowval, nzval)
+    end
+end
+function JLSparseMatrixCSC(m::Integer, n::Integer, colptr::JLVector, rowval::JLVector, nzval::JLVector)
+    Tv = eltype(nzval)
+    Ti = promote_type(eltype(colptr), eltype(rowval))
+    SparseArrays.sparse_check_Ti(m, n, Ti)
+    # SparseArrays.sparse_check(n, colptr, rowval, nzval) # TODO: this uses scalar indexing
+    # silently shorten rowval and nzval to usable index positions.
+    maxlen = abs(widemul(m, n))
+    isbitstype(Ti) && (maxlen = min(maxlen, typemax(Ti) - 1))
+    length(rowval) > maxlen && resize!(rowval, maxlen)
+    length(nzval) > maxlen && resize!(nzval, maxlen)
+    JLSparseMatrixCSC{Tv,Ti}(m, n, colptr, rowval, nzval)
+end
+
+JLSparseMatrixCSC(A::SparseMatrixCSC) = JLSparseMatrixCSC(A.m, A.n, JLVector(A.colptr), JLVector(A.rowval), JLVector(A.nzval))
+
+Base.copy(A::JLSparseMatrixCSC) = JLSparseMatrixCSC(A.m, A.n, copy(A.colptr), copy(A.rowval), copy(A.nzval))

src/GPUArrays.jl

@@ -4,6 +4,9 @@ using KernelAbstractions
 using Serialization
 using Random
 using LinearAlgebra
+using SparseArrays
+using SparseArrays: getcolptr, getrowval, getnzval
+
 using Printf
 
 using LinearAlgebra.BLAS
@@ -15,8 +18,6 @@ using LLVM.Interop
 using Reexport
 @reexport using GPUArraysCore
 
-using KernelAbstractions
-
 # device functionality
 include("device/abstractarray.jl")
 
@@ -33,6 +34,7 @@ include("host/math.jl")
 include("host/random.jl")
 include("host/quirks.jl")
 include("host/uniformscaling.jl")
+include("host/sparse.jl")
 include("host/statistics.jl")
 
 

src/host/sparse.jl

@@ -0,0 +1,79 @@
+## General Sparse Matrix
+
+Base.size(A::AbstractGPUSparseMatrix) = (A.m, A.n)
+
+SparseArrays.nonzeros(A::AbstractGPUSparseMatrix) = A.nzval
+SparseArrays.getnzval(A::AbstractGPUSparseMatrix) = nonzeros(A)
+SparseArrays.nnz(A::AbstractGPUSparseMatrix) = length(nzval(A))
+
+function LinearAlgebra.rmul!(A::AbstractGPUSparseMatrix, x::Number)
+    rmul!(SparseArrays.getnzval(A), x)
+    return A
+end
+
+function LinearAlgebra.lmul!(x::Number, A::AbstractGPUSparseMatrix)
+    lmul!(x, SparseArrays.getnzval(A))
+    return A
+end
+
+## CSC Matrix
+
+SparseArrays.getcolptr(A::AbstractGPUSparseMatrixCSC) = A.colptr
+SparseArrays.rowvals(A::AbstractGPUSparseMatrixCSC) = A.rowval
+SparseArrays.getrowval(A::AbstractGPUSparseMatrixCSC) = rowvals(A)
+# SparseArrays.nzrange(A::AbstractGPUSparseMatrixCSC, col::Integer) = getcolptr(A)[col]:(getcolptr(A)[col+1]-1) # TODO: this uses scalar indexing
+
+function _goodbuffers_csc(m, n, colptr, rowval, nzval)
+    return (length(colptr) == n + 1 && length(rowval) == length(nzval))
+    # TODO: also add the condition that colptr[end] - 1 == length(nzval) (allowscalar?)
+end
+
+@inline function LinearAlgebra.mul!(C::AbstractGPUVector, A::AnyGPUSparseMatrixCSC, B::AbstractGPUVector, α::Number, β::Number)
+    return LinearAlgebra.generic_matvecmul!(C, LinearAlgebra.wrapper_char(A), LinearAlgebra._unwrap(A), B, LinearAlgebra.MulAddMul(α, β))
+end
+
+@inline function LinearAlgebra.generic_matvecmul!(C::AbstractGPUVector, tA, A::AbstractGPUSparseMatrixCSC, B::AbstractGPUVector, _add::LinearAlgebra.MulAddMul)
+    return SparseArrays.spdensemul!(C, tA, 'N', A, B, _add)
+end
+
+Base.@constprop :aggressive function SparseArrays.spdensemul!(C::AbstractGPUVecOrMat, tA, tB, A::AbstractGPUSparseMatrixCSC, B::AbstractGPUVecOrMat, _add::LinearAlgebra.MulAddMul)
+    if tA == 'N'
+        return _spmatmul!(C, A, wrap(B, tB), _add.alpha, _add.beta)
+    else
+        throw(ArgumentError("tA different from 'N' not yet supported"))
+    end
+end
+
+function _spmatmul!(C::AbstractGPUVecOrMat, A::AbstractGPUSparseMatrixCSC, B::AbstractGPUVecOrMat, α::Number, β::Number)
+    size(A, 2) == size(B, 1) ||
+        throw(DimensionMismatch("second dimension of A, $(size(A,2)), does not match the first dimension of B, $(size(B,1))"))
+    size(A, 1) == size(C, 1) ||
+        throw(DimensionMismatch("first dimension of A, $(size(A,1)), does not match the first dimension of C, $(size(C,1))"))
+    size(B, 2) == size(C, 2) ||
+        throw(DimensionMismatch("second dimension of B, $(size(B,2)), does not match the second dimension of C, $(size(C,2))"))
+
+    A_colptr = getcolptr(A)
+    A_rowval = rowvals(A)
+    A_nzval = getnzval(A)
+    β != one(β) && LinearAlgebra._rmul_or_fill!(C, β)
+
+    @kernel function kernel_spmatmul!(C, @Const(A_colptr), @Const(A_rowval), @Const(A_nzval), @Const(B))
+        k, col = @index(Global, NTuple)
+
+        @inbounds axj = B[col, k] * α
+        @inbounds for j in A_colptr[col]:(A_colptr[col+1]-1) # nzrange(A, col)
+            KernelAbstractions.@atomic C[A_rowval[j], k] += A_nzval[j] * axj
+        end
+    end
+
+    backend_C = KernelAbstractions.get_backend(C)
+    backend_A = KernelAbstractions.get_backend(A_nzval)
+    backend_B = KernelAbstractions.get_backend(B)
+
+    backend_A == backend_B == backend_C || throw(ArgumentError("All arrays must be on the same backend"))
+
+    kernel! = kernel_spmatmul!(backend_A)
+    kernel!(C, A_colptr, A_rowval, A_nzval, B; ndrange=(size(C, 2), size(A, 2)))
+
+    return C
+end