Skip to content

Commit 79c1ba4

Browse files
committed
Accelerate reduce over all dims
Test-fix: `Broadcasted` should behave like a cartesian-indexed Array. (we don't always call `foldl` for performance) Update reduce.jl Update reduce.jl Update reduce.jl Update reduce.jl
1 parent 084585b commit 79c1ba4

File tree

2 files changed

+16
-5
lines changed

2 files changed

+16
-5
lines changed

base/reduce.jl

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -360,8 +360,8 @@ mapreduce_empty(::typeof(identity), op, T) = reduce_empty(op, T)
360360
mapreduce_empty(::typeof(abs), op, T) = abs(reduce_empty(op, T))
361361
mapreduce_empty(::typeof(abs2), op, T) = abs2(reduce_empty(op, T))
362362

363-
mapreduce_empty(f::typeof(abs), ::typeof(max), T) = abs(zero(T))
364-
mapreduce_empty(f::typeof(abs2), ::typeof(max), T) = abs2(zero(T))
363+
mapreduce_empty(::typeof(abs), ::typeof(max), T) = abs(zero(T))
364+
mapreduce_empty(::typeof(abs2), ::typeof(max), T) = abs2(zero(T))
365365

366366
# For backward compatibility:
367367
mapreduce_empty_iter(f, op, itr, ItrEltype) =
@@ -412,6 +412,10 @@ mapreduce_first(f, op, x) = reduce_first(op, f(x))
412412

413413
_mapreduce(f, op, A::AbstractArrayOrBroadcasted) = _mapreduce(f, op, IndexStyle(A), A)
414414

415+
_mapreduce(f, op, A::AbstractVector) = _mapreduce(f, op, IndexLinear(), A)
416+
417+
_mapreduce(f, op, A::AbstractZeroDimArray) = mapreduce_first(f, op, A[])
418+
415419
function _mapreduce(f, op, ::IndexLinear, A::AbstractArrayOrBroadcasted)
416420
inds = LinearIndices(A)
417421
n = length(inds)
@@ -437,7 +441,14 @@ end
437441

438442
mapreduce(f, op, a::Number) = mapreduce_first(f, op, a)
439443

440-
_mapreduce(f, op, ::IndexCartesian, A::AbstractArrayOrBroadcasted) = mapfoldl(f, op, A)
444+
function _mapreduce(f, op, ::IndexCartesian, A::AbstractArrayOrBroadcasted)
445+
(isempty(A) || length(axes1(A)) < 16) && return mapfoldl(f, op, A)
446+
mapfoldl(op, CartesianIndices(tail(axes(A)))) do IA
447+
@inline elf(i) = @inbounds f(A[i, IA])
448+
ax1 = axes1(A)
449+
mapreduce_impl(elf, op, ax1, firstindex(ax1), lastindex(ax1))
450+
end
451+
end
441452

442453
"""
443454
reduce(op, itr; [init])

test/broadcast.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -911,9 +911,9 @@ end
911911
@test IndexStyle(bc) == IndexLinear()
912912
@test reduce(paren, bc) == reduce(paren, xs)
913913
# If `Broadcasted` does not have `IndexLinear` style, it should
914-
# hit the `foldl` branch:
914+
# behave like a cartesian-indexed Array (PR #43618)
915915
@test IndexStyle(bcraw) == IndexCartesian()
916-
@test reduce(paren, bcraw) == foldl(paren, xs)
916+
@test reduce(paren, bcraw) == reduce(paren, view(xs, 1:length(xs), 1:1))
917917

918918
# issue #41055
919919
bc = Broadcast.instantiate(Broadcast.broadcasted(Base.literal_pow, Ref(^), [1,2], Ref(Val(2))))

0 commit comments

Comments
 (0)