diff --git a/src/linalg.jl b/src/linalg.jl index 6ee97c9f..34b2aa40 100644 --- a/src/linalg.jl +++ b/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 diff --git a/src/sparsematrix.jl b/src/sparsematrix.jl index 95fea008..16a163f2 100644 --- a/src/sparsematrix.jl +++ b/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) diff --git a/test/linalg.jl b/test/linalg.jl index 4ee4d99b..a39e6eab 100644 --- a/test/linalg.jl +++ b/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) diff --git a/test/linalg_products.jl b/test/linalg_products.jl index 582afdda..a59e4340 100644 --- a/test/linalg_products.jl +++ b/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 diff --git a/test/util/mulcount.jl b/test/util/mulcount.jl new file mode 100644 index 00000000..e080379e --- /dev/null +++ b/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)))