Decouple sum/prod promotion from reduce

Author: TotalVerb

base/reduce.jl

@@ -5,45 +5,29 @@
 ###### Generic (map)reduce functions ######
 
 if Int === Int32
-const SmallSigned = Union{Int8,Int16}
-const SmallUnsigned = Union{UInt8,UInt16}
+    const SmallSigned = Union{Int8,Int16}
+    const SmallUnsigned = Union{UInt8,UInt16}
 else
-const SmallSigned = Union{Int8,Int16,Int32}
-const SmallUnsigned = Union{UInt8,UInt16,UInt32}
+    const SmallSigned = Union{Int8,Int16,Int32}
+    const SmallUnsigned = Union{UInt8,UInt16,UInt32}
 end
 
-const CommonReduceResult = Union{UInt64,UInt128,Int64,Int128,Float16,Float32,Float64}
-const WidenReduceResult = Union{SmallSigned, SmallUnsigned}
-
-promote_sys_size{T}(::Type{T}) = T
-promote_sys_size{T<:SmallSigned}(::Type{T}) = Int
-promote_sys_size{T<:SmallUnsigned}(::Type{T}) = UInt
-# r_promote_type: promote T to the type of reduce(op, ::Array{T})
-# (some "extra" methods are required here to avoid ambiguity warnings)
-r_promote_type(op, ::Type{T}) where {T} = T
-r_promote_type(op, ::Type{T}) where {T<:WidenReduceResult} = promote_sys_size(T)
-r_promote_type(::typeof(+), ::Type{T}) where {T<:WidenReduceResult} = promote_sys_size(T)
-r_promote_type(::typeof(*), ::Type{T}) where {T<:WidenReduceResult} = promote_sys_size(T)
-r_promote_type(::typeof(+), ::Type{T}) where {T<:Number} = typeof(zero(T)+zero(T))
-r_promote_type(::typeof(*), ::Type{T}) where {T<:Number} = typeof(one(T)*one(T))
-r_promote_type(::typeof(scalarmax), ::Type{T}) where {T<:WidenReduceResult} = T
-r_promote_type(::typeof(scalarmin), ::Type{T}) where {T<:WidenReduceResult} = T
-r_promote_type(::typeof(max), ::Type{T}) where {T<:WidenReduceResult} = T
-r_promote_type(::typeof(min), ::Type{T}) where {T<:WidenReduceResult} = T
-
-# r_promote: promote x to the type of reduce(op, [x])
-r_promote(op, x::T) where {T} = convert(r_promote_type(op, T), x)
+# Certain reductions like sum and prod may wish to promote the items being
+# reduced over to an appropriate size.
+promote_sys_size(x) = x
+promote_sys_size(x::SmallSigned) = Int(x)
+promote_sys_size(x::SmallUnsigned) = UInt(x)
 
 ## foldl && mapfoldl
 
 @noinline function mapfoldl_impl(f, op, v0, itr, i)
     # Unroll the while loop once; if v0 is known, the call to op may
     # be evaluated at compile time
     if done(itr, i)
-        return r_promote(op, v0)
+        return v0
     else
         (x, i) = next(itr, i)
-        v = op(r_promote(op, v0), f(x))
+        v = op(v0, f(x))
         while !done(itr, i)
             @inbounds (x, i) = next(itr, i)
             v = op(v, f(x))
@@ -108,10 +92,10 @@ function mapfoldr_impl(f, op, v0, itr, i::Integer)
     # Unroll the while loop once; if v0 is known, the call to op may
     # be evaluated at compile time
     if isempty(itr) || i == 0
-        return r_promote(op, v0)
+        return v0
     else
         x = itr[i]
-        v  = op(f(x), r_promote(op, v0))
+        v  = op(f(x), v0)
         while i > 1
             x = itr[i -= 1]
             v = op(f(x), v)
@@ -180,12 +164,12 @@ foldr(op, itr) = mapfoldr(identity, op, itr)
 @noinline function mapreduce_impl(f, op, A::AbstractArray, ifirst::Integer, ilast::Integer, blksize::Int)
     if ifirst == ilast
         @inbounds a1 = A[ifirst]
-        return r_promote(op, f(a1))
+        return f(a1)
     elseif ifirst + blksize > ilast
         # sequential portion
         @inbounds a1 = A[ifirst]
         @inbounds a2 = A[ifirst+1]
-        v = op(r_promote(op, f(a1)), r_promote(op, f(a2)))
+        v = op(f(a1), f(a2))
         @simd for i = ifirst + 2 : ilast
             @inbounds ai = A[i]
             v = op(v, f(ai))
@@ -245,13 +229,14 @@ pairwise_blocksize(::typeof(abs2), ::typeof(+)) = 4096
 # handling empty arrays
 _empty_reduce_error() = throw(ArgumentError("reducing over an empty collection is not allowed"))
 mr_empty(f, op, T) = _empty_reduce_error()
-# use zero(T)::T to improve type information when zero(T) is not defined
-mr_empty(::typeof(identity), op::typeof(+), T) = r_promote(op, zero(T)::T)
-mr_empty(::typeof(abs), op::typeof(+), T) = r_promote(op, abs(zero(T)::T))
-mr_empty(::typeof(abs2), op::typeof(+), T) = r_promote(op, abs2(zero(T)::T))
-mr_empty(::typeof(identity), op::typeof(*), T) = r_promote(op, one(T)::T)
-mr_empty(::typeof(abs), op::typeof(scalarmax), T) = abs(zero(T)::T)
-mr_empty(::typeof(abs2), op::typeof(scalarmax), T) = abs2(zero(T)::T)
+mr_empty(::typeof(identity), op::typeof(+), T) = zero(T)
+mr_empty(::typeof(abs), op::typeof(+), T) = abs(zero(T))
+mr_empty(::typeof(abs2), op::typeof(+), T) = abs2(zero(T))
+mr_empty(::typeof(identity), op::typeof(*), T) = one(T)
+mr_empty(::typeof(promote_sys_size), op, T) =
+    promote_sys_size(mr_empty(identity, op, T))
+mr_empty(::typeof(abs), op::typeof(scalarmax), T) = abs(zero(T))
+mr_empty(::typeof(abs2), op::typeof(scalarmax), T) = abs2(zero(T))
 mr_empty(::typeof(abs), op::typeof(max), T) = mr_empty(abs, scalarmax, T)
 mr_empty(::typeof(abs2), op::typeof(max), T) = mr_empty(abs2, scalarmax, T)
 mr_empty(f, op::typeof(&), T) = true
@@ -271,12 +256,12 @@ function _mapreduce(f, op, ::IndexLinear, A::AbstractArray{T}) where T
         return mr_empty(f, op, T)
     elseif n == 1
         @inbounds a1 = A[inds[1]]
-        return r_promote(op, f(a1))
+        return f(a1)
     elseif n < 16 # process short array here, avoid mapreduce_impl() compilation
         @inbounds i = inds[1]
         @inbounds a1 = A[i]
         @inbounds a2 = A[i+=1]
-        s = op(r_promote(op, f(a1)), r_promote(op, f(a2)))
+        s = op(f(a1), f(a2))
         while i < last(inds)
             @inbounds Ai = A[i+=1]
             s = op(s, f(Ai))
@@ -352,7 +337,7 @@ julia> sum(abs2, [2; 3; 4])
 29
 ```
 """
-sum(f::Callable, a) = mapreduce(f, +, a)
+sum(f::Callable, a) = mapreduce(promote_sys_size ∘ f, +, a)
 
 """
     sum(itr)
@@ -364,7 +349,7 @@ julia> sum(1:20)
 210
 ```
 """
-sum(a) = mapreduce(identity, +, a)
+sum(a) = mapreduce(promote_sys_size, +, a)
 sum(a::AbstractArray{Bool}) = countnz(a)
 
 
@@ -379,7 +364,7 @@ summation algorithm for additional accuracy.
 """
 function sum_kbn(A)
     T = _default_eltype(typeof(A))
-    c = r_promote(+, zero(T)::T)
+    c = promote_sys_size(zero(T)::T)
     i = start(A)
     if done(A, i)
         return c
@@ -410,7 +395,7 @@ julia> prod(abs2, [2; 3; 4])
 576
 ```
 """
-prod(f::Callable, a) = mapreduce(f, *, a)
+prod(f::Callable, a) = mapreduce(promote_sys_size ∘ f, *, a)
 
 """
     prod(itr)
@@ -422,7 +407,7 @@ julia> prod(1:20)
 2432902008176640000
 ```
 """
-prod(a) = mapreduce(identity, *, a)
+prod(a) = mapreduce(promote_sys_size, *, a)
 
 ## maximum & minimum
 

base/reducedim.jl

@@ -116,19 +116,22 @@ reducedim_init(f, op::typeof(|), A::AbstractArray, region) = reducedim_initarray
 
 # specialize to make initialization more efficient for common cases
 
-for (IT, RT) in ((CommonReduceResult, :(eltype(A))), (SmallSigned, :Int), (SmallUnsigned, :UInt))
-    T = Union{[AbstractArray{t} for t in uniontypes(IT)]..., [AbstractArray{Complex{t}} for t in uniontypes(IT)]...}
-    @eval begin
-        reducedim_init(f::typeof(identity), op::typeof(+), A::$T, region) =
-            reducedim_initarray(A, region, zero($RT))
-        reducedim_init(f::typeof(identity), op::typeof(*), A::$T, region) =
-            reducedim_initarray(A, region, one($RT))
-        reducedim_init(f::Union{typeof(abs),typeof(abs2)}, op::typeof(+), A::$T, region) =
-            reducedim_initarray(A, region, real(zero($RT)))
-        reducedim_init(f::Union{typeof(abs),typeof(abs2)}, op::typeof(*), A::$T, region) =
-            reducedim_initarray(A, region, real(one($RT)))
-    end
+let
+    const BitIntFloat = Union{BitInteger, Math.IEEEFloat}
+    const T = Union{
+        [AbstractArray{t} for t in uniontypes(BitIntFloat)]...,
+        [AbstractArray{Complex{t}} for t in uniontypes(BitIntFloat)]...}
+
+    global reducedim_init(f::typeof(identity), op::typeof(+), A::T, region) =
+        reducedim_initarray(A, region, zero(eltype(A)))
+    global reducedim_init(f::typeof(identity), op::typeof(*), A::T, region) =
+        reducedim_initarray(A, region, one(eltype(A)))
+    global reducedim_init(f::Union{typeof(abs),typeof(abs2)}, op::typeof(+), A::T, region) =
+        reducedim_initarray(A, region, real(zero(eltype(A))))
+    global reducedim_init(f::Union{typeof(abs),typeof(abs2)}, op::typeof(*), A::T, region) =
+        reducedim_initarray(A, region, real(one(eltype(A))))
 end
+
 reducedim_init(f::Union{typeof(identity),typeof(abs),typeof(abs2)}, op::typeof(+), A::AbstractArray{Bool}, region) =
     reducedim_initarray(A, region, 0)
 
@@ -574,14 +577,22 @@ any!(r, A)
 for (fname, op) in [(:sum, :+), (:prod, :*),
                     (:maximum, :scalarmax), (:minimum, :scalarmin),
                     (:all, :&), (:any, :|)]
+    function compose_pss(x)
+        if fname in [:sum, :prod]
+            :(promote_sys_size ∘ $x)
+        else
+            x
+        end
+    end
     fname! = Symbol(fname, '!')
     @eval begin
+        # TODO: check this
         $(fname!)(f::Function, r::AbstractArray, A::AbstractArray; init::Bool=true) =
             mapreducedim!(f, $(op), initarray!(r, $(op), init), A)
         $(fname!)(r::AbstractArray, A::AbstractArray; init::Bool=true) = $(fname!)(identity, r, A; init=init)
 
         $(fname)(f::Function, A::AbstractArray, region) =
-            mapreducedim(f, $(op), A, region)
+            mapreducedim($(compose_pss(:f)), $(op), A, region)
         $(fname)(A::AbstractArray, region) = $(fname)(identity, A, region)
     end
 end

base/tuple.jl

@@ -305,9 +305,13 @@ reverse(t::Tuple) = revargs(t...)
 
 # TODO: these definitions cannot yet be combined, since +(x...)
 # where x might be any tuple matches too many methods.
+# TODO: this is inconsistent with the regular sum in cases where the arguments
+# require size promotion to system size.
 sum(x::Tuple{Any, Vararg{Any}}) = +(x...)
 
 # NOTE: should remove, but often used on array sizes
+# TODO: this is inconsistent with the regular prod in cases where the arguments
+# require size promotion to system size.
 prod(x::Tuple{}) = 1
 prod(x::Tuple{Any, Vararg{Any}}) = *(x...)
 

test/reduce.jl

@@ -2,7 +2,7 @@
 
 # fold(l|r) & mapfold(l|r)
 @test foldl(+, Int64[]) === Int64(0) # In reference to issues #7465/#20144 (PR #20160)
-@test foldl(+, Int16[]) === Int(0)
+@test foldl(+, Int16[]) === Int16(0) # In reference to issues #21536
 @test foldl(-, 1:5) == -13
 @test foldl(-, 10, 1:5) == -5
 
@@ -19,7 +19,7 @@
 @test Base.mapfoldl((x)-> x ⊻ true, |, false, [true false true false false]) == true
 
 @test foldr(+, Int64[]) === Int64(0) # In reference to issue #20144 (PR #20160)
-@test foldr(+, Int16[]) === Int(0)
+@test foldr(+, Int16[]) === Int16(0) # In reference to issues #21536
 @test foldr(-, 1:5) == 3
 @test foldr(-, 10, 1:5) == -7
 @test foldr(+, [1]) == 1 # Issue #21493
@@ -29,7 +29,7 @@
 
 # reduce
 @test reduce(+, Int64[]) === Int64(0) # In reference to issue #20144 (PR #20160)
-@test reduce(+, Int16[]) === Int(0)
+@test reduce(+, Int16[]) === Int16(0) # In reference to issues #21536
 @test reduce((x,y)->"($x+$y)", 9:11) == "((9+10)+11)"
 @test reduce(max, [8 6 7 5 3 0 9]) == 9
 @test reduce(+, 1000, 1:5) == (1000 + 1 + 2 + 3 + 4 + 5)
@@ -69,12 +69,21 @@
       typeof(mapreduce(abs, +, Float32[10, 11, 12, 13]))
 
 # sum
+@testset "sums promote to at least machine size" begin
+    @testset for T in [Int8, Int16, Int32]
+        @test sum(T[]) === Int(0)
+    end
+    @testset for T in [UInt8, UInt16, UInt32]
+        @test sum(T[]) === UInt(0)
+    end
+    @testset for T in [Int, Int64, Int128, UInt, UInt64, UInt128,
+                       Float16, Float32, Float64]
+        @test sum(T[]) === T(0)
+    end
+    @test sum(BigInt[]) == big(0) && sum(BigInt[]) isa BigInt
+end
 
-@test sum(Int8[]) === Int(0)
-@test sum(Int[]) === Int(0)
-@test sum(Float64[]) === 0.0
-
-@test sum(Int8(3)) === Int8(3)
+@test sum(Int8(3)) === Int(3)
 @test sum(3) === 3
 @test sum(3.0) === 3.0