Why selectdim is type instable

Hi, I just found that selectdim does not infer the concrete SubArray:

julia> using InteractiveUtils

julia> a = rand(5, 5);

julia> @code_warntype selectdim(a, 2, 1)
Body::SubArray
1 ─ %1  = (Base.arraysize)(A, 1)::Int64
│   %2  = (Base.arraysize)(A, 2)::Int64
│   %3  = (Base.slt_int)(%1, 0)::Bool
│   %4  = (Base.ifelse)(%3, 0, %1)::Int64
│   %5  = %new(Base.OneTo{Int64}, %4)::Base.OneTo{Int64}
│   %6  = (Base.slt_int)(%2, 0)::Bool
│   %7  = (Base.ifelse)(%6, 0, %2)::Int64
│   %8  = %new(Base.OneTo{Int64}, %7)::Base.OneTo{Int64}
│   %9  = %new(Base.Slice{Base.OneTo{Int64}}, %5)::Base.Slice{Base.OneTo{Int64}}
│   %10 = %new(Base.Slice{Base.OneTo{Int64}}, %8)::Base.Slice{Base.OneTo{Int64}}
│   %11 = (d === 1)::Bool
│   %12 = (Base.ifelse)(%11, i, %9)::Union{Slice{OneTo{Int64}}, Int64}
│   %13 = (Base.sub_int)(d, 1)::Int64
│   %14 = (%13 === 1)::Bool
│   %15 = (Base.ifelse)(%14, i, %10)::Union{Slice{OneTo{Int64}}, Int64}
│   %16 = (Core.tuple)(%12, %15)::Tuple{Union{Slice{OneTo{Int64}}, Int64},Union{Slice{OneTo{Int64}}, Int64}}
│   %17 = (Base._selectdim)(A, d, i, %16)::SubArray
└──       return %17

The view works well:

julia> @code_warntype view(a, :, 1)
Body::SubArray{Float64,1,Array{Float64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true}
1 ─ %1  = (Base.arraysize)(A, 1)::Int64
│         (Base.arraysize)(A, 2)
│   %3  = (Base.slt_int)(%1, 0)::Bool
│   %4  = (Base.ifelse)(%3, 0, %1)::Int64
│   %5  = %new(Base.OneTo{Int64}, %4)::Base.OneTo{Int64}
│   %6  = %new(Base.Slice{Base.OneTo{Int64}}, %5)::Base.Slice{Base.OneTo{Int64}}
│   %7  = (getfield)(I, 2)::Int64
└──       goto #6 if not $(Expr(:boundscheck))
2 ─ %9  = (Core.tuple)(%6, %7)::Tuple{Base.Slice{Base.OneTo{Int64}},Int64}
│         (Base.arraysize)(A, 1)
│   %11 = (Base.arraysize)(A, 2)::Int64
│   %12 = (Base.slt_int)(%11, 0)::Bool
│   %13 = (Base.ifelse)(%12, 0, %11)::Int64
│   %14 = (Base.sle_int)(1, %7)::Bool
│   %15 = (Base.sle_int)(%7, %13)::Bool
│   %16 = (Base.and_int)(%14, %15)::Bool
│   %17 = (Base.and_int)(%16, true)::Bool
│   %18 = (Base.and_int)(true, %17)::Bool
└──       goto #4 if not %18
3 ─       goto #5
4 ─       invoke Base.throw_boundserror(_2::Array{Float64,2}, %9::Tuple{Base.Slice{Base.OneTo{Int64}},Int64})
└──       $(Expr(:unreachable))
5 ┄       nothing
6 ┄ %24 = (Core.tuple)(%6, %7)::Tuple{Base.Slice{Base.OneTo{Int64}},Int64}
│         (Base.arraysize)(A, 1)
│         (Base.arraysize)(A, 2)
│   %27 = (Base.arraysize)(A, 1)::Int64
│         (Base.arraysize)(A, 2)
│   %29 = (Base.slt_int)(%27, 0)::Bool
│   %30 = (Base.ifelse)(%29, 0, %27)::Int64
│   %31 = (Base.sub_int)(%30, 0)::Int64
│   %32 = (Base.mul_int)(1, %31)::Int64
│   %33 = (Base.sub_int)(%7, 1)::Int64
│   %34 = (Base.mul_int)(%33, %32)::Int64
│   %35 = (Base.add_int)(1, %34)::Int64
│         (Base.arraysize)(A, 1)
│         (Base.arraysize)(A, 2)
│   %38 = (Base.sub_int)(%35, 1)::Int64
│   %39 = %new(SubArray{Float64,1,Array{Float64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true}, A, %24, %38, 1)::SubArray{Float64,1,Array{Float64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true}
└──       return %39

I check that it arise from later part of this line. I am not sure how it causes this. Should I worry about this type instability, say if I am solving an ODE with a core derivative function having those selectdim?

Thanks!

1 Like

I am trying to dig a bit further into this. I am basically taking the code from the line you pointed out to.

Here is a function that mimicks selectdim(X, 1, k), taking an array X and returning a view of X when dimension 1 is fixed at an arbitrary index k:

function dim_view_concrete(X, k)
    dim_tup = map(Base.Slice, axes(X))
    idx = Base.setindex(dim_tup, k, 1)
    vv = view(X, idx...)
end
julia> X = rand(2,2);

julia> @code_warntype dim_view_concrete(X,1)
Body::SubArray{Float64,1,Array{Float64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true}
...

Here the SubArray’s type is correctly inferred. However, if we make the dimension arbitrary (instead of fixing at 1), the type isn’t correctly inferred anymore:

function dim_view_abstract(X,d, k)
    dim_tup = map(Base.Slice, axes(X))
    idx = Base.setindex(dim_tup, k, d)
    vv = view(X, idx...)
end
julia> @code_warntype dim_view_abstract(X,1,1)
Body::SubArray
...

Any help is appreciated!

Sorry I missed this at first! selectdim creates a view — and views specialize on their indices. This is hugely beneficial for performance, but it can result in a type-instability if the computation of the view’s indices aren’t type-stable. This is exactly the case with selectdim — for example, selectdim(A, 1, 1) creates a @view A[1, :, :] whereas selectdim(A, 2, 2) creates a @view A[:, 2, :]. While the returned dimensionality is constant (it’ll be a matrix in both cases), the indices are either an Int, Colon, Colon or a Colon, Int, Colon — and those have different types. Which one you get depends upon the value of the dimension argument and not its type — this is the definition of a type instability.

Now, we’ve very carefully structured things so this plays nicely with Julia’s constant propagation! When you ask for @code_warntype selectdim(A, 1, 1), you’re asking for the code for arbitrary Ints:

julia> @code_warntype selectdim(A, 1, 1)
Variables
  #self#::Core.Compiler.Const(selectdim, false)
  A::Array{Float64,2}
  d::Int64
  i::Int64

Body::SubArray
…

But, if we hard-code the dimension inside another function, Julia can actually figure it out:

julia> selectdim1(A, i) = selectdim(A, 1, i)

julia> @code_warntype selectdim1(A, 1)
Variables
  #self#::Core.Compiler.Const(selectdim1, false)
  A::Array{Float64,2}
  i::Int64

Body::SubArray{Float64,1,Array{Float64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true}
…

So, in some senses, @code_warntype is showing you an overly pessimistic answer and assuming that the dimension isn’t constant. In real code, though, that dimension argument is often a hard-coded constant and will give you fast, type-stable code.

7 Likes

Thanks @mbauman for the very clear response.

What follows is probably related to constant propagation as you mentioned so a bit off-topic; anyway for my use case I need the dimension to be arbitrary. I tried something like

d1 = 1
selectdim2 = let
    d = d1
    (A, i) -> selectdim(A, d, i)
end

and the result is still type instable code. Is there a simple solution for this?

The cost that a type instability incurs is a dynamic dispatch upon each subsequent function call within that function body. If there’s only one or two dynamic dispatches, then it’s not a big deal, but if you go on to use the unstable variable within a hot for loop or somesuch, then it’s a problem.

The easiest way to limit this cost is by using a function barrier. This is a simple approach that just moves everything below the instability into an inner function. Julia will pay the dynamic dispatch cost once — to get into the inner function — but once you’re inside that inner function everything will be specialized and stable.

Alternatively, if you’re taking your slices inside that hot loop, you can use a Val to hoist the dimension into the type domain, which then demands specialization for that particular value:

outer(d) = inner(Val(d))
function inner(::Val{dim}) where dim
    for i in blah
         slice = selectdim(A, dim, i)
    …

That makes perfect sense. Thank you so much.

Can we have a selectdim(A, Val(dim), i) method? Similar to how ntuple(f, Val(n)) helps reduce type instability.

That’ll only avoid a type instability in cases where it was already type stable.

Sorry, but this comment is not clear to me. Can you give an example? In particular, why is this different from ntuple(f, ::Val) which does benefit from Val?

This is actually the same as ntuple (for n ≤ 10) — and I’m not sure we’d have its Val method were it newly introduced today. The key is that both of them rely upon constant propagation to be type-stable. That is, if you write ntuple(f, 4) in a compiled context, you get this bit of awesomeness:

julia> g(fn) = ntuple(fn, 4)
g (generic function with 1 method)

julia> @code_typed g(x->x+1)
CodeInfo(
1 ─     return (2, 3, 4, 5)
) => NTuple{4,Int64}

(Look ma, no Val!)

The only advantage of using Val is that it’ll jump through non-inlined functions.

3 Likes