Type stability, passing type as argument, where keyword

I’m working on a general interface for a package. I were surprised when I checked code stability. Now I’m feeling that I’m missing something very important about parametric methods.

My question is

Why foobar_clsr_1 is type stable whereas foobar_clsr_2 is not?

abstract type AbstractFoo end

struct Foo{T} <: AbstractFoo
    x::T
end

function foofunc(::Type{Foo}, x::AbstractVector)
    return 2x
end

function barfunc(f::Foo)
    return 2f.x
end

function foobar_clsr_1(::Type{T}, b::AbstractVector) where {T}
    fooclsr(x::AbstractVector) = b + foofunc(T, x)
    barclsr(x::AbstractVector) = b + barfunc(T(x))

    return fooclsr, barclsr
end

function foobar_clsr_2(T::Type{<:AbstractFoo}, b::AbstractVector)
    fooclsr(x::AbstractVector) = b + foofunc(T, x)
    barclsr(x::AbstractVector) = b + barfunc(T(x))

    return fooclsr, barclsr
end

foo1, bar1 = foobar_clsr_1(Foo, [10, 10])
foo2, bar2 = foobar_clsr_2(Foo, [10, 10])

@code_warntype gives this (I’m using julia 1.6.2 on macOS)

% julia -i test.jl
......
julia> x = [2, 2];

julia> @code_warntype foo1(x)
Variables
  #self#::var"#fooclsr#1"{Foo, Vector{Int64}}
  x::Vector{Int64}

Body::Vector{Int64}
1 ─ %1 = Core.getfield(#self#, :b)::Vector{Int64}
│   %2 = Main.foofunc($(Expr(:static_parameter, 1)), x)::Vector{Int64}
│   %3 = (%1 + %2)::Vector{Int64}
└──      return %3

julia> @code_warntype foo2(x)
Variables
  #self#::var"#fooclsr#3"{UnionAll, Vector{Int64}}
  x::Vector{Int64}

Body::Vector{Int64}
1 ─ %1 = Core.getfield(#self#, :b)::Vector{Int64}
│   %2 = Core.getfield(#self#, :T)::UnionAll  # (!) UnionAll
│   %3 = Main.foofunc(%2, x)::Vector{Int64}
│   %4 = (%1 + %3)::Vector{Int64}
└──      return %4

julia> @code_warntype bar1(x)
Variables
  #self#::var"#barclsr#2"{Foo, Vector{Int64}}
  x::Vector{Int64}

Body::Vector{Int64}
1 ─ %1 = Core.getfield(#self#, :b)::Vector{Int64}
│   %2 = ($(Expr(:static_parameter, 1)))(x)::Foo{Vector{Int64}}
│   %3 = Main.barfunc(%2)::Vector{Int64}
│   %4 = (%1 + %3)::Vector{Int64}
└──      return %4

julia> @code_warntype bar2(x)
Variables
  #self#::var"#barclsr#4"{UnionAll, Vector{Int64}}
  x::Vector{Int64}

Body::Any
1 ─ %1 = Core.getfield(#self#, :b)::Vector{Int64}
│   %2 = Core.getfield(#self#, :T)::UnionAll  # (!) UnionAll
│   %3 = (%2)(x)::Any
│   %4 = Main.barfunc(%3)::Any
│   %5 = (%1 + %4)::Any
└──      return %5

The reason is that julia avoids specializing on ::Type except in the case where you do Type{T} where T like in foobar_clsr_1

3 Likes

Thank you for the link. What does “specialize” mean? Is it “hey, compiler, infer types when possible”? So, Type{T} where T forces Julia to “specializing”?

2 Likes

yes, exactly that

If the compiler specialize then it compiles a specific version of the method body for each specific combination of input types, i.e., there are multiple compiled method bodies and, at the function call, dynamic dispatch does not select just the method body but the specific specialization of it.

If the compiler does not specialize, then it compiles a generic version of the method body, and each combination of input types that leads to that method body will call the same compiled code that should work for any combination of types that lead the dynamic dispatch to that method body.

1 Like

Another way of saying what the others have said above:

Julia compiles a method based on the types of the parameters passed in to the function, not based on the values.

In foobar_clsr_2, T is a value which happens to have a type of Type. At compile time, the compiler only knows that this is Type. And based on that knowledge, it can’t avoid type instability, which means it’s going to compile a really generic, poorly optimized version of this function that can handle any type T.

On the other hand, foobar_clsr_1 has a type T elevated to the type of the first variable thanks to the parameterization, so the compiler knows about T at compile time and can specialize on it. For every new type T that you pass this function, Julia will compile a new optimized version of the function specific to that type.

1 Like