IterTools.jl Partition type stability?

General Usage Performance
1

Trying some code like:

function test_loop(event_times)
    out = 0.0
    for (i, event_time) in partition(event_times, 2, 1)
        out += i + event_time
    end
    return out
end

@code_warntype test_loop([1.1, 1.8, 2.1, 3.2, 5.1])

And I get lots of scary highlighting. Am I doing something iffy? I just want a nice sliding window like function. Looked at the IterTools code a bit but don’t fully understand what might be causing the issue.

Thanks so much

2

partition uses the runtime value of the second argument as one of the parameters in the resulting type, which is type unstable.

Note that for constant values this can be fixed by inlining the constructor and exploiting constant propagation as shown by:

julia> @code_warntype f(event_times)
Body::IterTools.Partition{Array{Float64,1},_1} where _1
1 1 ─ %1 = invoke Main.partition(_2::Array{Float64,1}, 2::Int64, 1::Int64)::IterTools.Partition{Array{Float64,1},_1} where _1            │
  └──      return %1                                                                                                                     │


julia> @eval IterTools begin # only adding @inline to the constructor, rest is the same
       @inline function partition(xs::I, n::Int, step::Int) where I
           if step < 1
               throw(ArgumentError("Partition step must be at least 1."))
           end

           if n < 1
               throw(ArgumentError("Partition size n must be at least 1"))
           end

           Partition{I, n}(xs, step)
       end
       end

julia> @code_warntype f(event_times)
Body::IterTools.Partition{Array{Float64,1},2}
1 1 ─ %1 = %new(IterTools.Partition{Array{Float64,1},2}, event_times, 1)::IterTools.Partition{Array{Float64,1},2}            │╻╷ partition
  └──      return %1

With this, the original example infers fine as well.

2 Likes
3

Oh man this is above my level of understanding, but it works sweet. Is this something to try to push into IterTools itself? Or would there be a reason not to do this?

4

Inlining too many things can bloat the code, but if it is a common thing to call partition with a literal for n then it probably makes sense since having things infer is nice. There are other possible solutions that involves changing the API for how partition is called or how the Partition struct is defined but this seems like a pretty simple solution, at least for this use-case.

1 Like