Type instability when indexing

Usage Performance
,
#1

I’m trying to understand why type instability appears for certain indexing operations. Suppose I define three functions (MWE for a longer code where type parameter d is used to index into part of a tuple/array)

struct bar{d} end
foo1(a::bar{d},x) where {d} = x[1:d]
foo2(a::bar{d},x) where {d} = x[2:d+1]
foo3(a::bar{d},x) where {d} = ntuple(i -> x[i+1], d)

If I check if the output type can be inferred, then I get

using Test
@inferred foo1(bar{2}(),(1,2,3,4)) # type stable
@inferred foo2(bar{2}(),(1,2,3,4)) # not type stable
@inferred foo3(bar{2}(),(1,2,3,4)) # type stable

Can someone help me with the logic of why the second one isn’t type stable but the first and last are?

#2

It wasn’t clear to me that range indexing for tuples would ever be type stable with constant propagation, so I had a look with @which getindex((1, 2), 1:2) and the method in Base is special cased in a tricky, but not fool-proof, way. If you work hard enough, you might be able to make your own fool-proof method with recursion. A good start would be wrapping the range start and ends with Val to trap them in the type domain. It would be hard on the compiler, and probably slow for very long tuples.

1 Like
#3

The last example will also become type-unstable if you increase d (and the length of the tuple).

julia> @code_warntype foo3(bar{11}(),(1,2,3,4,5,6,7,8,9,0,1,2))
Variables
  #self#::Core.Const(foo3)
  a::Core.Const(bar{11}())
  x::NTuple{12, Int64}
  #1::var"#1#2"{NTuple{12, Int64}}

Body::Tuple{Vararg{Int64, N} where N}
1 ─ %1 = Main.:(var"#1#2")::Core.Const(var"#1#2")
│   %2 = Core.typeof(x)::Core.Const(NTuple{12, Int64})
│   %3 = Core.apply_type(%1, %2)::Core.Const(var"#1#2"{NTuple{12, Int64}})
│        (#1 = %new(%3, x))
│   %5 = #1::var"#1#2"{NTuple{12, Int64}}
│   %6 = Main.ntuple(%5, $(Expr(:static_parameter, 1)))::Tuple{Vararg{Int64, N} where N}
└──      return %6

The first example is type-stable while the other one is not since it’s handled differently in getindex, see julia/range.jl at 399f8ba175b1add5f8b8286097afd2deaada0a41 · JuliaLang/julia · GitHub. Whenever the first and/or last index of the tuple are included in the range, there is a special case. The general case is just more demanding for the compiler. As @bramtayl wrote above, it might be possible to write another version based on recursion (similar to Base.front and Base.tail).

1 Like
#4

@bramtayl and @ranocha - thanks for the help! To make sure I’m not misunderstanding - are the first and last functions are type stable basically because the implementations of getindex and ntuple manually “unroll” (using the term loosely) the range indexing?

#5

My understanding is that the compiler simply is not able to perform computations such as d+1 with type parameters d. This is also why something like ComputedFieldTypes.jl exists.

#6

I think that’s not completely true. For example,

julia> @code_warntype foo2(bar{3}(),(1,2,3,4))
Variables
  #self#::Core.Const(foo2)
  a::Core.Const(bar{3}())
  x::NTuple{4, Int64}

Body::Tuple{Int64, Int64, Int64}
1 ─ %1 = ($(Expr(:static_parameter, 1)) + 1)::Core.Const(4)
│   %2 = (2:%1)::Core.Const(2:4)
│   %3 = Base.getindex(x, %2)::Tuple{Int64, Int64, Int64}
└──      return %3

is type-stable (since it hits another special case in getindex).

#7

Yes, that’s basically my understanding (although there is no “manual” unrolling, but the compiler is able to unroll everything based on the recursive implementations used internally).

1 Like
#8

I don’t think this means the compiler is calculating d+1 in general. Otherwise, foo2(a::bar{d},x) where {d} = x[2:d+1] would Just Work.

#9

Let me put it this way: the compiler’s ability to compute with type parameters is very limited.