Skip to content

avoid intermediate map allocations in multi-arg mapreduce #55301

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

avoid intermediate map allocations in multi-arg mapreduce #55301

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5d45b21
avoid intermediate map allocations in multi-arg mapreduce
mbauman Jul 29, 2024
07f320a
restore support for args of different dimensionalities
mbauman Jul 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion base/broadcast.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -578,7 +578,7 @@ an `Int`.
"""
Base.@propagate_inbounds newindex(arg, I::CartesianIndex) = to_index(_newindex(axes(arg), I.I))
Base.@propagate_inbounds newindex(arg, I::Integer) = to_index(_newindex(axes(arg), (I,)))
Base.@propagate_inbounds _newindex(ax::Tuple, I::Tuple) = (ifelse(length(ax[1]) == 1, ax[1][1], I[1]), _newindex(tail(ax), tail(I))...)
Base.@propagate_inbounds _newindex(ax::Tuple, I::Tuple) = (ifelse(length(ax[1]) == 1, ax[1][begin], I[1]), _newindex(tail(ax), tail(I))...)
Base.@propagate_inbounds _newindex(ax::Tuple{}, I::Tuple) = ()
Base.@propagate_inbounds _newindex(ax::Tuple, I::Tuple{}) = (ax[1][1], _newindex(tail(ax), ())...)
Base.@propagate_inbounds _newindex(ax::Tuple{}, I::Tuple{}) = ()
Expand Down
29 changes: 23 additions & 6 deletions base/reducedim.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,10 +66,17 @@ _realtype(T::Type) = T
_realtype(::Union{typeof(abs),typeof(abs2)}, T) = _realtype(T)
_realtype(::Any, T) = T

function reducedim_init(f, op::Union{typeof(+),typeof(add_sum)}, A::AbstractArray, region)
"""
reducedim_init(f, op, A, region)

Internal function to create the correctly-sized array holding the initial
values (possibly the return values, or possibly erroring) for a reduction
over the region (dims) when no initial value was explicitly passed.
"""
function reducedim_init(f, op::Union{typeof(+),typeof(add_sum)}, A::AbstractArrayOrBroadcasted, region)
_reducedim_init(f, op, zero, sum, A, region)
end
function reducedim_init(f, op::Union{typeof(*),typeof(mul_prod)}, A::AbstractArray, region)
function reducedim_init(f, op::Union{typeof(*),typeof(mul_prod)}, A::AbstractArrayOrBroadcasted, region)
_reducedim_init(f, op, one, prod, A, region)
end
function _reducedim_init(f, op, fv, fop, A, region)
Expand All @@ -87,7 +94,7 @@ end

# initialization when computing minima and maxima requires a little care
for (f1, f2, initval, typeextreme) in ((:min, :max, :Inf, :typemax), (:max, :min, :(-Inf), :typemin))
@eval function reducedim_init(f, op::typeof($f1), A::AbstractArray, region)
@eval function reducedim_init(f, op::typeof($f1), A::AbstractArrayOrBroadcasted, region)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change seems incomplete as view in L106 has no Broadcasted support.
we might need

reduced_view(A::AbstractArray, ri) = view(A, ri...)
reduced_view(bc::Broadcasted, ri) = Broadcasted(bc.style, bc.f, bc.args, ri)

# First compute the reduce indices. This will throw an ArgumentError
# if any region is invalid
ri = reduced_indices(A, region)
Expand Down Expand Up @@ -129,7 +136,7 @@ for (f1, f2, initval, typeextreme) in ((:min, :max, :Inf, :typemax), (:max, :min
end
end

function reducedim_init(f::ExtremaMap, op::typeof(_extrema_rf), A::AbstractArray, region)
function reducedim_init(f::ExtremaMap, op::typeof(_extrema_rf), A::AbstractArrayOrBroadcasted, region)
# First compute the reduce indices. This will throw an ArgumentError
# if any region is invalid
ri = reduced_indices(A, region)
Expand Down Expand Up @@ -326,8 +333,18 @@ julia> mapreduce(isodd, |, a, dims=1)
"""
mapreduce(f, op, A::AbstractArrayOrBroadcasted; dims=:, init=_InitialValue()) =
_mapreduce_dim(f, op, init, A, dims)
mapreduce(f, op, A::AbstractArrayOrBroadcasted, B::AbstractArrayOrBroadcasted...; kw...) =
reduce(op, map(f, A, B...); kw...)
function mapreduce(f, op, A::AbstractArrayOrBroadcasted, B::AbstractArrayOrBroadcasted...; kwargs...)
Adims = ndims(A)
if any(b->Adims != ndims(b), B)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if any(b->Adims != ndims(b), B)
if Adims != 1 && any(b->Adims != ndims(b), B)

# The map documentation explicitly allows arrays of different shapes when ndims don't match.
# We fall back to `map` in these cases — which will allocate — but perhaps this allocation will
# help someone identify their bug: if you hit this, the dimensionalities of your args don't match!
return reduce(op, map(f, A, B...); kwargs...)
end
Aax = axes(A)
all(b->Aax==axes(b), B) || throw(ArgumentError("all arguments must have the same axes"))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perhaps DimensionMismatch is better than ArgumentError here?

mapreduce(splat(f), op, Broadcast.instantiate(broadcasted(tuple, A, B...)); kwargs...)
end

_mapreduce_dim(f, op, nt, A::AbstractArrayOrBroadcasted, ::Colon) =
mapfoldl_impl(f, op, nt, A)
Expand Down