Skip to content

Commit d319536

Browse files
committed
limit mapreduce_impl to min/max
1 parent 8cf8acb commit d319536

File tree

1 file changed

+12
-5
lines changed

1 file changed

+12
-5
lines changed

base/reducedim.jl

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -267,13 +267,20 @@ function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArrayOrBroadcasted)
267267
if reducedim1(R, A)
268268
# keep the accumulator as a local variable when reducing along the first dimension
269269
i1 = first(axes1(R))
270+
ax1 = axes1(A)
270271
@inbounds for IA in CartesianIndices(indsAt)
271272
IR = Broadcast.newindex(IA, keep, Idefault)
272-
@inline elf(i) = @inbounds f(A[i, IA])
273-
ax1 = axes(A, 1)
274-
# use mapreduce_impl for performance
275-
r = mapreduce_impl(elf, op, ax1, firstindex(ax1), lastindex(ax1))
276-
R[i1,IR] = op(R[i1,IR], r)
273+
r = R[i1,IR]
274+
if op === min || op === max #|| op === _extrema_op # TODO: uncomment this
275+
@inline elf(i) = @inbounds f(A[i, IA])
276+
r′ = mapreduce_impl(elf, op, ax1, firstindex(ax1), lastindex(ax1))
277+
r = op(r, r′)
278+
else # the original version seems faster for Intrinsic based reduction
279+
@simd for i in ax1
280+
r = op(r, f(A[i, IA]))
281+
end
282+
end
283+
R[i1,IR] = r
277284
end
278285
else
279286
@inbounds for IA in CartesianIndices(indsAt)

0 commit comments

Comments
 (0)