Skip to content

Commit

Permalink
Add divide-and-conquer _mapreduce
Browse files Browse the repository at this point in the history
  • Loading branch information
tkf committed Feb 29, 2020
1 parent 621d50a commit 17a1e2f
Showing 1 changed file with 22 additions and 0 deletions.
22 changes: 22 additions & 0 deletions src/mapreduce.jl
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,28 @@ end

@inline _mapreduce(args::Vararg{Any,N}) where N = _mapfoldl(args...)

@generated function _mapreduce(f, op, dims::Colon, init, ::Size{S}, a::StaticArray...) where {S}
function op_expr(idx)
if length(idx) == 1
i, = idx
tmp = [:(a[$j][$i]) for j 1:length(a)]
expr = :(f($(tmp...)))
if init === _InitialValue
return :(Base.reduce_first(op, $expr))
else
return :(op(init, $expr))
end
end
left = op_expr(idx[1:end÷2])
right = op_expr(idx[end÷2+1:end])
return :(op($left, $right))
end
return quote
@_inline_meta
@inbounds return $(op_expr(1:prod(S)))
end
end

@generated function _mapfoldl(f, op, dims::Colon, init, ::Size{S}, a::StaticArray...) where {S}
tmp = [:(a[$j][1]) for j 1:length(a)]
expr = :(f($(tmp...)))
Expand Down

0 comments on commit 17a1e2f

Please sign in to comment.