Merge pull request #4 from rasmushenningsson/mul_fallbacks

Better mul! fallbacks

Author: rasmushenningsson

src/ThreadedSparseArrays.jl

@@ -53,6 +53,14 @@ for f in [:rowvals, :nonzeros, :getcolptr]
     @eval SparseArrays.$(f)(A::ThreadedSparseMatrixCSC) = SparseArrays.$(f)(A.A)
 end
 
+# For non-threaded implementations, fallback to sparse methods and not generic matmul.
+mul!(C::AbstractVector, A::ThreadedSparseMatrixCSC, B::AbstractVector, α::Number, β::Number) = mul!(C, A.A, B, α, β)
+mul!(C::AbstractMatrix, A::ThreadedSparseMatrixCSC, B::AbstractMatrix, α::Number, β::Number) = mul!(C, A.A, B, α, β)
+mul!(C::AbstractVector, A::Adjoint{<:Any,<:ThreadedSparseMatrixCSC}, B::AbstractVector, α::Number, β::Number) = mul!(C, adjoint(A.parent.A), B, α, β)
+mul!(C::AbstractMatrix, A::Adjoint{<:Any,<:ThreadedSparseMatrixCSC}, B::AbstractMatrix, α::Number, β::Number) = mul!(C, adjoint(A.parent.A), B, α, β)
+mul!(C::AbstractVector, A::Transpose{<:Any,<:ThreadedSparseMatrixCSC}, B::AbstractVector, α::Number, β::Number) = mul!(C, transpose(A.parent.A), B, α, β)
+mul!(C::AbstractMatrix, A::Transpose{<:Any,<:ThreadedSparseMatrixCSC}, B::AbstractMatrix, α::Number, β::Number) = mul!(C, transpose(A.parent.A), B, α, β)
+
 function mul!(C::StridedVecOrMat, A::ThreadedSparseMatrixCSC, B::Union{StridedVector,AdjOrTransDenseMatrix}, α::Number, β::Number)
     size(A, 2) == size(B, 1) || throw(DimensionMismatch())
     size(A, 1) == size(C, 1) || throw(DimensionMismatch())

test/runtests.jl

@@ -49,4 +49,51 @@ using Test
         @test norm(ref-out) == 0
     end
 
+
+    # These test below are here to ensure we don't hit ambiguity warnings.
+    # The implementations are not (currently) threaded.
+    sx = sprand(Bool,n,0.05)
+    @testset "$(Mat)_L_sparsevec" for Mat in [ThreadedSparseMatrixCSC]
+        Ct = Mat(C)
+
+        out = similar(sx, T, N)
+        LinearAlgebra.mul!(out, Ct, sx)
+        ref = C*sx
+        @test norm(ref-out) == 0
+        @test typeof(ref)==typeof(out)
+    end
+
+    sx = sprand(Bool,N,0.05)
+    @testset "$(Mat)_L_$(op)_sparsevec" for op in [adjoint,transpose], Mat in [ThreadedSparseMatrixCSC]
+        Ct = Mat(C)
+
+        out = similar(sx, T, n)
+        LinearAlgebra.mul!(out, op(Ct), sx)
+        ref = op(C)*sx
+        @test norm(ref-out) == 0
+        @test typeof(ref)==typeof(out)
+    end
+
+    sx = sparse(rand(1:n,10),1:10,true,n,10)
+    @testset "$(Mat)_L_sparse" for Mat in [ThreadedSparseMatrixCSC]
+        Ct = Mat(C)
+
+        out = similar(sx, T, N, 10)
+        LinearAlgebra.mul!(out, Ct, sx)
+        ref = C*sx
+        @test norm(ref-out) == 0
+        @test typeof(ref)==typeof(out)
+    end
+
+    sx = sparse(rand(1:N,10),1:10,true,N,10)
+    @testset "$(Mat)_L_$(op)_sparse" for op in [adjoint,transpose], Mat in [ThreadedSparseMatrixCSC]
+        Ct = Mat(C)
+
+        out = similar(sx, T, n, 10)
+        LinearAlgebra.mul!(out, op(Ct), sx)
+        ref = op(C)*sx
+        @test norm(ref-out) == 0
+        @test typeof(ref)==typeof(out)
+    end
+
 end