Clean up some multiplication code (#388)

Author: dkarrasch

src/SparseArrays.jl

@@ -10,7 +10,7 @@ using Base: ReshapedArray, promote_op, setindex_shape_check, to_shape, tail,
 using Base.Order: Forward
 using LinearAlgebra
 using LinearAlgebra: AdjOrTrans, AdjointFactorization, TransposeFactorization, matprod,
-    AbstractQ, AdjointQ, HessenbergQ, QRCompactWYQ, QRPackedQ, LQPackedQ,
+    AbstractQ, AdjointQ, HessenbergQ, QRCompactWYQ, QRPackedQ, LQPackedQ, MulAddMul,
     UpperOrLowerTriangular
 
 

src/linalg.jl

@@ -95,7 +95,7 @@ _sparsem(A::AbstractSparseMatrix) = A
 _sparsem(A::AbstractSparseVector) = A
 
 # Transpose/Adjoint of sparse vector (returning sparse matrix)
-function _sparsem(A::Union{Transpose{<:Any,<:AbstractSparseVector},Adjoint{<:Any,<:AbstractSparseVector}})
+function _sparsem(A::AdjOrTrans{<:Any,<:AbstractSparseVector})
     B = parent(A)
     n = length(B)
     Ti = eltype(nonzeroinds(B))
@@ -217,8 +217,7 @@ function _sparsem(A::AbstractTriangularSparse{Tv}) where Tv
 end
 
 # 8 cases: (Transpose|Adjoint){Tv,[Unit](Upper|Lower)Triangular}
-function _sparsem(taA::Union{Transpose{Tv,<:AbstractTriangularSparse},
-                             Adjoint{Tv,<:AbstractTriangularSparse}}) where {Tv}
+function _sparsem(taA::AdjOrTrans{Tv,<:AbstractTriangularSparse}) where {Tv}
 
     sA = taA.parent
     A = sA.data

src/sparseconvert.jl

@@ -209,8 +209,7 @@ julia> nnz(A)
 """
 nnz(S::AbstractSparseMatrixCSC) = Int(getcolptr(S)[size(S, 2) + 1]) - 1
 nnz(S::ReshapedArray{<:Any,1,<:AbstractSparseMatrixCSC}) = nnz(parent(S))
-nnz(S::Adjoint{<:Any,<:AbstractSparseMatrixCSC}) = nnz(parent(S))
-nnz(S::Transpose{<:Any,<:AbstractSparseMatrixCSC}) = nnz(parent(S))
+nnz(S::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}) = nnz(parent(S))
 nnz(S::UpperTriangular{<:Any,<:AbstractSparseMatrixCSC}) = nnz1(S)
 nnz(S::LowerTriangular{<:Any,<:AbstractSparseMatrixCSC}) = nnz1(S)
 nnz(S::SparseMatrixCSCView) = nnz1(S)

src/sparsematrix.jl

@@ -4,6 +4,7 @@
 
 import Base: sort!, findall, copy!
 import LinearAlgebra: promote_to_array_type, promote_to_arrays_
+using LinearAlgebra: adj_or_trans
 
 ### The SparseVector
 
@@ -1768,23 +1769,52 @@ end
 
 const _StridedOrTriangularMatrix{T} = Union{StridedMatrix{T}, LowerTriangular{T}, UnitLowerTriangular{T}, UpperTriangular{T}, UnitUpperTriangular{T}}
 
+_fliptri(A::UpperTriangular) = LowerTriangular(parent(parent(A)))
+_fliptri(A::UnitUpperTriangular) = UnitLowerTriangular(parent(parent(A)))
+_fliptri(A::LowerTriangular) = UpperTriangular(parent(parent(A)))
+_fliptri(A::UnitLowerTriangular) = UnitUpperTriangular(parent(parent(A)))
+
 function (*)(A::_StridedOrTriangularMatrix{Ta}, x::AbstractSparseVector{Tx}) where {Ta,Tx}
     require_one_based_indexing(A, x)
     m, n = size(A)
     length(x) == n || throw(DimensionMismatch())
     Ty = promote_op(matprod, eltype(A), eltype(x))
     y = Vector{Ty}(undef, m)
-    mul!(y, A, x)
+    mul!(y, A, x, true, false)
 end
 
-function mul!(y::AbstractVector, A::_StridedOrTriangularMatrix, x::AbstractSparseVector, α::Number, β::Number)
+function LinearAlgebra.generic_matvecmul!(y::AbstractVector, tA, A::StridedMatrix, x::AbstractSparseVector,
+                                            _add::MulAddMul = MulAddMul())
+    if tA == 'N'
+        _spmul!(y, A, x, _add.alpha, _add.beta)
+    elseif tA == 'T'
+        _At_or_Ac_mul_B!(transpose, y, A, x, _add.alpha, _add.beta)
+    elseif tA == 'C'
+        _At_or_Ac_mul_B!(adjoint, y, A, x, _add.alpha, _add.beta)
+    else
+        _spmul!(y, LinearAlgebra.wrap(A, tA), x, _add.alpha, _add.beta)
+    end
+    return y
+end
+function LinearAlgebra.generic_matvecmul!(y::AbstractVector, tA, A::UpperOrLowerTriangular, x::AbstractSparseVector,
+                                             _add::MulAddMul = MulAddMul())
+    @assert tA == 'N'
+    Adata = parent(A)
+    if Adata isa Transpose
+        _At_or_Ac_mul_B!(transpose, y, _fliptri(A), x, _add.alpha, _add.beta)
+    elseif Adata isa Adjoint
+        _At_or_Ac_mul_B!(adjoint, y, _fliptri(A), x, _add.alpha, _add.beta)
+    else # Adata is plain
+        _spmul!(y, A, x, _add.alpha, _add.beta)
+    end
+    return y
+end
+function _spmul!(y::AbstractVector, A::AbstractMatrix, x::AbstractSparseVector, α::Number, β::Number)
     require_one_based_indexing(y, A, x)
     m, n = size(A)
     length(x) == n && length(y) == m || throw(DimensionMismatch())
     m == 0 && return y
-    if β != one(β)
-        β == zero(β) ? fill!(y, zero(eltype(y))) : rmul!(y, β)
-    end
+    β != one(β) && LinearAlgebra._rmul_or_fill!(y, β)
     α == zero(α) && return y
 
     xnzind = nonzeroinds(x)
@@ -1802,80 +1832,47 @@ function mul!(y::AbstractVector, A::_StridedOrTriangularMatrix, x::AbstractSpars
     return y
 end
 
-# * and mul!(C, transpose(A), B)
-
-function *(tA::Transpose{<:Any,<:_StridedOrTriangularMatrix{Ta}}, x::AbstractSparseVector{Tx}) where {Ta,Tx}
-    require_one_based_indexing(tA, x)
-    m, n = size(tA)
-    length(x) == n || throw(DimensionMismatch())
-    Ty = promote_op(matprod, eltype(tA), eltype(x))
-    y = Vector{Ty}(undef, m)
-    mul!(y, tA, x)
-end
-
-function mul!(y::AbstractVector, tA::Transpose{<:Any,<:_StridedOrTriangularMatrix}, x::AbstractSparseVector, α::Number, β::Number)
-    require_one_based_indexing(y, tA, x)
-    m, n = size(tA)
+function _At_or_Ac_mul_B!(tfun::Function,
+                            y::AbstractVector, A::_StridedOrTriangularMatrix, x::AbstractSparseVector,
+                            α::Number, β::Number)
+    require_one_based_indexing(y, A, x)
+    n, m = size(A)
     length(x) == n && length(y) == m || throw(DimensionMismatch())
     m == 0 && return y
-    if β != one(β)
-        β == zero(β) ? fill!(y, zero(eltype(y))) : rmul!(y, β)
-    end
+    β != one(β) && LinearAlgebra._rmul_or_fill!(y, β)
     α == zero(α) && return y
 
     xnzind = nonzeroinds(x)
     xnzval = nonzeros(x)
     _nnz = length(xnzind)
     _nnz == 0 && return y
 
-    A = tA.parent
     Ty = promote_op(matprod, eltype(A), eltype(x))
     @inbounds for j = 1:m
         s = zero(Ty)
         for i = 1:_nnz
-            s += transpose(A[xnzind[i], j]) * xnzval[i]
+            s += tfun(A[xnzind[i], j]) * xnzval[i]
         end
         y[j] += s * α
     end
     return y
 end
 
-# * and mul!(C, adjoint(A), B)
-
-function *(adjA::Adjoint{<:Any,<:_StridedOrTriangularMatrix{Ta}}, x::AbstractSparseVector{Tx}) where {Ta,Tx}
-    require_one_based_indexing(adjA, x)
-    m, n = size(adjA)
+function *(A::AdjOrTrans{<:Any,<:StridedMatrix}, x::AbstractSparseVector)
+    require_one_based_indexing(A, x)
+    m, n = size(A)
     length(x) == n || throw(DimensionMismatch())
-    Ty = promote_op(matprod, eltype(adjA), eltype(x))
+    Ty = promote_op(matprod, eltype(A), eltype(x))
     y = Vector{Ty}(undef, m)
-    mul!(y, adjA, x)
+    mul!(y, A, x, true, false)
 end
-
-function mul!(y::AbstractVector, adjA::Adjoint{<:Any,<:_StridedOrTriangularMatrix}, x::AbstractSparseVector, α::Number, β::Number)
-    require_one_based_indexing(y, adjA, x)
-    m, n = size(adjA)
-    length(x) == n && length(y) == m || throw(DimensionMismatch())
-    m == 0 && return y
-    if β != one(β)
-        β == zero(β) ? fill!(y, zero(eltype(y))) : rmul!(y, β)
-    end
-    α == zero(α) && return y
-
-    xnzind = nonzeroinds(x)
-    xnzval = nonzeros(x)
-    _nnz = length(xnzind)
-    _nnz == 0 && return y
-
-    A = adjA.parent
+function *(A::LinearAlgebra.HermOrSym{<:Any,<:StridedMatrix}, x::AbstractSparseVector)
+    require_one_based_indexing(A, x)
+    m, n = size(A)
+    length(x) == n || throw(DimensionMismatch())
     Ty = promote_op(matprod, eltype(A), eltype(x))
-    @inbounds for j = 1:m
-        s = zero(Ty)
-        for i = 1:_nnz
-            s += adjoint(A[xnzind[i], j]) * xnzval[i]
-        end
-        y[j] += s * α
-    end
-    return y
+    y = Vector{Ty}(undef, m)
+    mul!(y, A, x, true, false)
 end
 
 
@@ -1906,15 +1903,26 @@ function densemv(A::AbstractSparseMatrixCSC, x::AbstractSparseVector; trans::Abs
 end
 
 # * and mul!
+function LinearAlgebra.generic_matvecmul!(y::AbstractVector, tA, A::AbstractSparseMatrixCSC, x::AbstractSparseVector,
+                            _add::MulAddMul = MulAddMul())
+    if tA == 'N'
+        _spmul!(y, A, x, _add.alpha, _add.beta)
+    elseif tA == 'T'
+        _At_or_Ac_mul_B!((a,b) -> transpose(a) * b, y, A, x, _add.alpha, _add.beta)
+    elseif tA == 'C'
+        _At_or_Ac_mul_B!((a,b) -> adjoint(a) * b, y, A, x, _add.alpha, _add.beta)
+    else
+        LinearAlgebra._generic_matvecmul!(y, 'N', LinearAlgebra.wrap(A, tA), x, _add)
+    end
+    return y
+end
 
-function mul!(y::AbstractVector, A::AbstractSparseMatrixCSC, x::AbstractSparseVector, α::Number, β::Number)
+function _spmul!(y::AbstractVector, A::AbstractSparseMatrixCSC, x::AbstractSparseVector, α::Number, β::Number)
     require_one_based_indexing(y, A, x)
     m, n = size(A)
     length(x) == n && length(y) == m || throw(DimensionMismatch())
     m == 0 && return y
-    if β != one(β)
-        β == zero(β) ? fill!(y, zero(eltype(y))) : rmul!(y, β)
-    end
+    β != one(β) && LinearAlgebra._rmul_or_fill!(y, β)
     α == zero(α) && return y
 
     xnzind = nonzeroinds(x)
@@ -1936,23 +1944,14 @@ function mul!(y::AbstractVector, A::AbstractSparseMatrixCSC, x::AbstractSparseVe
     return y
 end
 
-# * and *(Transpose(A), B)
-mul!(y::AbstractVector, tA::Transpose{<:Any,<:AbstractSparseMatrixCSC}, x::AbstractSparseVector, α::Number, β::Number) =
-    _At_or_Ac_mul_B!((a,b) -> transpose(a) * b, y, tA.parent, x, α, β)
-
-mul!(y::AbstractVector, adjA::Adjoint{<:Any,<:AbstractSparseMatrixCSC}, x::AbstractSparseVector, α::Number, β::Number) =
-    _At_or_Ac_mul_B!((a,b) -> adjoint(a) * b, y, adjA.parent, x, α, β)
-
 function _At_or_Ac_mul_B!(tfun::Function,
                           y::AbstractVector, A::AbstractSparseMatrixCSC, x::AbstractSparseVector,
                           α::Number, β::Number)
     require_one_based_indexing(y, A, x)
     m, n = size(A)
     length(x) == m && length(y) == n || throw(DimensionMismatch())
     n == 0 && return y
-    if β != one(β)
-        β == zero(β) ? fill!(y, zero(eltype(y))) : rmul!(y, β)
-    end
+    β != one(β) && LinearAlgebra._rmul_or_fill!(y, β)
     α == zero(α) && return y
 
     xnzind = nonzeroinds(x)
@@ -1981,11 +1980,8 @@ function *(A::AbstractSparseMatrixCSC, x::AbstractSparseVector)
     _dense2sparsevec(y, initcap)
 end
 
-*(tA::Transpose{<:Any,<:AbstractSparseMatrixCSC}, x::AbstractSparseVector) =
-    _At_or_Ac_mul_B((a,b) -> transpose(a) * b, tA.parent, x, promote_op(matprod, eltype(tA), eltype(x)))
-
-*(adjA::Adjoint{<:Any,<:AbstractSparseMatrixCSC}, x::AbstractSparseVector) =
-    _At_or_Ac_mul_B((a,b) -> adjoint(a) * b, adjA.parent, x, promote_op(matprod, eltype(adjA), eltype(x)))
+*(xA::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, x::AbstractSparseVector) =
+    _At_or_Ac_mul_B((a,b) -> adj_or_trans(xA)(a) * b, xA.parent, x, promote_op(matprod, eltype(xA), eltype(x)))
 
 function _At_or_Ac_mul_B(tfun::Function, A::AbstractSparseMatrixCSC{TvA,TiA}, x::AbstractSparseVector{TvX,TiX},
                          Tv = promote_op(matprod, TvA, TvX)) where {TvA,TiA,TvX,TiX}

src/sparsevector.jl

@@ -190,7 +190,7 @@ end
     @testset "symmetric/Hermitian sparse multiply with $S($U)" for S in (Symmetric, Hermitian), U in (:U, :L), (A, B) in ((Areal,Breal), (Acomplex,Bcomplex))
         Asym = S(A, U)
         As = sparse(Asym) # takes most time
-        @test which(mul!, (typeof(B), typeof(Asym), typeof(B))).module == SparseArrays
+        # @test which(mul!, (typeof(B), typeof(Asym), typeof(B))).module == SparseArrays
         @test norm(Asym * B - As * B, Inf) <= eps() * n * p * 10
     end
 end
@@ -207,7 +207,7 @@ end
     @testset "symmetric/Hermitian sparseview multiply with $S($U)" for S in (Symmetric, Hermitian), U in (:U, :L), (A, B) in ((Areal,Breal), (Acomplex,Bcomplex))
         Asym = S(A, U)
         As = sparse(Asym) # takes most time
-        @test which(mul!, (typeof(B), typeof(Asym), typeof(B))).module == SparseArrays
+        # @test which(mul!, (typeof(B), typeof(Asym), typeof(B))).module == SparseArrays
         @test norm(Asym * B - As * B, Inf) <= eps() * n * p * 10
     end
 end
@@ -662,11 +662,14 @@ end
         @test Array(f*b) == f*Array(b)
         A = rand(2n, 2n)
         sA = view(A, 1:2:2n, 1:2:2n)
-        @test Array(sA*b) ≈ Array(sA)*Array(b)
-        @test Array(a*sA) ≈ Array(a)*Array(sA)
+        @test Array((sA*b)::Matrix) ≈ Array(sA)*Array(b)
+        @test Array((a*sA)::Matrix) ≈ Array(a)*Array(sA)
+        @test Array((sA'b)::Matrix) ≈ Array(sA')*Array(b)
         c = sprandn(ComplexF32, n, n, q)
-        @test Array(sA*c') ≈ Array(sA)*Array(c)'
-        @test Array(c'*sA) ≈ Array(c)'*Array(sA)
+        @test Array((sA*c')::Matrix) ≈ Array(sA)*Array(c)'
+        @test Array((c'*sA)::Matrix) ≈ Array(c)'*Array(sA)
+        @test Array((sA'c)::Matrix) ≈ Array(sA')*Array(c)
+        @test Array((sA'c')::Matrix) ≈ Array(sA')*Array(c)'
     end
 end
 

test/linalg.jl

@@ -175,16 +175,15 @@ dA = Array(sA)
     p28227 = sparse(Real[0 0.5])
 
     for arr in (se33, sA, pA, p28227, spzeros(3, 3))
+        farr = Array(arr)
         for f in (sum, prod, minimum, maximum)
-            farr = Array(arr)
             @test f(arr) ≈ f(farr)
             @test f(arr, dims=1) ≈ f(farr, dims=1)
             @test f(arr, dims=2) ≈ f(farr, dims=2)
             @test f(arr, dims=(1, 2)) ≈ [f(farr)]
             @test isequal(f(arr, dims=3), f(farr, dims=3))
         end
         for f in (+, *, min, max)
-            farr = Array(arr)
             @test mapreduce(identity, f, arr) ≈ mapreduce(identity, f, farr)
             @test mapreduce(x -> x + 1, f, arr) ≈ mapreduce(x -> x + 1, f, farr)
         end

test/sparsematrix_ops.jl

@@ -1098,6 +1098,22 @@ end
                 @test isa(y, Vector{T})
                 @test y ≈ *(adjoint(A), xf)
             end
+
+            let A = randn(TA, 16, 16), x = sprand(Tx, 16, 0.7)
+                xf = Array(x)
+                for wrap in (M -> Symmetric(M, :U), M -> Symmetric(M, :L),
+                        M -> Hermitian(M, :U), M -> Hermitian(M, :L))
+                    for α in (0.0, 1.0, 2.0), β in (0.0, 0.5, 1.0)
+                        y = rand(T, 16)
+                        rr = α*wrap(A)*xf + β*y
+                        @test mul!(y, wrap(A), x, α, β) === y
+                        @test y ≈ rr
+                    end
+                    y = *(wrap(A), x)
+                    @test isa(y, Vector{T})
+                    @test y ≈ *(wrap(A), xf)
+                end
+            end
         end
     end
     @testset "sparse A * sparse x -> dense y" begin
@@ -1129,6 +1145,29 @@ end
             @test y ≈ *(transpose(Af), xf)
         end
 
+        let A = sprandn(16, 16, 0.5), x = sprand(16, 0.7)
+            Af = Array(A)
+            xf = Array(x)
+            for wrap in (M -> Symmetric(M, :U), M -> Symmetric(M, :L),
+                M -> Hermitian(M, :U), M -> Hermitian(M, :L),
+                M -> UpperTriangular(M), M -> UnitUpperTriangular(M),
+                M -> LowerTriangular(M), M -> UnitLowerTriangular(M),
+                M -> UpperTriangular(transpose(M)), M -> UnitUpperTriangular(transpose(M)),
+                M -> LowerTriangular(transpose(M)), M -> UnitLowerTriangular(transpose(M)),
+                M -> UpperTriangular(adjoint(M)), M -> UnitUpperTriangular(adjoint(M)),
+                M -> LowerTriangular(adjoint(M)), M -> UnitLowerTriangular(adjoint(M)),
+                M -> UpperTriangular(Symmetric(M)))
+                for α in (0.0, 1.0, 2.0), β in (0.0, 0.5, 1.0)
+                    y = rand(16)
+                    rr = α*wrap(Af)*xf + β*y
+                    @test mul!(y, wrap(A), x, α, β) === y
+                    @test y ≈ rr
+                end
+                y = wrap(A) * x
+                @test y ≈ *(wrap(Af), xf)
+            end
+        end
+
         let A = complex.(sprandn(7, 8, 0.5), sprandn(7, 8, 0.5)),
             x = complex.(sprandn(8, 0.6), sprandn(8, 0.6)),
             x2 = complex.(sprandn(7, 0.75), sprandn(7, 0.75))