Merge pull request #258 from JuliaGPU/tb/mapreduce_vararg

Specialize varargs version of mapreduce.

Author: maleadt

src/host/mapreduce.jl

@@ -2,7 +2,7 @@
 
 # GPUArrays' mapreduce methods build on `Base.mapreducedim!`, but with an additional
 # argument `init` value to avoid eager initialization of `R` (if set to something).
-mapreducedim!(f, op, R::AbstractGPUArray, A::AbstractArray, init=nothing) = error("Not implemented") # COV_EXCL_LINE
+mapreducedim!(f, op, R::AbstractGPUArray, As::AbstractArray...; init=nothing) = error("Not implemented") # COV_EXCL_LINE
 Base.mapreducedim!(f, op, R::AbstractGPUArray, A::AbstractArray) = mapreducedim!(f, op, R, A)
 
 neutral_element(op, T) =
@@ -18,11 +18,11 @@ neutral_element(::typeof(Base.mul_prod), T) = one(T)
 neutral_element(::typeof(Base.min), T) = typemax(T)
 neutral_element(::typeof(Base.max), T) = typemin(T)
 
-function Base.mapreduce(f, op, A::AbstractGPUArray; dims=:, init=nothing)
+function Base.mapreduce(f, op, As::AbstractGPUArray...; dims=:, init=nothing)
     # figure out the destination container type by looking at the initializer element,
     # or by relying on inference to reason through the map and reduce functions.
     if init === nothing
-        ET = Base.promote_op(f, eltype(A))
+        ET = Base.promote_op(f, map(eltype, As)...)
         ET = Base.promote_op(op, ET, ET)
         (ET === Union{} || ET === Any) &&
             error("mapreduce cannot figure the output element type, please pass an explicit init value")
@@ -32,10 +32,14 @@ function Base.mapreduce(f, op, A::AbstractGPUArray; dims=:, init=nothing)
         ET = typeof(init)
     end
 
+    # TODO: Broadcast-semantics after JuliaLang-julia#31020
+    A = first(As)
+    all(B -> size(A) == size(B), As) || throw(DimensionMismatch("dimensions of containers must be identical"))
+
     sz = size(A)
     red = ntuple(i->(dims==Colon() || i in dims) ? 1 : sz[i], ndims(A))
     R = similar(A, ET, red)
-    mapreducedim!(f, op, R, A, init)
+    mapreducedim!(f, op, R, As...; init=init)
 
     if dims==Colon()
         @allowscalar R[]

src/reference.jl

@@ -295,11 +295,11 @@ Adapt.adapt_storage(::Adaptor, x::JLArray{T,N}) where {T,N} =
 GPUArrays.unsafe_reinterpret(::Type{T}, A::JLArray, size::Tuple) where T =
     reshape(reinterpret(T, A.data), size)
 
-function GPUArrays.mapreducedim!(f, op, R::JLArray, A::AbstractArray, init=nothing)
+function GPUArrays.mapreducedim!(f, op, R::JLArray, As::AbstractArray...; init=nothing)
     if init !== nothing
         fill!(R, init)
     end
-    @allowscalar Base.mapreducedim!(f, op, R.data, A)
+    @allowscalar Base.reducedim!(op, R.data, map(f, As...))
 end
 
 end

test/testsuite/mapreduce.jl

@@ -136,6 +136,11 @@ function test_mapreduce(AT)
                                       AT, rand(range, N, N))
                         @test compare(x->mapreduce(_addone, +, x; dims = 2, init = _zero),
                                       AT, rand(range, N, N))
+
+                        @test compare(x->mapreduce(+, +, x; dims = 2),
+                                      AT, rand(range, N, N), rand(range, N, N))
+                        @test compare(x->mapreduce(+, +, x; dims = 2, init = _zero),
+                                      AT, rand(range, N, N). rand(range, N, N))
                     end
                 end
                 @testset "sum maximum minimum prod" begin