Skip to content

Commit bd93724

Browse files
committed
Collapse spbroadcast_args! into broadcast!.
As suggested by @Sacha0. Also add nothing argument to broadcast! calls in sparse broadcast as suggested by @timholy.
1 parent d0c0c14 commit bd93724

File tree

1 file changed

+15
-25
lines changed

1 file changed

+15
-25
lines changed

base/sparse/higherorderfns.jl

+15-25
Original file line numberDiff line numberDiff line change
@@ -107,18 +107,6 @@ broadcast(f::Tf, A::SparseMatrixCSC) where {Tf} = _noshapecheck_map(f, A)
107107
end
108108
return C
109109
end
110-
@inline function broadcast!(f::Tf, dest::SparseVecOrMat, ::Void, As::Vararg{Any,N}) where {Tf,N}
111-
if f isa typeof(identity) && N == 1
112-
A = As[1]
113-
if A isa Number
114-
return fill!(dest, A)
115-
elseif A isa AbstractArray && Base.axes(dest) == Base.axes(A)
116-
return copyto!(dest, A)
117-
end
118-
end
119-
spbroadcast_args!(f, dest, Broadcast.combine_styles(As...), As...)
120-
return dest
121-
end
122110

123111
# the following three similar defs are necessary for type stability in the mixed vector/matrix case
124112
broadcast(f::Tf, A::SparseVector, Bs::Vararg{SparseVector,N}) where {Tf,N} =
@@ -1015,25 +1003,27 @@ broadcast(f, ::PromoteToSparse, ::Void, ::Void, As::Vararg{Any,N}) where {N} =
10151003
# For broadcast! with ::Any inputs, we need a layer of indirection to determine whether
10161004
# the inputs can be promoted to SparseVecOrMat. If it's just SparseVecOrMat and scalars,
10171005
# we can handle it here, otherwise see below for the promotion machinery.
1018-
function spbroadcast_args!(f::Tf, C, ::SPVM, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N}) where {Tf,N}
1019-
_aresameshape(C, A, Bs...) && return _noshapecheck_map!(f, C, A, Bs...)
1020-
Base.Broadcast.check_broadcast_indices(axes(C), A, Bs...)
1006+
function broadcast!(f::Tf, dest::SparseVecOrMat, ::SPVM, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N}) where {Tf,N}
1007+
if f isa typeof(identity) && N == 0 && Base.axes(dest) == Base.axes(A)
1008+
return copyto!(dest, A)
1009+
end
1010+
_aresameshape(dest, A, Bs...) && return _noshapecheck_map!(f, dest, A, Bs...)
1011+
Base.Broadcast.check_broadcast_indices(axes(dest), A, Bs...)
10211012
fofzeros = f(_zeros_eltypes(A, Bs...)...)
10221013
fpreszeros = _iszero(fofzeros)
1023-
return fpreszeros ? _broadcast_zeropres!(f, C, A, Bs...) :
1024-
_broadcast_notzeropres!(f, fofzeros, C, A, Bs...)
1014+
fpreszeros ? _broadcast_zeropres!(f, dest, A, Bs...) :
1015+
_broadcast_notzeropres!(f, fofzeros, dest, A, Bs...)
1016+
return dest
10251017
end
1026-
function spbroadcast_args!(f::Tf, dest, ::SPVM, mixedsrcargs::Vararg{Any,N}) where {Tf,N}
1018+
function broadcast!(f::Tf, dest::SparseVecOrMat, ::SPVM, mixedsrcargs::Vararg{Any,N}) where {Tf,N}
10271019
# mixedsrcargs contains nothing but SparseVecOrMat and scalars
10281020
parevalf, passedsrcargstup = capturescalars(f, mixedsrcargs)
1029-
return broadcast!(parevalf, dest, passedsrcargstup...)
1030-
end
1031-
function spbroadcast_args!(f::Tf, dest, ::PromoteToSparse, mixedsrcargs::Vararg{Any,N}) where {Tf,N}
1032-
broadcast!(f, dest, map(_sparsifystructured, mixedsrcargs)...)
1021+
broadcast!(parevalf, dest, nothing, passedsrcargstup...)
1022+
return dest
10331023
end
1034-
function spbroadcast_args!(f::Tf, dest, ::Any, mixedsrcargs::Vararg{Any,N}) where {Tf,N}
1035-
# Fallback. From a performance perspective would it be best to densify?
1036-
Broadcast._broadcast!(f, dest, mixedsrcargs...)
1024+
function broadcast!(f::Tf, dest::SparseVecOrMat, ::PromoteToSparse, mixedsrcargs::Vararg{Any,N}) where {Tf,N}
1025+
broadcast!(f, dest, nothing, map(_sparsifystructured, mixedsrcargs)...)
1026+
return dest
10371027
end
10381028

10391029
_sparsifystructured(M::AbstractMatrix) = SparseMatrixCSC(M)

0 commit comments

Comments
 (0)