diff --git a/src/higherorderfns.jl b/src/higherorderfns.jl index 3f47cdaa..be693c63 100644 --- a/src/higherorderfns.jl +++ b/src/higherorderfns.jl @@ -1216,7 +1216,7 @@ end # broadcast entry points for combinations of sparse arrays and other (scalar) types @inline function copy(bc::Broadcasted{<:SparseVecOrMatStyle}) - _isdensesum(bc) && return _densebroadcast(bc) + _isdensesum(bc) && return _densesum(bc) bcf = flatten(bc) _asksthezeros(bc) && _fillszeros(bcf.f, bcf.args...) && return _densenonzero(bc, bcf.f, bcf.args...) return _copy(bcf.f, bcf.args...) @@ -1360,7 +1360,7 @@ broadcast(f::Tf, A::AbstractSparseMatrixCSC, ::Type{T}) where {Tf,T} = broadcast # which is sparse, and inferable, for every element type. function copy(bc::Broadcasted{PromoteToSparse}) - _isdensesum(bc) && return _densebroadcast(bc) + _isdensesum(bc) && return _densesum(bc) bcf = flatten(bc) _asksthezeros(bc) && _fillszeros(bcf.f, bcf.args...) && return _densenonzero(bc, bcf.f, bcf.args...) if can_skip_sparsification(bcf.f, bcf.args...) @@ -1382,6 +1382,78 @@ end end _isdensesum(bc::Broadcasted) = _issumtree(bc) && _anydensifying(bc.args...) + +# A sum of sparse arrays and scalars with a dense result, such as `A .+ 1`, skips the dense +# copy of each sparse argument. It broadcasts once with every sparse argument replaced by its +# implicit zero, which is already right wherever no stored entry reaches, then recomputes the +# entries that a stored entry reaches. The function is made of `+` and `-`, so computing an +# entry twice is harmless. The rest goes to `_densebroadcast`: sums with a dense array, where +# the copy costs little, non-`Number` elements, whose zero may be of another type, and results +# whose element type depends on the values. +function _densesum(bc::Broadcasted) + ax = axes(bc) + (any(isempty, ax) || !_sparseandscalars(bc)) && return _densebroadcast(bc) + # recomputing an entry looks up every sparse argument, which costs about as much as copying + # 32 entries, so beyond that many reached entries the dense copy is cheaper + len = prod(map(length, ax)) + nsparse, reached = _sparsereach(len, bc.args...) + 32 * nsparse * reached <= len || return _densebroadcast(bc) + zbc = Broadcasted{Broadcast.DefaultArrayStyle{length(ax)}}(bc.f, map(_zerosparse, bc.args), ax) + # the zeros alone must give the element type that dense broadcast infers from all the values + T = Broadcast.combine_eltypes(zbc.f, zbc.args) + isconcretetype(T) && T === Broadcast.combine_eltypes(bc.f, bc.args) || return _densebroadcast(bc) + C = copy(Broadcast.instantiate(zbc)) + return _storedentries!(C, bc, bc.args...) +end +_zerosparse(bc::Broadcasted) = Broadcasted(bc.f, map(_zerosparse, bc.args)) +_zerosparse(S::SparseVecOrMat) = _densezero(S) +_zerosparse(x) = x +# the number of sparse arguments, and of the result entries their stored entries reach +_sparsereach(len) = (0, 0) +function _sparsereach(len, x, rest::Vararg{Any,N}) where {N} + n, r = _sparsereach(len, x) + n2, r2 = _sparsereach(len, rest...) + return (n + n2, r + r2) +end +_sparsereach(len, x::Broadcasted) = _sparsereach(len, x.args...) +_sparsereach(len, S::SparseVecOrMat) = (1, nnz(S) * (len รท length(S))) +_sparsereach(len, x) = (0, 0) +_sparseandscalars(bc::Broadcasted) = all(_sparseandscalars, bc.args) +_sparseandscalars(S::SparseVecOrMat) = eltype(S) <: Number +_sparseandscalars(x::Union{AbstractArray,Tuple}) = false +_sparseandscalars(x) = true + +# recompute the entries of `C` that the stored entries of the sparse arguments reach +_storedentries!(C, bc) = C +function _storedentries!(C, bc, x, rest::Vararg{Any,N}) where {N} + _storedentries!(C, bc, x) + return _storedentries!(C, bc, rest...) +end +_storedentries!(C, bc, x::Broadcasted) = _storedentries!(C, bc, x.args...) +_storedentries!(C, bc, x) = C +function _storedentries!(C, bc, S::AbstractCompressedVector) + for i in nonzeroinds(S) + _reevaluate!(C, bc, S, (Int(i),)) + end + return C +end +function _storedentries!(C, bc, S::AbstractSparseMatrixCSC) + rv = getrowval(S) + for j in axes(S, 2), k in nzrange(S, j) + _reevaluate!(C, bc, S, (Int(rv[k]), j)) + end + return C +end +# the stored entry at `I` of `S` reaches all of `C` along the dimensions where `S` has size 1, +# and index `I[d]` along each other dimension `d` +function _reevaluate!(C, bc, S, I::Tuple{Vararg{Int}}) + ranges = ntuple(d -> size(S, d) == 1 ? UnitRange(axes(C, d)) : (I[d]:I[d]), Val(ndims(C))) + for J in CartesianIndices(ranges) + C[J] = bc[J] + end + return C +end + function _densebroadcast(bc::Broadcasted) # an empty result reads no entry, so the sparse arguments are left as they are dbc = any(isempty, axes(bc)) ? bc : Broadcast.instantiate(_densifysparse(bc)) diff --git a/test/higherorderfns.jl b/test/higherorderfns.jl index 7ef4578d..c8ad4bdd 100644 --- a/test/higherorderfns.jl +++ b/test/higherorderfns.jl @@ -808,6 +808,10 @@ densesum(x, y) = x .+ y @test @inferred(broadcast(+, s, d))::Vector{Float64} == fs + d @test @inferred(broadcast(-, D, S))::Matrix{Float64} == D - fS @test @inferred(broadcast(+, s, 1))::Vector{Float64} == fs .+ 1 + # a sum of a sparse array with few stored entries and a scalar makes no dense copy of it + L = sparse([2, 400], [1, 3], [-0.0, NaN], 400, 3) + @test isequal(@inferred(densesum(L, 1))::Matrix{Float64}, Array(L) .+ 1) + @test @allocated(densesum(L, 1)) < 2 * sizeof(Array(L)) # the other broadcasts with a dense array or a scalar, and sums without one, stay sparse @test (S .* D)::SparseMatrixCSC == fS .* D @test (S .- S)::SparseMatrixCSC == fS .- fS @@ -830,6 +834,10 @@ densesum(x, y) = x .+ y @test @inferred(broadcast(-, 1, S))::Matrix{Float64} == 1 .- fS @test (s .+ fill(1.0))::Vector{Float64} == fs .+ 1 @test (S .+ s .- Ref(2))::Matrix{Float64} == fS .+ fs .- 2 + l = sparsevec([5], [2.0], 400) + @test isequal((L .+ l .- 2)::Matrix{Float64}, Array(L) .+ Array(l) .- 2) + R = SparseMatrixCSC{Real,Int}(400, 3, [1, 2, 2, 2], [7], Real[1.5]) + @test (R .+ 1)::Matrix{Real} == Array(R) .+ 1 @test (.-s)::SparseVector == .-fs # a sparse destination keeps the result sparse x = copy(s); x .+= d; x .-= 1