I’m trying to write an iteration for my type. The iteration needs to keep a Vector and an Int in its state. I cannot get rid of an allocation.
I condensed the problem into the following MWE:
function iter((i,v))
v[i] += 1
i == length(v) ? nothing : (i+1,v)
end
which gives:
julia> state=(1,[1,2]); @btime iter($state);
3.805 ns (1 allocation: 32 bytes)
Is there any way to get rid of this allocation? Thanks a lot for any insight!
I don’t understand that allocation either, but I have a similar iterator, which does not allocate. Here is a simplified version of it:
julia> struct Chunk{T<:AbstractArray}
x::T
n::Int
end
julia> function iterate(c::Chunk, state=nothing)
if isnothing(state)
return (c, 1)
elseif state < c.n
return (c, state + 1)
end
return nothing
end
iterate (generic function with 2 methods)
julia> @btime iterate($c)
1.468 ns (0 allocations: 0 bytes)
(Chunk{Vector{Int64}}([1, 2], 1), 1)
Thank you for the reply!
It seems that it does allocate after the first iteration for me though:
julia> c = Chunk([1,2,3], 3)
Chunk{Vector{Int64}}([1, 2, 3], 3)
julia> @btime iterate($c)
1.893 ns (0 allocations: 0 bytes)
(Chunk{Vector{Int64}}([1, 2, 3], 3), 1)
julia> @btime iterate($c,1)
3.595 ns (1 allocation: 32 bytes)
(Chunk{Vector{Int64}}([1, 2, 3], 3), 2)
I don´t know. Seems to be some artifact of benchmarking? In a “real” test it does not allocate:
julia> using ChunkSplitters
julia> x = rand(1000);
julia> function test(x)
s = 0.0
for inds in chunks(x; n = 100)
for i in inds
s += x[i]
end
end
return s
end
test (generic function with 1 method)
julia> test(x)
509.1937895395777
julia> @btime test($x)
716.511 ns (0 allocations: 0 bytes)
509.1937895395777
The exact implementation of that is here: ChunkSplitters.jl/src/ChunkSplitters.jl at 6f7ea5e8baf4b0de84f6d04e4320959c4c3f6f6a · m3g/ChunkSplitters.jl · GitHub
One difference, relative to the benchmark, is that in the “real” test the return value of the function to the REPL is always of the same type, while in the small example it can return nothing to the REPL. That may cause an artificial allocation, maybe.
You’re right! Unfortunately, in my case the allocation is real. I must have condensed the wrong MWE. Thanks again!
You could try to slap a @inline in front of your iteration functions. Sometimes that fixes allocations (which might stem from tuple packing/unpacking or so…)
This is what I thought at first, but iterate(x::Vector) does not allocate, even if it can return nothing. So it seems that inlining definitely has also to do with it…
