Specialized reduction of sparse arrays

Internals & Design
1

Hi All,

I have created a specialized version of reduction for sparse arrays, as the default falls back to generic arrays, which is too slow.

The code is pretty simple as

function Base.reduce(::typeof(hcat), A::AbstractVector{T}) where {T<: SparseMatrixCSC{Tv, Ti} where {Tv, Ti}}
  length(A) <= 1 && return(A)
  Tv, Ti = eltype(A[1].nzval), eltype(A[1].colptr)
  coloffset = 0
  Is, Js, Vs = Vector{Vector{Ti}}(undef, length(A)), Vector{Vector{Ti}}(undef, length(A)), Vector{Vector{Tv}}(undef, length(A))
  for (i, a) in enumerate(A)
    (I, J, V) = findnz(a, coloffset)
    Is[i] =  I
    Js[i] =  J
    Vs[i] =  V
    coloffset += a.n
  end
  sparse(reduce(vcat, Is), reduce(vcat, Js), reduce(vcat, Vs))
end

where I have slightly modified the default findnz function

function SparseArrays.findnz(S::SparseMatrixCSC{Tv,Ti}, coloffset) where {Tv,Ti}
    numnz = nnz(S)
    I = Vector{Ti}(undef, numnz)
    J = Vector{Ti}(undef, numnz)
    V = Vector{Tv}(undef, numnz)

    count = 1
    @inbounds for col = 1 : S.n, k = S.colptr[col] : (S.colptr[col+1]-1)
        I[count] = S.rowval[k]
        J[count] = col + coloffset
        V[count] = S.nzval[k]
        count += 1
    end

    return (I, J, V)
end

My question is, if it is of sufficient interest to be pushed to julia’s Sparse Arrays package. Also, if it possible to controbute only to SparseArrays and not to the whole Julia base.

Thanks for answer.

1 Like