Improve broadcast cache performance

broadcast_cache_perf.jl

let broadcast_cache = Dict{Function,Dict{Int,Dict{Int,Function}}}()
global broadcast!
function broadcast!(f::Function, B, As::Union(Array,BitArray)...)
    nd = ndims(B)
    narrays = length(As)
    
    idx_f = Base.ht_keyindex2(broadcast_cache, f)
    if idx_f < 0
        idx_f = -idx_f
        Base._setindex!(broadcast_cache, Dict{Int,Dict{Int,Function}}(), f, idx_f)
    end
    cache_f = broadcast_cache.vals[idx_f]

    idx_f_nd = Base.ht_keyindex2(cache_f, nd)
    if idx_f_nd < 0
        idx_f_nd = -idx_f_nd
        Base._setindex!(cache_f, Dict{Int,Function}(), nd, idx_f_nd)
    end
    cache_f_nd = cache_f.vals[idx_f_nd]

    idx_f_nd_na = Base.ht_keyindex2(cache_f_nd, narrays)
    if idx_f_nd_na < 0
        idx_f_nd_na = -idx_f_nd_na
        func = gen_broadcast_function(gen_broadcast_body_iter, nd, narrays, f)
        Base._setindex!(cache_f_nd, func, narrays, idx_f_nd_na)
    else
        func = cache_f_nd.vals[idx_f_nd_na]
    end

    func(B, As...)
    B
end
end  # let broadcast_cache

broadcast_cache_perf_new.jl

macro get_cached1(cache, key, default)
    quote
        idx = Base.ht_keyindex2($(esc(cache)), $(esc(key)))
        if idx < 0
            idx = -idx
            Base._setindex!($(esc(cache)), $(esc(default)), $(esc(key)), idx)
        end
        $(esc(cache)).vals[idx]
    end
end

macro get_cached(cache, f, nd, narrays, genfunc)
    quote
        cache_f = @get_cached1($(esc(cache)), $(esc(f)), Dict{Int,Dict{Int,Function}}())
        cache_f_nd = @get_cached1(cache_f, $(esc(nd)), Dict{Int,Function}())
        func = @get_cached1(cache_f_nd, $(esc(narrays)), $(esc(genfunc)))
    end
end

let broadcast_cache = Dict{Function,Dict{Int,Dict{Int,Function}}}()
global broadcast!
function broadcast!(f::Function, B, As::Union(Array,BitArray)...)
    nd = ndims(B)
    narrays = length(As)
    func = @get_cached(broadcast_cache, f, nd, narrays,
        gen_broadcast_function(gen_broadcast_body_iter, nd, narrays, f))
    func(B, As...)
    B
end
end  # let broadcast_cache