Help to track down a type instability

General Usage
, ,
1

Hi there,

I have a piece of code that is sometimes type-stable, sometimes not, and can’t understand why… I’ve been scratching my head on this for a while, a fresh look would be very much appreciated…

I’ve boiled it down to the following MWE: given

  1. args : a tuple that contains Ints or <:AbstractVector{Int};
  2. inds : a tuple that also contains only Ints or <:AbstractVector{Int}, whose length that matches the number of <:AbstractVector{Int} elements of args.

The following code should apply getindex to each non-Int element of args with the corresponding inds.

using Base: tail, OneTo

getindices(res, ::Tuple{}, ::Tuple{}) = res

getindices(res, args::Tuple{Int,Vararg}, inds) =
    getindices((res..., first(args)), tail(args), inds)

getindices(res, args::Tuple{AbstractVector{Int},Vararg}, inds) =
    getindices((res..., getindex(first(args), first(inds))), tail(args), tail(inds))

Sometimes it’s type stable

julia> args = (1, 2, OneTo(3), 4:5);

julia> inds = (2:3, 1);

julia> @code_warntype getindices((), args, inds)
MethodInstance for getindices(::Tuple{}, ::Tuple{Int64, Int64, OneTo{Int64}, UnitRange{Int64}}, ::Tuple{UnitRange{Int64}, Int64})
  from getindices(res, args::Tuple{Int64, Vararg}, inds) @ Main ~/.julia/dev/minimal.jl:5
Arguments
  #self#::Core.Const(getindices)
  res::Core.Const(())
  args::Tuple{Int64, Int64, OneTo{Int64}, UnitRange{Int64}}
  inds::Tuple{UnitRange{Int64}, Int64}
Body::Tuple{Int64, Int64, UnitRange{Int64}, Int64}
1 ─ %1 = Main.first(args)::Int64
│   %2 = Core.tuple(%1)::Tuple{Int64}
│   %3 = Core._apply_iterate(Base.iterate, Core.tuple, res, %2)::Tuple{Int64}
│   %4 = Main.tail(args)::Tuple{Int64, OneTo{Int64}, UnitRange{Int64}}
│   %5 = Main.getindices(%3, %4, inds)::Tuple{Int64, Int64, UnitRange{Int64}, Int64}
└──      return %5
args = (1, 2, 3, OneTo(5), -1, OneTo(6))
inds = (1, 2)

But other times it’s not

julia> args = (1, 2, 3, OneTo(5), -1, OneTo(6));

julia> inds = (1, 2:3);

julia> @code_warntype getindices((), args, inds)
MethodInstance for getindices(::Tuple{}, ::Tuple{Int64, Int64, Int64, OneTo{Int64}, Int64, OneTo{Int64}}, ::Tuple{Int64, UnitRange{Int64}})
  from getindices(res, args::Tuple{Int64, Vararg}, inds) @ Main ~/.julia/dev/ProductSandbox/minimal.jl:5
Arguments
  #self#::Core.Const(getindices)
  res::Core.Const(())
  args::Tuple{Int64, Int64, Int64, OneTo{Int64}, Int64, OneTo{Int64}}
  inds::Tuple{Int64, UnitRange{Int64}}
Body::Tuple
1 ─ %1 = Main.first(args)::Int64
│   %2 = Core.tuple(%1)::Tuple{Int64}
│   %3 = Core._apply_iterate(Base.iterate, Core.tuple, res, %2)::Tuple{Int64}
│   %4 = Main.tail(args)::Tuple{Int64, Int64, OneTo{Int64}, Int64, OneTo{Int64}}
│   %5 = Main.getindices(%3, %4, inds)::Tuple
└──      return %5

I’ve tried different versions of Julia (include 1.9.0-beta2), the results are the same… Thanks in advance for any tip to help me solve this!

2

%4 has 3 elements in the stable one and 5 in the second, so the second has deeper recursion. Could be that the compiler gives up in the second one because of that. Not sure how to verify that though. I personally would probably use a generated function for such a thing, also because I find the intent easier to express than with recursion (case in point, I don’t get your code at all :wink: )

3

Hi @jules thanks for your help, I think you’re right about the length of the input Tuple: when they are small enough (5 or less), there is no type instability… A @generated function would probably work, but it shouldn’t be needed: the MWE simply applies getindex when possible, otherwise ignores leaves the element untouched.

So two questions your point raises are: what’s happening with the compiler, and is there a way to force it to stay type-stable without using @generated?

Thanks again for your help!

4

Another question: as it is defined now getindices returns a Union of Ints and Vector{Int}s. Is this the goal, or is the goal to have a flattened output of a Tuple of Ints? i.e. getindices((1,2),(1:5,5),(3:4,)) == (1,2,3:4,5) like it is now, or == (1,2,3,4,5) ?
If the latter, then it may influence the type stability exploration here (and the function needs to be changed).

5

Hi @Dan a Tuple of Union{Int,Vector{Int}} is indeed the goal.

Following the comment from @jules I came up with the following implementation :

@generated function getindices(vals, inds)
    expr = Expr(:call, :tuple)
    j = 0
    for (i, T) in enumerate(vals.parameters)
        if T <: Integer
            push!(expr.args, :(vals[$i]))
        elseif T <: AbstractVector{<:Integer}
            j += 1
            push!(expr.args, :(getindex(vals[$i], inds[$j])))
        end
    end
    expr
end

It seems to do the trick : no type instability raised in the previous examples, e.g.

julia> args = (1, 2, Base.OneTo(3), 4:5, Base.OneTo(3), 2, 2:5);

julia> inds = (2:3, 1, 2:3, 1:2);

julia> @code_warntype getindices(args, inds)
MethodInstance for getindices(::Tuple{Int64, Int64, Base.OneTo{Int64}, UnitRange{Int64}, Base.OneTo{Int64}, Int64, UnitRange{Int64}}, ::Tuple{UnitRange{Int64}, Int64, UnitRange{Int64}, UnitRange{Int64}})
  from getindices(vals, inds) @ Main ~/.julia/dev/minimal.jl:30
Arguments
  #self#::Core.Const(getindices)
  vals::Tuple{Int64, Int64, Base.OneTo{Int64}, UnitRange{Int64}, Base.OneTo{Int64}, Int64, UnitRange{Int64}}
  inds::Tuple{UnitRange{Int64}, Int64, UnitRange{Int64}, UnitRange{Int64}}
Body::Tuple{Int64, Int64, UnitRange{Int64}, Int64, UnitRange{Int64}, Int64, UnitRange{Int64}}
1 ─ %1  = Base.getindex(vals, 1)::Int64
│   %2  = Base.getindex(vals, 2)::Int64
│   %3  = Base.getindex(vals, 3)::Base.OneTo{Int64}
│   %4  = Base.getindex(inds, 1)::UnitRange{Int64}
│   %5  = Main.getindex(%3, %4)::UnitRange{Int64}
│   %6  = Base.getindex(vals, 4)::UnitRange{Int64}
│   %7  = Base.getindex(inds, 2)::Int64
│   %8  = Main.getindex(%6, %7)::Int64
│   %9  = Base.getindex(vals, 5)::Base.OneTo{Int64}
│   %10 = Base.getindex(inds, 3)::UnitRange{Int64}
│   %11 = Main.getindex(%9, %10)::UnitRange{Int64}
│   %12 = Base.getindex(vals, 6)::Int64
│   %13 = Base.getindex(vals, 7)::UnitRange{Int64}
│   %14 = Base.getindex(inds, 4)::UnitRange{Int64}
│   %15 = Main.getindex(%13, %14)::UnitRange{Int64}
│   %16 = Main.tuple(%1, %2, %5, %8, %11, %12, %15)::Tuple{Int64, Int64, UnitRange{Int64}, Int64, UnitRange{Int64}, Int64, UnitRange{Int64}}
└──       return %16

However not being familiar with @generated functions (my basic understanding was that they should be avoided as much as possible, hence my initial implementation), I would appreciate if someone (@jules ?)
could comment on the implementation (how to access the Tuple’s parameters in particular. Thanks!

1 Like
6

You can do fieldtypes(vals) to get the types. Other than that, I think it’s fine. You’re using very simple logic in your generated function, so I don’t think you need to worry about the more complex reasons why one would want to avoid generated functions. You can search a bit in this forum, I know there have been long threads where the pros and cons were discussed. I find them quite useful for cases pretty similar to yours, usually where I really want to make sure that inference succeeds for tuple types.