Update reduce.jl

JuliaLang · Dec 31, 2021 · 2dc55e9 · 2dc55e9
1 parent 00202f7
commit 2dc55e9
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/base/reduce.jl b/base/reduce.jl
@@ -437,7 +437,19 @@ end
 
 mapreduce(f, op, a::Number) = mapreduce_first(f, op, a)
 
-_mapreduce(f, op, ::IndexCartesian, A::AbstractArrayOrBroadcasted) = mapfoldl(f, op, A)
+function _mapreduce(f, op, ::IndexCartesian, A::AbstractArrayOrBroadcasted)
+    ndims(A) <= 1 && return _mapreduce(f, op, IndexLinear(), A)
+    size(A, 1) < 16 && return mapfoldl(f, op, A)
+    ax1, axo = Base.axes1(A), tail(axes(A))
+    IA, outer = Iterators.peel(CartesianIndices(axo))
+    @inline elf(i) = @inbounds f(A[i, IA])
+    r = mapreduce_impl(elf, op, ax1, firstindex(ax1), lastindex(ax1))
+    for IA in outer
+        r′ = mapreduce_impl(elf, op, ax1, firstindex(ax1), lastindex(ax1))
+        r = op(r, r′)
+    end
+    r
+end
 
 """
     reduce(op, itr; [init])