Faster multi-arg mapreduce

[oscardssmith]

Partially addresses #38558
If there's an easy way to make this handle the kwarkgs, let me know.

[oscardssmith]

Also does anyone have a good idea for how to test for O(1) allocations?

[Keno]

Also does anyone have a good idea for how to test for O(1) allocations?

Add is as a benchmark, so nanosoldier will track it?

[simeonschaub]

Also does anyone have a good idea for how to test for O(1) allocations?

I think in this case just adding an explicit @allocated test (after you call the function once before to precompile) should be fine.

[oscardssmith]

I'm reverting to the first version of this that I made. It doesn't handle keword arguments, but I think it's probably good to merge as is since it's a strict improvement compared to master.

[StefanKarpinski]

Bump?

[oscardssmith]

good call. I'd forgotten about this PR.

[mcabbott]

Should this allow mapreduce(+, +, -3:3, -3:3, init=1) to take the zip path?

[oscardssmith]

Yes it should. Is there an easy way to do that?

[mcabbott]

I think the only two legal ones are dims, init, but am not sure if init = _InitialValue() is understood by all the relevant methods. If not, then this would work:

function mapreduce(f, op, A::AbstractArrayOrBroadcasted...; dims=:, kw...)
    dims==(:) && return mapreduce(A->f(A...), op, zip(A...); kw...)
    ...

[mcabbott]

Last commit needs ; kw...) -> ; dims=:, kw...), I think.

[simeonschaub]

LGTM

[KristofferC]

As a note, I think in PRs like this it is great if the first post contains a benchmark that shows the before/after performance and can easily be run by someone looking at the PR.

[simeonschaub]

The failure here is real. Locally, replacing A -> f(A...) with Base.splat(f) and then defining

mapreduce_empty(f::typeof(splat(+)).name.wrapper, op, ::Type{T}) where {T<:Tuple} =
    reduce_empty(op, Core.Compiler.return_type(op, T))

worked as a quick fix. I think for a proper solution, we should make Base.splat return a functor of type Base.Splat though.