Add direct constructors for sparse arrays

[nHackel]

This PR adds direct constructors for the sparse matrix types defined in GPUArrays.jl or rather adds a function for backends to generically define methods for, see also #677.

Changes:

• Introduced a generic function for backends to define constructors for sparse formats.
• Added direct constructors for sparse matrices with a specified format.
• Added a GPUSparseMatrix convenience constructor that defaults to GPUSparseMatrixCOO, which in turn falls back to SparseArrays.sparse for completeness.
• Added tests for CSC and CSR, which can also be extended to the other formats. I was unsure if the constructor tests should also assert behaviour, so far I've only checked for appropriate array types and size.

There currently is no COO and BSR implementation for JLArrays. I could add one as well, at least for the BSR I'd need to take a look at the format first or adapt from maybe CUDA.jl

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic master) to apply these changes.

Click here to view the suggested changes.

diff --git a/lib/JLArrays/src/JLArrays.jl b/lib/JLArrays/src/JLArrays.jl
index 3bd3571..b76a237 100644
--- a/lib/JLArrays/src/JLArrays.jl
+++ b/lib/JLArrays/src/JLArrays.jl
@@ -150,7 +150,7 @@ mutable struct JLSparseMatrixCSC{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCSC
         new{Tv, Ti}(colPtr, rowVal, nzVal, dims, length(nzVal))
     end
 end
-function GPUSparseMatrixCSC(colPtr::JLArray{Ti, 1}, rowVal::JLArray{Ti, 1}, nzVal::JLArray{Tv, 1}, dims::NTuple{2,<:Integer}) where {Tv, Ti <: Integer}
+function GPUSparseMatrixCSC(colPtr::JLArray{Ti, 1}, rowVal::JLArray{Ti, 1}, nzVal::JLArray{Tv, 1}, dims::NTuple{2, <:Integer}) where {Tv, Ti <: Integer}
     return JLSparseMatrixCSC(colPtr, rowVal, nzVal, dims)
 end
 function JLSparseMatrixCSC(colPtr::JLArray{Ti, 1}, rowVal::JLArray{Ti, 1}, nzVal::JLArray{Tv, 1}, dims::NTuple{2,<:Integer}) where {Tv, Ti <: Integer}
@@ -184,7 +184,7 @@ end
 function JLSparseMatrixCSR(rowPtr::JLArray{Ti, 1}, colVal::JLArray{Ti, 1}, nzVal::JLArray{Tv, 1}, dims::NTuple{2,<:Integer}) where {Tv, Ti <: Integer}
     return JLSparseMatrixCSR{Tv, Ti}(rowPtr, colVal, nzVal, dims)
 end
-function GPUSparseMatrixCSR(rowPtr::JLArray{Ti, 1}, colVal::JLArray{Ti, 1}, nzVal::JLArray{Tv, 1}, dims::NTuple{2,<:Integer}) where {Tv, Ti <: Integer}
+function GPUSparseMatrixCSR(rowPtr::JLArray{Ti, 1}, colVal::JLArray{Ti, 1}, nzVal::JLArray{Tv, 1}, dims::NTuple{2, <:Integer}) where {Tv, Ti <: Integer}
     return JLSparseMatrixCSR(rowPtr, colVal, nzVal, dims)
 end
 function SparseArrays.SparseMatrixCSC(x::JLSparseMatrixCSR) 
diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl
index 146801e..ecfee75 100644
--- a/test/testsuite/sparse.jl
+++ b/test/testsuite/sparse.jl
@@ -154,11 +154,11 @@ end
 
 # Helper function to derive direct matrix formats:
 # Create colptr, rowval, nzval for m x n matrix with 3 values per column
-function csc_vectors(m::Int, n::Int, ::Type{ET}; I::Type{<:Integer}=Int32) where {ET}
+function csc_vectors(m::Int, n::Int, ::Type{ET}; I::Type{<:Integer} = Int32) where {ET}
     # Fixed, deterministic 3 nnz per column; random nz values
     colptr = Vector{I}(undef, n + 1)
     rowval = Vector{I}()
-    nzval  = Vector{ET}()
+    nzval = Vector{ET}()
 
     colptr[1] = I(1)
     nnz_acc = 0
@@ -172,29 +172,29 @@ function csc_vectors(m::Int, n::Int, ::Type{ET}; I::Type{<:Integer}=Int32) where
     end
     return colptr, rowval, nzval
 end
-function csr_vectors(m::Int, n::Int, ::Type{ET}; I::Type{<:Integer}=Int32) where {ET}
+function csr_vectors(m::Int, n::Int, ::Type{ET}; I::Type{<:Integer} = Int32) where {ET}
     # Build CSC for (n, m), then interpret as CSR for (m, n)
-    colptr_nm, rowval_nm, nzval_nm = csc_vectors(n, m, ET; I=I)
+    colptr_nm, rowval_nm, nzval_nm = csc_vectors(n, m, ET; I = I)
     rowptr = colptr_nm
     colind = rowval_nm
-    nzval  = nzval_nm
+    nzval = nzval_nm
     return rowptr, colind, nzval
 end
 # Construct appropriate sparse arrays
-function construct_sparse_matrix(AT::Type{<:GPUArrays.AbstractGPUSparseMatrixCSC}, ::Type{ET}, m::Int, n::Int; I::Type{<:Integer}=Int32) where {ET}
-    colptr, rowval, nzval = csc_vectors(m, n, ET; I=I)
+function construct_sparse_matrix(AT::Type{<:GPUArrays.AbstractGPUSparseMatrixCSC}, ::Type{ET}, m::Int, n::Int; I::Type{<:Integer} = Int32) where {ET}
+    colptr, rowval, nzval = csc_vectors(m, n, ET; I = I)
     dense_AT = GPUArrays.dense_array_type(AT)
     d_colptr = dense_AT(colptr)
     d_rowval = dense_AT(rowval)
-    d_nzval  = dense_AT(nzval)
+    d_nzval = dense_AT(nzval)
     return GPUSparseMatrixCSC(d_colptr, d_rowval, d_nzval, (m, n))
 end
-function construct_sparse_matrix(AT::Type{<:GPUArrays.AbstractGPUSparseMatrixCSR}, ::Type{ET}, m::Int, n::Int; I::Type{<:Integer}=Int32) where {ET}
-    rowptr, colind, nzval = csr_vectors(m, n, ET; I=I)
+function construct_sparse_matrix(AT::Type{<:GPUArrays.AbstractGPUSparseMatrixCSR}, ::Type{ET}, m::Int, n::Int; I::Type{<:Integer} = Int32) where {ET}
+    rowptr, colind, nzval = csr_vectors(m, n, ET; I = I)
     dense_AT = GPUArrays.dense_array_type(AT)
     d_rowptr = dense_AT(rowptr)
     d_colind = dense_AT(colind)
-    d_nzval  = dense_AT(nzval)
+    d_nzval = dense_AT(nzval)
     return GPUSparseMatrixCSR(d_rowptr, d_colind, d_nzval, (m, n))
 end
 function direct_vector_construction(AT::Type{<:GPUArrays.AbstractGPUSparseMatrix}, eltypes)
@@ -205,6 +205,7 @@ function direct_vector_construction(AT::Type{<:GPUArrays.AbstractGPUSparseMatrix
         @test x isa AT{ET}
         @test size(x) == (m, n)
     end
+    return
 end
 function direct_vector_construction(AT, eltypes)
     # NOP

[kshyatt]

Sorry for neglecting this, I'll try to take a look today

[kshyatt]

So one question I have is whether COO is really the best default format, or whether we should allow users/developers to set the format somehow with a default_sparse_format reference value or something. For example for CUSPARSE, CSR is the preferred format.

[nHackel]

@kshyatt that's a good question. From a user perspective, I think a changeable default matrix would make using the default constructor difficult unless the user can overwrite the default. Otherwise since the inputs are interpreted differently between the formats, a user would need to check the default format and construct the appropriate inputs. At which point they could also just call the appropriate "direct" constructor anyways.

So maybe we have a default of:

function GPUSparseMatrix(args...; kwargs..., format = :coo)
	if format == :coo
		GPUSparseMatrixCOO(args...; kwargs...)
	elseif format == :csc
	# ...
	else
		throw(ArgumentError("Unexpected format specifier $format"))
	end
end

# Optionally in this setting CUDA could overwrite like this 
GPUSparseMatrix(i::CuVector, args...; kwargs..., format = :csc) = ...

or with a default option from a backend via a function:

function GPUSparseMatrix(gpu::GPUArray, args...; kwargs..., format = default_sparse_format(typeof(GPU)))
	if format == :coo
		GPUSparseMatrixCOO(args...; kwargs...)
	elseif format == :csc
	# ...
	else
		throw(ArgumentError("Unexpected format specifier $format"))
	end
end

or we just drop the GPUSparseArray all together, since the current default is just sparse(...) anyway and only offer GPUSparseMatrixXYZ

[kshyatt]

I think default_sparse_format(typeof(GPU)) could be great, maybe it should dispatch on a subtype of GPUArray rather than typeof(GPU)?

[nHackel]

Should a default constructor always expect COO/normal sparse method inputs and know how to translate them:

function GPUSparseMatrix(I, J, V; format = default_sparse_format(I))
	if format == :coo
		GPUSparseMatrixCOO(I, J, V)
	elseif format == :csc
		# calculate colptr and use in function call
		GPUSparseMatrixCSC(...)
	elseif format == :csr
		# calculate rowptr and use in function call
		GPUSparseMatrixCSR(...)
	# ...
	end
end

or should it just pass along arguments as in the version in my previous comment? I think CUDA has a similar interface for sparse(I, J, V)

[kshyatt]

Well I guess we could also pass a Val type so dispatch could handle this more effectively? Leave the keyword arg version as an outer wrapper just in case?

[nHackel]

I can do that, my main worry/question is what we should do with the inputs for the default setting.

If we just pass along the arguments, a user can run into this issue:

colptr = ...
rowval = ...
vals = ...

GPUSparseMatrix(colptr, rowval, vals; format = :csc) # result should be as expected by user

# Would most likely result in transpose of desired matrix
# Since CSR expects arguments (rowptr, colval, vals)
# In this case the user chose that, in the case of a backend specific default, the user doesn't necessarily
# know the default backend
GPUSparseMatrix(colptr, rowval, vals; format = :csr)

If we always expect arguments to be in COO format, we could try to compute the approriate indices:

GPUSparseMatrix(I, J, V; format = default_sparse_format(I)) = GPUSparseMatrix(format, I, J, V)
GPUSparseMatrix(format::Symbol, args...) = GPUSparseMatrix(Val(format), args...)

function GPUSparseMatrix(::Val{:csc}, I, J, V)
	colptr = # based on I
	return GPUSparseMatrixCSC(colptr, J, V)
end

function GPUSparseMatrix(::Val{:csr}, I, J, V)
	rowptr = # based on J
	return GPUSparseMatrixCSC(rowptr, I, V)
end

or we drop the GPUSparseMatrix and just provde GPUSparseMatrixCOO, GPUSparseMatrixCSC, GPUSparseMatrixCSR and GPUSparseMatrixBSR

[kshyatt]

I agree users could be tricked here... but on the other hand at some point they have to use the constructor responsibly. It's also the case that they can inadvertently create a transposed matrix with GPUSparseMatrixCSR by passing args in the wrong order.

[kshyatt]

Maybe the default could indeed be COO but then if the user passes a Val type bypassing this interface they're explicitly opting into a format?

[kshyatt]

GPUSparseMatrix(I, J, V) = GPUSparseMatrix(Val(:coo), I, J, V)

Then
GPUSparseMatrix(::Val{:csr}, rowPtr, colVal, nzVal) = GPUSparseMatrixCSR(rowPtr, colVal, nzVal)

etc?

[nHackel]

That sounds like a good middle-ground between the options! I'll try that out as soon as I can