Specialize Diagonal products and dot for adjoint/transpose of sparse matrices

src/linalg.jl

@@ -56,12 +56,28 @@ matop_dest(::typeof(*), A::QuasiStridedMatrix, b::AbstractSparseVector) =
     Vector{promote_op(matprod, eltype(A), eltype(b))}(undef, size(A, 1))
 matop_dest(::typeof(*), A, B::QuasiSparseMatrix) =
     similar(A, promote_op(matprod, eltype(A), eltype(B)), (size(A, 1), size(B, 2)))
-# sparse products with banded matrices should return sparse arrays (Diagonal is handled by fallback)
+# sparse products with banded matrices should return sparse arrays
 matop_dest(::typeof(*), A::BiTriSym, B::QuasiSparseMatrix) =
     similar(B, promote_op(matprod, eltype(A), eltype(B)), size(B))
 # needed for disambiguation with LinearAlgebra
 matop_dest(::typeof(*), A::Diagonal, B::QuasiSparseMatrix) =
     similar(B, promote_op(matprod, eltype(A), eltype(B)), size(B))
+# a `Diagonal` product keeps the structure of the sparse operand, so a fixed operand gets
+# a fixed destination with that structure up front, which `mul!` then only has to fill
+# (an empty fixed destination could not take the indices); the adjoint/transpose of a
+# sparse matrix gets an empty, writable destination, since its structure is not that of
+# the parent
+matop_dest(::typeof(*), A::Diagonal, B::AbstractSparseMatrixCSC) =
+    _is_fixed(B) ? similar(B, promote_op(matprod, eltype(A), eltype(B))) :
+                   similar(B, promote_op(matprod, eltype(A), eltype(B)), size(B))
+matop_dest(::typeof(*), A::Diagonal, B::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}) =
+    _adjtrans_dest(B, promote_op(matprod, eltype(A), eltype(B)))
+matop_dest(::typeof(*), A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, B::Diagonal) =
+    _adjtrans_dest(A, promote_op(matprod, eltype(A), eltype(B)))
+function _adjtrans_dest(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, ::Type{T}) where T
+    P = parent(A)
+    return sizehint!(spzeros(T, indtype(P), size(A)...), nnz(P))
+end
 matop_dest(::typeof(*), A::QuasiSparseMatrix, B::BiTriSym) =
     similar(A, promote_op(matprod, eltype(A), eltype(B)), (size(A, 1), size(B, 2)))
 
@@ -755,6 +771,63 @@ function dot(A::AbstractSparseMatrixCSC, B::Union{DenseMatrixUnion,MatrixWrapper
     return conj(dot(B, A))
 end
 
+# Frobenius dot of the adjoint/transpose of a CSC matrix with a CSC matrix (issue #627).
+# With `P = parent(A)`, `dot(A, B) = Σ dot(op(P[j,i]), B[i,j])`, so the stored entries of
+# one operand are matched against those of the other at transposed positions. Walking the
+# sparser operand with one cursor per column of the other keeps the work at
+# O(nnz(P) + nnz(B) + n) with O(n) extra memory, where `n` counts the columns of the other
+# operand; a binary search per stored entry is used instead when the other operand is far
+# denser, since the cursors would then sweep all of its entries.
+function dot(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, B::AbstractSparseMatrixCSC)
+    m, n = size(A)
+    size(B) == (m, n) || throw(DimensionMismatch(lazy"A has size ($m, $n) but B has size $(size(B))"))
+    P = parent(A)
+    op = LinearAlgebra.wrapperop(A)
+    r = dot(op(zero(eltype(P))), zero(eltype(B)))
+    (iszero(nnz(P)) || iszero(nnz(B))) && return r
+    if nnz(B) <= nnz(P)
+        return _dot_transposed_walk((b, p) -> dot(op(p), b), B, P, r)
+    else
+        return _dot_transposed_walk((p, b) -> dot(op(p), b), P, B, r)
+    end
+end
+
+# `r + Σ f(X[i,j], Y[j,i])` over the stored entries of `X` that have a stored counterpart
+# in `Y`. Walking the columns of `X` in order, the row index `j` looked for in column `i`
+# of `Y` is nondecreasing, so one cursor per column of `Y` suffices. The cursors visit
+# every stored entry of `Y`, so once `Y` holds well over an order of magnitude more entries
+# (or columns) than `X`, a binary search per entry of `X` is cheaper; the crossover is at
+# a ratio of about 20-50 in measurements.
+function _dot_transposed_walk(f::F, X::AbstractSparseMatrixCSC, Y::AbstractSparseMatrixCSC, r) where F
+    Xrows, Xvals = rowvals(X), nonzeros(X)
+    Yrows, Yvals, Ycolptr = rowvals(Y), nonzeros(Y), getcolptr(Y)
+    if size(Y, 2) + nnz(Y) > 32 * nnz(X)
+        @inbounds for j in axes(X, 2), k in nzrange(X, j)
+            i = Xrows[k]
+            rng = nzrange(Y, i)
+            p = searchsortedfirst(view(Yrows, rng), j) + first(rng) - 1
+            if p <= last(rng) && Yrows[p] == j
+                r += f(Xvals[k], Yvals[p])
+            end
+        end
+        return r
+    end
+    cursor = Ycolptr[1:size(Y, 2)]   # cursor[i] indexes into column i of Y
+    @inbounds for j in axes(X, 2), k in nzrange(X, j)
+        i = Xrows[k]
+        p = cursor[i]
+        pend = Ycolptr[i+1]
+        while p < pend && Yrows[p] < j
+            p += 1
+        end
+        cursor[i] = p
+        if p < pend && Yrows[p] == j
+            r += f(Xvals[k], Yvals[p])
+        end
+    end
+    return r
+end
+
 function dot(x::AbstractSparseVector, D::Diagonal, y::AbstractVector)
     d = D.diag
     if length(x) != length(y) || length(y) != length(d)
@@ -2178,6 +2251,41 @@ function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagona
     C
 end
 
+# Adjoint/transpose of a sparse matrix with a `Diagonal` (issue #619): the generic
+# `Diagonal` kernel in LinearAlgebra visits every element of `C`. With `beta == 0` the
+# adjoint is formed directly in `C` (one `halfperm!`, O(nnz)) and scaled in place;
+# otherwise it is materialized once and handed to the CSC kernels above, which also
+# covers a destination that shares storage with the parent, or whose index type or fixed
+# structure `halfperm!` cannot write.
+function _adjtrans_into!(C::AbstractSparseMatrixCSC, A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC})
+    P = parent(A)
+    return halfperm!(C, P, axes(P, 2), _adjtrans_fun(A))
+end
+_adjtrans_direct(C::AbstractSparseMatrixCSC, A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, beta) =
+    iszero(beta) && !Base.mightalias(C, parent(A)) && !_is_fixed(C) && indtype(C) === indtype(parent(A))
+
+function mul!(C::AbstractSparseMatrixCSC, A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, D::Diagonal, alpha::Number, beta::Number)
+    m, n = size(A)
+    lb = length(D.diag)
+    n == lb || throw(DimensionMismatch(lazy"A has size ($m, $n) but D has size ($lb, $lb)"))
+    size(C) == (m, n) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))"))
+    _adjtrans_direct(C, A, beta) || return mul!(C, copy(A), D, alpha, beta)
+    rmul!(_adjtrans_into!(C, A), D)
+    isone(alpha) || rmul!(C, alpha)
+    return C
+end
+
+function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, alpha::Number, beta::Number)
+    m, n = size(A)
+    lb = length(D.diag)
+    m == lb || throw(DimensionMismatch(lazy"D has size ($lb, $lb) but A has size ($m, $n)"))
+    size(C) == (m, n) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))"))
+    _adjtrans_direct(C, A, beta) || return mul!(C, D, copy(A), alpha, beta)
+    lmul!(D, _adjtrans_into!(C, A))
+    isone(alpha) || rmul!(C, alpha)
+    return C
+end
+
 function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCSC, alpha::Number, beta::Number)
     m, n = size(A)
     b    = D.diag

src/sparsematrix.jl

@@ -1496,10 +1496,12 @@ end
 
 adjoint(A::AbstractSparseMatrixCSC) = Adjoint(A)
 transpose(A::AbstractSparseMatrixCSC) = Transpose(A)
+_adjtrans_fun(::Adjoint) = x -> adjoint(copy(x))
+_adjtrans_fun(::Transpose) = x -> transpose(copy(x))
 Base.copy(A::Adjoint{<:Any,<:AbstractSparseMatrixCSC}) =
-    ftranspose(parent(A), x -> adjoint(copy(x)), eltype(A))
+    ftranspose(parent(A), _adjtrans_fun(A), eltype(A))
 Base.copy(A::Transpose{<:Any,<:AbstractSparseMatrixCSC}) =
-    ftranspose(parent(A), x -> transpose(copy(x)), eltype(A))
+    ftranspose(parent(A), _adjtrans_fun(A), eltype(A))
 function Base.permutedims(A::AbstractSparseMatrixCSC, (a,b))
     (a, b) == (2, 1) && return ftranspose(A, identity)
     (a, b) == (1, 2) && return copy(A)

test/linalg.jl

@@ -2,10 +2,11 @@ module SparseLinalgTests
 
 using Test
 using SparseArrays
-using SparseArrays: nonzeroinds, getcolptr
+using SparseArrays: AbstractSparseMatrixCSC, nonzeroinds, getcolptr, rowvals, nonzeros, fixed, _is_fixed
 using LinearAlgebra
 using Random
 include("forbidproperties.jl")
+include("util/mulcount.jl")
 
 # an AbstractSparseVector outside the types the sparse product kernel handles
 struct WrappedSparseVector <: AbstractSparseVector{Float64,Int}
@@ -715,6 +716,71 @@ end
     @test Diagonal(b) * dA == mul!(sC, Diagonal(b), sA)
     @test Diagonal(b) * dA == lmul!(Diagonal(b), copy(sA))
 
+    # adjoint/transpose of a sparse matrix with a Diagonal (issue #619)
+    for T in (Float64, ComplexF64), W in (adjoint, transpose)
+        S = sprand(T, 7, 3, 0.5); M = Matrix(S)
+        Dl = Diagonal(randn(T, 3)); Dr = Diagonal(randn(T, 7))
+        @test W(S) * Dr isa SparseMatrixCSC
+        @test Dl * W(S) isa SparseMatrixCSC
+        @test W(S) * Dr ≈ W(M) * Dr
+        @test Dl * W(S) ≈ Dl * W(M)
+        @test Dl * W(S) * Dr ≈ Dl * W(M) * Dr
+        @test_throws DimensionMismatch W(S) * Dl
+        @test_throws DimensionMismatch Dr * W(S)
+        # mixed eltypes promote
+        Di = Diagonal(1:7)
+        @test W(S) * Di ≈ W(M) * Di
+        # 3- and 5-argument mul! reach the same kernels
+        C = similar(W(S))
+        @test mul!(C, W(S), Dr) === C
+        @test C ≈ W(M) * Dr
+        @test mul!(C, Dl, W(S)) === C
+        @test C ≈ Dl * W(M)
+        C0 = sprand(T, 3, 7, 0.5)
+        @test mul!(copy(C0), W(S), Dr, 2, 3) ≈ 2 * W(M) * Dr + 3 * Matrix(C0)
+        @test mul!(copy(C0), Dl, W(S), 2, 3) ≈ 2 * Dl * W(M) + 3 * Matrix(C0)
+        @test mul!(copy(C0), W(S), Dr, 2, 0) ≈ 2 * W(M) * Dr
+        @test_throws DimensionMismatch mul!(C, W(S), Dl)
+        @test_throws DimensionMismatch mul!(similar(S), Dl, W(S))
+        # a destination with another index type goes through a materialized copy
+        C32 = SparseMatrixCSC{T,Int32}(spzeros(3, 7))
+        @test mul!(C32, Dl, W(S)) ≈ Dl * W(M)
+        # so does a destination aliasing the parent
+        Q = sprand(T, 5, 5, 0.5); MQ = Matrix(Q); Dq = Diagonal(randn(T, 5))
+        @test mul!(Q, W(Q), Dq) ≈ W(MQ) * Dq
+        Q = sprand(T, 5, 5, 0.5); MQ = Matrix(Q)
+        @test mul!(Q, Dq, W(Q)) ≈ Dq * W(MQ)
+        # or sharing its storage
+        Q = sprand(T, 5, 5, 0.5); MQ = Matrix(Q)
+        Cs = SparseMatrixCSC(5, 5, copy(getcolptr(Q)), copy(rowvals(Q)), nonzeros(Q))
+        @test mul!(Cs, W(Q), Dq) ≈ W(MQ) * Dq
+        # fixed operands are read, never written
+        F = fixed(S)
+        @test W(F) * Dr isa AbstractSparseMatrixCSC
+        @test W(F) * Dr ≈ W(M) * Dr
+        @test Dl * W(F) isa AbstractSparseMatrixCSC
+        @test Dl * W(F) ≈ Dl * W(M)
+        @test F == S
+    end
+    # a Diagonal times a fixed matrix keeps the structure, and the fixedness, of the input
+    F = fixed(sA)
+    let Dl = Diagonal(randn(3)), Dr = Diagonal(randn(7))
+        @test Dl * F ≈ Dl * dA
+        @test F * Dr ≈ dA * Dr
+        @test _is_fixed(Dl * F) && _is_fixed(F * Dr)
+    end
+    # the kernels touch only the stored entries: exactly nnz(S) scalar multiplications,
+    # whereas the generic Diagonal kernel visits every element of the result
+    S = mulcount_sparse(sprand(20, 30, 0.2))
+    Dl = Diagonal(MulCount.(rand(30))); Dr = Diagonal(MulCount.(rand(20)))
+    for W in (adjoint, transpose)
+        @test mulcount(() -> W(S) * Dr) == nnz(S)
+        @test mulcount(() -> Dl * W(S)) == nnz(S)
+        C = similar(W(S))
+        @test mulcount(() -> mul!(C, W(S), Dr)) == nnz(S)
+        @test mulcount(() -> mul!(C, Dl, W(S))) == nnz(S)
+    end
+
     @test dA * 0.5            == sA * 0.5
     @test dA * 0.5            == mul!(sC, sA, 0.5)
     @test dA * 0.5            == rmul!(copy(sA), 0.5)

test/linalg_products.jl

@@ -4,10 +4,11 @@ module SparseLinalgProductTests
 
 using Test
 using SparseArrays
-using SparseArrays: nonzeroinds, getcolptr
+using SparseArrays: nonzeroinds, getcolptr, rowvals, nonzeros, fixed
 using LinearAlgebra
 using Random
 include("forbidproperties.jl")
+include("util/mulcount.jl")
 
 sA = sprandn(3, 7, 0.5)
 sC = similar(sA)
@@ -217,6 +218,11 @@ end
             @test dot(WA,TB) ≈ dot(WA, Matrix(TB))
             @test dot(TA,WB) ≈ dot(Matrix(TA), WB)
             @test dot(TA,WC) ≈ dot(Matrix(TA), WC)
+            # lazy adjoint/transpose of a sparse matrix (issue #627)
+            @test dot(W(A), TB) ≈ dot(WA, Matrix(TB))
+            @test dot(TA, W(B)) ≈ dot(Matrix(TA), WB)
+            @test dot(W(A), sparse(WC)) ≈ dot(WA, WC)
+            @test_throws DimensionMismatch dot(W(A), B)
         end
         for M in (A, B, C)
             D = Diagonal(M * M')
@@ -236,6 +242,45 @@ end
     @test_throws DimensionMismatch dot(sprand(5,5,0.2),sprand(5,6,0.2))
     @test_throws DimensionMismatch dot(rand(5,5),sprand(5,6,0.2))
     @test_throws DimensionMismatch dot(sprand(5,5,0.2),rand(5,6))
+    # stored zeros, empty columns, and non-square shapes with a lazy adjoint (issue #627)
+    for W in (adjoint, transpose)
+        A = sparse([1, 3, 3, 5], [1, 1, 4, 2], [1.0im, 0.0, 2.0, 3.0], 6, 4)
+        B = sparse([1, 2, 4, 4], [3, 3, 1, 6], [1.0, 0.0, 4.0im, 5.0], 4, 6)
+        @test dot(W(A), B) ≈ dot(W(Matrix(A)), Matrix(B))
+        @test dot(B, W(A)) ≈ dot(Matrix(B), W(Matrix(A)))
+        @test dot(W(spzeros(6, 4)), B) == 0
+        @test dot(W(A), spzeros(4, 6)) == 0
+        # Int eltype and small matrices with `Any`-free result type
+        Ai = sparse([1, 2], [2, 1], [1, 2], 2, 2)
+        @test dot(W(Ai), Ai) == dot(W(Matrix(Ai)), Matrix(Ai)) == 4
+        @test dot(W(Ai), Ai) isa Int
+    end
+    # the kernel walks the sparser operand and multiplies only where both operands store
+    # an entry (plus one multiplication seeding the accumulator), whereas the generic
+    # fallback multiplies every stored entry of the sparse operand
+    P = mulcount_sparse(sparse([1, 2, 3], [1, 2, 3], [1.0, 2.0, 3.0], 6, 4))
+    for W in (adjoint, transpose)
+        # disjoint patterns: `B[i, j]` is stored only where `P[j, i]` is not
+        B = mulcount_sparse(sparse([1, 2, 4, 4], [2, 3, 1, 6], [1.0, 2.0, 3.0, 4.0], 4, 6))
+        @test mulcount(() -> dot(W(P), B)) == 1
+        @test mulcount(() -> dot(B, W(P))) == 1
+        # two matching pairs, found from either side of the walk
+        B = mulcount_sparse(sparse([1, 1, 2, 3, 4, 4], [1, 2, 3, 3, 1, 6], 1.0:6.0, 4, 6))
+        @test nnz(B) > nnz(P)   # walks P
+        @test mulcount(() -> dot(W(P), B)) == 1 + 2
+        Pw = mulcount_sparse(sparse([1, 2, 3, 4, 5, 6, 6], [1, 2, 3, 4, 4, 1, 2], 1.0:7.0, 6, 4))
+        @test nnz(Pw) > nnz(B)  # walks B
+        @test mulcount(() -> dot(W(Pw), B)) == 1 + 2
+    end
+    # far more columns than stored entries: a binary search per entry, no cursor array
+    for W in (adjoint, transpose)
+        P = sparse([1], [1], [1.0], 2, 10^5); B = sparse([1], [1], [2.0], 10^5, 2)
+        @test dot(W(P), B) == 2
+        dot(W(P), B)
+        @test (@allocated dot(W(P), B)) < 1024
+    end
+    # fixed operands are read only
+    @test dot(fixed(sprand(5, 4, 0.5))', sprand(4, 5, 0.5)) isa Float64
 end
 
 @testset "generalized dot product" begin

test/util/mulcount.jl

@@ -0,0 +1,26 @@
+# Deterministic replacement for wall-clock guards (see #781): an eltype that records every
+# scalar multiplication, so that a kernel touching only the stored entries and a generic
+# fallback visiting every element are told apart by the count rather than by timing.
+struct MulCount{T} <: Number
+    x::T
+end
+const MULCOUNT = Ref(0)
+Base.:*(a::MulCount, b::MulCount) = (MULCOUNT[] += 1; MulCount(a.x * b.x))
+Base.:+(a::MulCount, b::MulCount) = MulCount(a.x + b.x)
+Base.:-(a::MulCount, b::MulCount) = MulCount(a.x - b.x)
+Base.:-(a::MulCount) = MulCount(-a.x)
+Base.zero(::Type{MulCount{T}}) where {T} = MulCount(zero(T))
+Base.zero(a::MulCount) = zero(typeof(a))
+Base.one(::Type{MulCount{T}}) where {T} = MulCount(one(T))
+Base.conj(a::MulCount) = MulCount(conj(a.x))
+Base.adjoint(a::MulCount) = conj(a)
+Base.transpose(a::MulCount) = a
+Base.:(==)(a::MulCount, b::MulCount) = a.x == b.x
+Base.iszero(a::MulCount) = iszero(a.x)
+Base.isone(a::MulCount) = isone(a.x)
+Base.promote_rule(::Type{MulCount{T}}, ::Type{MulCount{U}}) where {T,U} = MulCount{promote_type(T, U)}
+# number of scalar multiplications performed by `f()`
+mulcount(f) = (MULCOUNT[] = 0; f(); MULCOUNT[])
+# the same sparse matrix with `MulCount` entries
+mulcount_sparse(S::SparseMatrixCSC) =
+    SparseMatrixCSC(size(S)..., copy(getcolptr(S)), copy(rowvals(S)), MulCount.(nonzeros(S)))