Initial commit

Author: jagot

Project.toml

@@ -3,6 +3,10 @@ uuid = "59d54670-b8ac-4d81-ab7a-bb56233e17ab"
 authors = ["Stefanos Carlström <[email protected]>"]
 version = "0.1.0"
 
+[deps]
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
 [compat]
 julia = "1"
 

README.md

@@ -3,3 +3,7 @@
 [![Build Status](https://travis-ci.com/jagot/ThreadedSparseArrays.jl.svg?branch=master)](https://travis-ci.com/jagot/ThreadedSparseArrays.jl)
 [![Build Status](https://ci.appveyor.com/api/projects/status/github/jagot/ThreadedSparseArrays.jl?svg=true)](https://ci.appveyor.com/project/jagot/ThreadedSparseArrays-jl)
 [![Codecov](https://codecov.io/gh/jagot/ThreadedSparseArrays.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/jagot/ThreadedSparseArrays.jl)
+
+Simple package providing a wrapper type enabling threaded sparse
+matrix–dense matrix multiplication. Based on [this
+PR](https://github.com/JuliaLang/julia/pull/29525).

src/ThreadedSparseArrays.jl

@@ -1,5 +1,92 @@
 module ThreadedSparseArrays
 
-greet() = print("Hello World!")
+using LinearAlgebra
+import LinearAlgebra: mul!
+using SparseArrays
+import SparseArrays: AdjOrTransStridedOrTriangularMatrix, getcolptr
+
+struct ThreadedSparseMatrixCSC{Tv,Ti,At} <: AbstractSparseMatrix{Tv,Ti}
+    A::At
+    ThreadedSparseMatrixCSC(A::At) where {Tv,Ti,At<:AbstractSparseMatrix{Tv,Ti}} =
+        new{Tv,Ti,At}(A)
+end
+
+Base.size(A::ThreadedSparseMatrixCSC, args...) = size(A.A, args...)
+Base.eltype(A::ThreadedSparseMatrixCSC) = eltype(A.A)
+Base.getindex(A::ThreadedSparseMatrixCSC, args...) = getindex(A.A, args...)
+
+# Need to override printing
+# Need to forward findnz, etc
+
+for f in [:rowvals, :nonzeros, :getcolptr]
+    @eval SparseArrays.$(f)(A::ThreadedSparseMatrixCSC) = SparseArrays.$(f)(A.A)
+end
+
+function mul!(C::StridedVecOrMat, adjA::Adjoint{<:Any,<:ThreadedSparseMatrixCSC}, B::Union{StridedVector,AdjOrTransStridedOrTriangularMatrix}, α::Number, β::Number)
+    A = adjA.parent
+    size(A, 2) == size(C, 1) || throw(DimensionMismatch())
+    size(A, 1) == size(B, 1) || throw(DimensionMismatch())
+    size(B, 2) == size(C, 2) || throw(DimensionMismatch())
+    colptrA = getcolptr(A)
+    nzv = nonzeros(A)
+    rv = rowvals(A)
+    if β != 1
+        β != 0 ? rmul!(C, β) : fill!(C, zero(eltype(C)))
+    end
+    for k = 1:size(C, 2)
+        Threads.@threads for col = 1:size(A, 2)
+            @inbounds begin
+                tmp = zero(eltype(C))
+                for j = colptrA[col]:(colptrA[col+1] - 1)
+                    tmp += adjoint(nzv[j])*B[rv[j],k]
+                end
+                C[col,k] += α*tmp
+            end
+        end
+    end
+    C
+end
+
+function mul!(C::StridedVecOrMat, transA::Transpose{<:Any,<:ThreadedSparseMatrixCSC}, B::Union{StridedVector,AdjOrTransStridedOrTriangularMatrix}, α::Number, β::Number)
+    A = transA.parent
+    size(A, 2) == size(C, 1) || throw(DimensionMismatch())
+    size(A, 1) == size(B, 1) || throw(DimensionMismatch())
+    size(B, 2) == size(C, 2) || throw(DimensionMismatch())
+    nzv = nonzeros(A)
+    rv = rowvals(A)
+    if β != 1
+        β != 0 ? rmul!(C, β) : fill!(C, zero(eltype(C)))
+    end
+    Threads.@threads for k = 1:size(C, 2)
+        @inbounds for col = 1:size(A, 2)
+            tmp = zero(eltype(C))
+            for j = getcolptr(A)[col]:(getcolptr(A)[col + 1] - 1)
+                tmp += transpose(nzv[j])*B[rv[j],k]
+            end
+            C[col,k] += tmp * α
+        end
+    end
+    C
+end
+
+function mul!(C::StridedVecOrMat, X::AdjOrTransStridedOrTriangularMatrix, A::ThreadedSparseMatrixCSC, α::Number, β::Number)
+    mX, nX = size(X)
+    nX == size(A, 1) || throw(DimensionMismatch())
+    mX == size(C, 1) || throw(DimensionMismatch())
+    size(A, 2) == size(C, 2) || throw(DimensionMismatch())
+    rv = rowvals(A)
+    nzv = nonzeros(A)
+    if β != 1
+        β != 0 ? rmul!(C, β) : fill!(C, zero(eltype(C)))
+    end
+    Threads.@threads for col = 1:size(A, 2)
+        @inbounds for multivec_row=1:mX, k=getcolptr(A)[col]:(getcolptr(A)[col+1]-1)
+            C[multivec_row, col] += α * X[multivec_row, rv[k]] * nzv[k] # perhaps suboptimal position of α?
+        end
+    end
+    C
+end
+
+export ThreadedSparseMatrixCSC
 
 end # module

test/runtests.jl

@@ -1,6 +1,19 @@
 using ThreadedSparseArrays
+using SparseArrays
+using LinearAlgebra
 using Test
 
 @testset "ThreadedSparseArrays.jl" begin
-    # Write your own tests here.
+    M,N = 5000,4000
+    n = 200
+    T = ComplexF64
+
+    C = sprand(T, N, n, 0.05)
+    Ct = ThreadedSparseMatrixCSC(C)
+
+    eye = Matrix(one(T)*I, N, N)
+    out = zeros(T, N, n)
+    LinearAlgebra.mul!(out, eye, Ct)
+    ref = eye*C
+    @test norm(ref-out) == 0
 end