Work around splatting to avoid unecessary allocations

General Usage Performance
1

Hello everyone,

I am trying to call Base.Iterators.product on a number of arguments which is not fixed in advance, i.e. to call the following function on integer vectors of arbitrary length.

function indices_below(Λ_max::Union{AbstractArray{U, 1}, Tuple}) where U <: Integer
    return Base.Iterators.product((0:Λi_max for Λi_max in Λ_max)...)
end

This function does the job, but it performs some allocations:

using BenchmarkTools
@btime indices_below($(5,4,6))
  403.880 ns (9 allocations: 480 bytes)
Base.Iterators.ProductIterator{Tuple{UnitRange{Int64},UnitRange{Int64},UnitRange{Int64}}}((0:5, 0:4, 0:6))

However, calling Base.Iterators.product directly on the desired UnitRanges produces no allocation:


@btime Base.Iterators.product(0:5, 0:4, 0:6)
  1.888 ns (0 allocations: 0 bytes)
Base.Iterators.ProductIterator{Tuple{UnitRange{Int64},UnitRange{Int64},UnitRange{Int64}}}((0:5, 0:4, 0:6))

Do you have any suggestion on how to improve the performance of this function ?

Many thanks

1 Like
2

map seems to work:

julia> f(x...) = Base.Iterators.product(map(N -> 0:N, x)...)
f (generic function with 1 method)
julia> @btime f(5,4,6)
  0.017 ns (0 allocations: 0 bytes)
3 Likes
3

Solved, many thanks for looking into it !

I had to remove the slurping ... for argument x to get the same interface as my original function, but the solution works.

Any idea why the generator allocates but map does not ?

Actually, leaving the type constraint on the arguments of the function produces 2 allocations, any idea why ?

f2(x::Union{AbstractArray{U, 1}, Tuple}) where U <: Integer = Base.Iterators.product(map(N -> 0:N, x)...)
f2 (generic function with 1 method)

@btime f2($(5,4,6))
  23.354 ns (2 allocations: 96 bytes)
Base.Iterators.ProductIterator{Tuple{UnitRange{Int64},UnitRange{Int64},UnitRange{Int64}}}((0:5, 0:4, 0:6))