Extend min/maximum optimization to much shorter length

Author: N5N3

base/reduce.jl

@@ -625,49 +625,41 @@ isbadzero(op, x) = false
 isgoodzero(::typeof(max), x) = isbadzero(min, x)
 isgoodzero(::typeof(min), x) = isbadzero(max, x)
 
-function mapreduce_impl(f, op::Union{typeof(max), typeof(min)},
-                        A::AbstractArrayOrBroadcasted, first::Int, last::Int)
-    # 1. This optimization gives different result from general fallback, if the inputs `f.(A)`
-    #    contains both 'missing' and 'Nan'.
-    # 2. For Integer cases, general fallback seems faster.
-    # Based the above reasons, only use this for AbstractFloat cases.
-    Eltype = _return_type(i -> f(A[i]), Tuple{Int})
+function mapreduce_impl(f, op::Union{typeof(max),typeof(min)},
+                        A::AbstractArrayOrBroadcasted, fi::Int, la::Int)
+    @inline elf(i) = @inbounds f(A[i])
+    # 1. If `f.(A)` contains both 'missing' and 'Nan', this might return `NaN`.
+    # 2. For Integer input, general fallback is about 2x faster.
+    # Thus limit this optimization to AbstractFloat.
+    Eltype = _return_type(elf, Tuple{Int})
     Eltype <: AbstractFloat ||
-        return invoke(mapreduce_impl,Tuple{Any,Any,AbstractArrayOrBroadcasted,Int,Int},f,op,A,first,last)
-    a1 = @inbounds A[first]
-    v1 = mapreduce_first(f, op, a1)
-    v2 = v3 = v4 = v1
-    chunk_len = 256
-    start = first + 1
-    simdstop  = start + chunk_len - 4
-    while simdstop <= last - 3
-        # short circuit in case of NaN or missing
-        v1 == v1 || return v1
-        v2 == v2 || return v2
-        v3 == v3 || return v3
-        v4 == v4 || return v4
-        @inbounds for i in start:4:simdstop
-            v1 = _fast(op, v1, f(A[i+0]))
-            v2 = _fast(op, v2, f(A[i+1]))
-            v3 = _fast(op, v3, f(A[i+2]))
-            v4 = _fast(op, v4, f(A[i+3]))
+        return invoke(mapreduce_impl,Tuple{Any,Any,AbstractArrayOrBroadcasted,Int,Int},f,op,A,fi,la)
+    v1 = v2 = v3 = v4 = elf(fi)
+    len = (la - fi) >> 2
+    i = fi
+    for I in Iterators.partition(1:len, 64)
+        for _ in I
+            v1 = _fast(op, v1, elf(i+=1))
+            v2 = _fast(op, v2, elf(i+=1))
+            v3 = _fast(op, v3, elf(i+=1))
+            v4 = _fast(op, v4, elf(i+=1))
         end
-        checkbounds(A, simdstop+3)
-        start += chunk_len
-        simdstop += chunk_len
+        # short circuit in case of NaN
+        isnan(v1) && return v1
+        isnan(v2) && return v2
+        isnan(v3) && return v3
+        isnan(v4) && return v4
     end
     v = op(op(v1,v2),op(v3,v4))
-    for i in start:last
-        @inbounds ai = A[i]
-        v = op(v, f(ai))
+    while i < la
+        v = op(v, elf(i+=1))
     end
-
     # enforce correct order of 0.0 and -0.0
     # e.g. maximum([0.0, -0.0]) === 0.0
     # should hold
     if isbadzero(op, v)
-        for i in first:last
-            x = @inbounds A[i]
+        for i in fi:la
+            x = elf(i)
             isgoodzero(op,x) && return x
         end
     end