limit mapreduce_impl to min/max

JuliaLang · Jan 1, 2022 · d319536 · d319536
1 parent 8cf8acb
commit d319536
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/base/reducedim.jl b/base/reducedim.jl
@@ -267,13 +267,20 @@ function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArrayOrBroadcasted)
     if reducedim1(R, A)
         # keep the accumulator as a local variable when reducing along the first dimension
         i1 = first(axes1(R))
+        ax1 = axes1(A)
         @inbounds for IA in CartesianIndices(indsAt)
             IR = Broadcast.newindex(IA, keep, Idefault)
-            @inline elf(i) = @inbounds f(A[i, IA])
-            ax1 = axes(A, 1)
-            # use mapreduce_impl for performance
-            r = mapreduce_impl(elf, op, ax1, firstindex(ax1), lastindex(ax1))
-            R[i1,IR] = op(R[i1,IR], r)
+            r = R[i1,IR]
+            if op === min || op === max #|| op === _extrema_op # TODO: uncomment this
+                @inline elf(i) = @inbounds f(A[i, IA])
+                r′ = mapreduce_impl(elf, op, ax1, firstindex(ax1), lastindex(ax1))
+                r = op(r, r′)
+            else # the original version seems faster for Intrinsic based reduction
+                @simd for i in ax1
+                    r = op(r, f(A[i, IA]))
+                end
+            end
+            R[i1,IR] = r
         end
     else
         @inbounds for IA in CartesianIndices(indsAt)