From 4d1cdaef87eea20359856ea9e97dcdcf8de13a3f Mon Sep 17 00:00:00 2001 From: "Viral B. Shah" Date: Fri, 16 Oct 2026 04:15:02 +0000 Subject: [PATCH 1/3] Skip the dense copy of the sparse arguments in a sparse-plus-scalar broadcast such as `A .+ 1` Co-Authored-By: Claude Opus 5.5 --- src/higherorderfns.jl | 73 ++++++++++++++++++++++++++++++++++++++++-- test/higherorderfns.jl | 6 ++++ 2 files changed, 77 insertions(+), 2 deletions(-) diff --git a/src/higherorderfns.jl b/src/higherorderfns.jl index 3f47cdaa..0c1a41d3 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,75 @@ end end _isdensesum(bc::Broadcasted) = _issumtree(bc) && _anydensifying(bc.args...) + +# A dense sum of sparse arrays and scalars evaluates the broadcast with each sparse argument +# replaced by its implicit zero, which gives every entry that no stored entry reaches without a +# dense copy of the argument, then evaluates it again where a stored entry is broadcast to. The +# function is a sum, so evaluating it twice there is unobservable. Number entries keep the +# element type of the copy that they replace; a result type known only from the values is left +# to dense broadcast, and so is a sum with a dense array, which the copy costs little beside. +function _densesum(bc::Broadcasted) + ax = axes(bc) + any(isempty, ax) && return _densebroadcast(bc) + _sparseandscalars(bc) || return _densebroadcast(bc) + # each entry evaluated again looks up every sparse argument, which costs as much as about + # thirty-two entries of the dense copy of an argument + 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) + isconcretetype(Broadcast.combine_eltypes(zbc.f, zbc.args)) || return _densebroadcast(bc) + C = copy(Broadcast.instantiate(zbc)) + _storedentries!(C, bc, bc.args...) + return C +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 entries of the result 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 + +_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, map(Int, size(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, map(Int, size(S)), (Int(rv[k]), j)) + end + return C +end +# the entries of `C` that the stored entry at `I` of an argument of size `sz` is broadcast to +function _reevaluate!(C, bc, sz::NTuple{M,Int}, I::NTuple{M,Int}) where {M} + ranges = ntuple(d -> d > M || sz[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..fa3d0b9d 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,8 @@ 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) @test (.-s)::SparseVector == .-fs # a sparse destination keeps the result sparse x = copy(s); x .+= d; x .-= 1 From 3f59b6d809731f1273ed44a7a9fab13a6d17325e Mon Sep 17 00:00:00 2001 From: "Viral B. Shah" Date: Tue, 20 Oct 2026 07:49:14 +0000 Subject: [PATCH 2/3] Simplify the comments and helpers of the sparse dense-sum broadcast Co-Authored-By: Claude Opus 5.5 --- src/higherorderfns.jl | 37 +++++++++++++++++++------------------ 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/src/higherorderfns.jl b/src/higherorderfns.jl index 0c1a41d3..b50b0e8d 100644 --- a/src/higherorderfns.jl +++ b/src/higherorderfns.jl @@ -1383,31 +1383,30 @@ end _isdensesum(bc::Broadcasted) = _issumtree(bc) && _anydensifying(bc.args...) -# A dense sum of sparse arrays and scalars evaluates the broadcast with each sparse argument -# replaced by its implicit zero, which gives every entry that no stored entry reaches without a -# dense copy of the argument, then evaluates it again where a stored entry is broadcast to. The -# function is a sum, so evaluating it twice there is unobservable. Number entries keep the -# element type of the copy that they replace; a result type known only from the values is left -# to dense broadcast, and so is a sum with a dense array, which the copy costs little beside. +# 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) && return _densebroadcast(bc) - _sparseandscalars(bc) || return _densebroadcast(bc) - # each entry evaluated again looks up every sparse argument, which costs as much as about - # thirty-two entries of the dense copy of an argument + (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) isconcretetype(Broadcast.combine_eltypes(zbc.f, zbc.args)) || return _densebroadcast(bc) C = copy(Broadcast.instantiate(zbc)) - _storedentries!(C, bc, bc.args...) - return C + 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 entries of the result their stored entries reach +# 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) @@ -1422,6 +1421,7 @@ _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) @@ -1431,20 +1431,21 @@ _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, map(Int, size(S)), (Int(i),)) + _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, map(Int, size(S)), (Int(rv[k]), j)) + _reevaluate!(C, bc, S, (Int(rv[k]), j)) end return C end -# the entries of `C` that the stored entry at `I` of an argument of size `sz` is broadcast to -function _reevaluate!(C, bc, sz::NTuple{M,Int}, I::NTuple{M,Int}) where {M} - ranges = ntuple(d -> d > M || sz[d] == 1 ? UnitRange(axes(C, d)) : (I[d]:I[d]), Val(ndims(C))) +# 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 From d5f299fcfc03d0bc5f3702ba06f2fece445b0c2c Mon Sep 17 00:00:00 2001 From: "Viral B. Shah" Date: Tue, 20 Oct 2026 08:09:02 +0000 Subject: [PATCH 3/3] Fall back to dense broadcast when the zeros alone do not give the result element type Co-Authored-By: Claude Opus 5.5 --- src/higherorderfns.jl | 4 +++- test/higherorderfns.jl | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/higherorderfns.jl b/src/higherorderfns.jl index b50b0e8d..be693c63 100644 --- a/src/higherorderfns.jl +++ b/src/higherorderfns.jl @@ -1399,7 +1399,9 @@ function _densesum(bc::Broadcasted) 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) - isconcretetype(Broadcast.combine_eltypes(zbc.f, zbc.args)) || return _densebroadcast(bc) + # 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 diff --git a/test/higherorderfns.jl b/test/higherorderfns.jl index fa3d0b9d..c8ad4bdd 100644 --- a/test/higherorderfns.jl +++ b/test/higherorderfns.jl @@ -836,6 +836,8 @@ densesum(x, y) = x .+ y @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