How to Make Tuple of Tuple generation Type Stable

General Usage Performance
1

I have code that needs to generate lists of indices. Cutting away all the fuzz, it boils down to

using JET

function foobar(K::Int, t::NTuple{N,T}) where {N,T}
    return [t[1:K]]
end

function foo(::Val{N}, k::Int) where {N}
    t = ntuple(i -> rand(Int) + k, N)
    ntuple(i -> foobar(i, t), N)
end

foo(Val(10), 3)
rep = @report_opt foo(Val(10), 3)
print(rep)

Output

═════ 1 possible error found ═════
┌ foo(::Val{10}, k::Int64) @ Main /Users/NENNIA/.julia/dev/TTApproximations/play.jl:54
│┌ ntuple(f::var"#62#64"{NTuple{10, Int64}}, n::Int64) @ Base ./ntuple.jl:19
││┌ (::var"#62#64"{NTuple{10, Int64}})(i::Int64) @ Main /Users/NENNIA/.julia/dev/TTApproximations/play.jl:54
│││┌ foobar(K::Int64, t::NTuple{10, Int64}) @ Main /Users/NENNIA/.julia/dev/TTApproximations/play.jl:49
││││ runtime dispatch detected: Base.vect(%189::Tuple{Vararg{Int64}})::Vector
│││└────────────────────

Calculating by hand, the return type of foo should be Tuple{Vec{NTuple{1, Int}}, Vec{NTuple{2, Int}}, ..., Vec{NTuple{N, Int}}}.

How do I make Julia understand that?

2

Here you go:

@inline function foobar(K::Int, t::NTuple{N,T}) where {N,T}
    return [t[1:K]]
end

function foo(::Val{N}, k::Int) where {N}
    t = ntuple(i -> rand(Int) + k, N)
    ntuple((@inline function(i); return foobar(i, t); end), N)
end

The difference is that I’ve forced inlining of all functions.

2 Likes