Why selectdim is type instable

General Usage
1

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
2

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!

3

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
4

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?

5

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)
    …
6

That makes perfect sense. Thank you so much.

7

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

8

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

9

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?

10

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