Broadcasting mapreducedim!

gistfile1.jl

using Base.Cartesian, Base.Test
using Base: check_reducdims, evaluate, @get!, mapreducedim!

function gen_mapreducedim_function(nd::Int, narrays::Int)
    As = [symbol("A_"*string(i)) for i = 1:narrays]
    @eval begin
        local _F_
        function _F_(f, op, R, $(As...))
            @nextract $nd sz d->($(narrays > 1 ? :(@ncall $narrays max k->size(A_k, d)) : :(size(A_1, d))))
            @nloops($nd, i, d->1:sz_d,
                d->(r_d = ifelse(size(R, d) == 1, 1, i_d);
                    @nexprs $narrays k->(j_d_k = ifelse(size(A_k, d) == 1, 1, i_d))), # pre
                begin # body
                    @nexprs $narrays k->(@inbounds v_k = @nref $nd A_k d->j_d_k)
                    @inbounds (@nref $nd R r) = evaluate(op, (@nref $nd R r), (@ncall $narrays evaluate f v))
                end)
            R
        end
        _F_
    end
end

@eval let cache = Dict{Int,Dict{Int,Function}}()
    function Base.mapreducedim!(f, op, R::AbstractArray, As::AbstractArray...)
        nd = ndims(R)
        narrays = length(As)

        cache_na = @get! cache    narrays Dict{Int,Function}()
        func     = @get! cache_na nd      gen_mapreducedim_function(nd, narrays)

        func(f, op, R, As...)
        R
    end
end

immutable Abs2MinusFun; end
Base.evaluate(::Abs2MinusFun, x, y) = abs2(x - y)

function var{T}(A::AbstractArray{T}, region)
    R = Base.reducedim_initarray(A, region, zero(Base.momenttype(T)))
    scale!(mapreducedim!(Abs2MinusFun(), Base.AddFun(), R, A, mean(A, region)), 1/(div(length(A), length(R)) - 1))
end