Allocations when manipulating an iterator

General Usage Performance
1

I distilled a significant performance issue in JuMP 0.19-alpha (https://github.com/JuliaOpt/JuMP.jl/issues/1403#issuecomment-437630140) to the following standalone example:

using BenchmarkTools

struct FlippedDictIterator{K,V}
    d::Dict{K,V}
end

reorder_iterator(::Nothing) = nothing
reorder_iterator(p::Pair, state::Int) = ((p.second, p.first), state)
reorder_iterator(x) = reorder_iterator(x...)

Base.iterate(fi::FlippedDictIterator) = reorder_iterator(iterate(fi.d))
Base.iterate(fi::FlippedDictIterator, state) = reorder_iterator(iterate(fi.d, state))

# The behavior isn't affected by these 4 lines, but I thought they could help.
Base.IteratorSize(fi::FlippedDictIterator) = Base.IteratorSize(fi.d)
Base.length(fi::FlippedDictIterator) = length(fi.d)
Base.IteratorEltype(fi::FlippedDictIterator) = Base.IteratorEltype(fi.d)
Base.eltype(fi::FlippedDictIterator{K,V}) where {K,V} = Tuple{V, K} 

function loop1(d)
    s = 0
    for (key, val) in d
        s += key * length(val)
    end
    return s
end

function loop2(d)
    s = 0
    for (val, key) in FlippedDictIterator(d)
        s += key * length(val)
    end
    return s
end

d = Dict{Int, String}()
for i in 1:1000
    d[i] = "test"
end

@show loop1(d)
@show loop2(d)

@btime loop1($d)
@btime loop2($d)

Yields (on Julia 1.0.1):

loop1(d) = 2002000
loop2(d) = 2002000
  14.323 μs (0 allocations: 0 bytes)
  20.913 μs (2000 allocations: 62.50 KiB)

Why are there two allocations per call to iterate(fi::FlippedDictIterator, state)?

3 Likes
2

The optimiser has to chase nothing through too many levels: try

julia> Base.iterate(fi::FlippedDictIterator, state) = let ϕ = iterate(fi.d, state); if !( ϕ === nothing ) return reverse(ϕ[1]), ϕ[2] else return nothing end end

julia> Base.iterate(fi::FlippedDictIterator) = let ϕ = iterate(fi.d); if !( ϕ === nothing ) return reverse(ϕ[1]), ϕ[2] else return nothing end end
4 Likes
3

That worked. Thanks!

4

Can you please describe how you found this out? It would be useful to learn to do this systematically.

5

systematically

Uh, unfortunately that was more intuition than science, but the gist is to bail early when an iterator returns nothing and not to delay handling of nothings to later instances.

6

Currently, my way of dealing with these things is to make small functions, stare at @code_warntype, and try patterns I have seen in Base; I was hoping I would learn something more organized :wink:

1 Like