Why does an iterator with a vector as state variable allocate?

General Usage Performance
1

Here is an example of a phenomenon that I have noticed more than once: The following type C uses a Vector as state variable for iterate. This of course needs one allocation at the beginning of the iteration. However, there are allocations in each iteration step although no new Vector is created.

struct C
    v::Vector{Int}
end

Base.length(c::C) = prod(c.v)

function Base.iterate(c)
    w = c.v .- 1
    w, w
end

function Base.iterate(c, w)
    i = findfirst(!iszero, w)
    i === nothing && return nothing
    w[i] -= 1
    w[1:i-1] .= view(c.v, 1:i-1) .- 1   # EDIT: now using `view`
    w, w
end

Iterating over some c::C is like iterating over a Cartesian product of linear ranges:

julia> [copy(v) for v in C([2,3])]
6-element Vector{Vector{Int64}}:
 [1, 2]
 [0, 2]
 [1, 1]
 [0, 1]
 [1, 0]
 [0, 0]

Can one avoid the many allocations?

julia> c = C([2,3,4]); @b [sum(v) for v in $c]
513.120 ns (25 allocs: 1.047 KiB)

julia> c = C([2,3,4,5]); @b [sum(v) for v in $c]
2.200 μs (121 allocs: 4.875 KiB)
2

Slicing in Julia copies - you’ll want to use @view c.v[1:i-1].

You’ll probably only be able to half them, since the explicit copy in your array comprehension is necessary to make sure the individual iterations don’t share the same array.

1 Like
3

It’s not clear why you use a vector in the C struct. Your example can be written with tuples:

struct CC{N, T<:NTuple{N}}
    v::T
end

Base.length(c::CC) = prod(c.v)

@inline function Base.iterate(c::CC)
    w = c.v .- 1
    return w, w
end

@inline function Base.iterate(c::CC{N}, state) where N
    i = findfirst(!iszero, state)
    i === nothing && return nothing
    w = ntuple(Val(N)) do j
        j < i && return c.v[j]-1
        j == i && return state[i]-1
        return state[j]
    end
    return w, w
end

cc = CC((2,3,4)); @b [sum(v) for v in $cc]
cc = CC((2,3,4,5)); @b [sum(v) for v in $cc]

If you need a vector as output from the iterator, you can let iterate do a return [w...], w, but it might be cleaner to do that conversion when it’s needed.

4

The missing view was a mistake on my part. I’ve updated the OP. Even with view, there is an allocation for each iteration in the example with sum that doesn’t use copy. Why is that?

5

Because it is an example. My goal is not the best possible code for this specific case. I want to understand the reason for the allocations that arise here and in other (more complicated) examples, even when using view.

6

I think the allocations are occurring when the loop code is destructuring the returned tuple.

If I run the code:

function test()
    iter = C([2, 3, 4, 5])
    next = iterate(iter)
    while next !== nothing
        #(item, state) = next
        next = iterate(iter, state)
    end
end
@btime test()

I get only 4 allocations, but uncommenting the destructuring line results in the 122 allocations.

Under the hood, I suspect that the generator code for for v in $c is doing something very similar. Not sure if there’s a way around it.

7

Without uncommenting the destructuring line the code doesn’t work because state is undefined.

In any case, I agree that it probably has something to do with the code that Julia adds when lowering a loop. But why are there allocations only for some types of state variables and not for others?

1 Like
8

The @inline annotations on iterate are important (just as they are in sguare’s code). They solve the return value allocating problem.

2 Likes
9

Great, thanks!