Is there something like broadcast_mapreduce?

General Usage
15

Just in case someone stumbles upon this thread and to answer my own question. Today I thought about the problem again and I found a nice and fast way to do this by using the mutating broadcast! on a custom AbstractArray which does the reduction upon calling setindex! Here is my final implementation of broadcast_reduce which has no allocations and is faster than the allocating version:

#Define abstract Array which consumes values on setindex!
type ArrayReducer{T,N,F} <: AbstractArray{T,N}
    v::T
    size::NTuple{N,Int}
    op::F
end
ArrayReducer{T,N}(size::NTuple{N},v0::T,op)=ArrayReducer{T,N,typeof(op)}(v0,size,op)
Base.setindex!(x::ArrayReducer,v,i...) = x.v=x.op(x.v,v)
Base.size(a::ArrayReducer)=a.size
get_stop(a::Base.OneTo)=a.stop

#Define the function
function broadcast_reduce(f,op,v0,A,Bs...)
    shape = Base.Broadcast.broadcast_indices(A,Bs...)
    reducer = ArrayReducer(get_stop.(shape),v0,op)
    broadcast!(f,reducer,A,Bs...)
    reducer.v
end

# First check if both give the same result
x1 = rand(200,5000,30)
x2 = rand(1,5000,1)
x3 = rand(1,1,1)

println(sum(x1.*x2.*x3))
println(broadcast_reduce(*,+,0.0,x1,x2,x3))

#Then benchmark
using BenchmarkTools
@btime sum(x1.*x2.*x3)
#  109.582 ms (57 allocations: 228.88 MiB)

@btime broadcast_reduce(*,+,0.0,x1,x2,x3)
# 35.359 ms (46 allocations: 1.81 KiB)
10 Likes