Type instability with StaticArrays.jl and “type alias”

[ranocha]

I see the following type instability in current Julia v1.4 and master (commit cebd4fa).

julia> using StaticArrays
[ Info: Precompiling StaticArrays [90137ffa-7385-5640-81b9-e52037218182]

julia> struct Foo{N} end

julia> foo = Foo{4}()
Foo{4}()

julia> bar(::Foo{N}) where N = N
bar (generic function with 1 method)

julia> function baz1(foo)
           MVector{bar(foo), Float64}(undef)
       end
baz1 (generic function with 1 method)

julia> @code_warntype baz1(foo)
Variables
  #self#::Core.Compiler.Const(baz1, false)
  foo::Core.Compiler.Const(Foo{4}(), false)

Body::MArray{Tuple{4},Float64,1,4}
1 ─ %1 = Main.bar(foo)::Core.Compiler.Const(4, false)
│   %2 = Core.apply_type(Main.MVector, %1, Main.Float64)::Core.Compiler.Const(MArray{Tuple{4},Float64,1,4}, false)
│   %3 = (%2)(Main.undef)::MArray{Tuple{4},Float64,1,4}
└──      return %3

julia> function baz2(foo)
           tmp = MVector{bar(foo), Float64}
           tmp(undef)
       end
baz2 (generic function with 1 method)

julia> @code_warntype baz2(foo)
Variables
  #self#::Core.Compiler.Const(baz2, false)
  foo::Core.Compiler.Const(Foo{4}(), false)
  tmp::Type{MArray{Tuple{4},Float64,1,4}}

Body::MArray{Tuple{4},Float64,1,4}
1 ─ %1 = Main.bar(foo)::Core.Compiler.Const(4, false)
│        (tmp = Core.apply_type(Main.MVector, %1, Main.Float64))
│   %3 = (tmp::Core.Compiler.Const(MArray{Tuple{4},Float64,1,4}, false))(Main.undef)::MArray{Tuple{4},Float64,1,4}
└──      return %3

julia> function baz3(foo)
           tmp = MVector{bar(foo), Float64}
           [tmp(undef) for _ in 1:bar(foo)]
       end
baz3 (generic function with 1 method)

julia> @code_warntype baz3(foo)
Variables
  #self#::Core.Compiler.Const(baz3, false)
  foo::Core.Compiler.Const(Foo{4}(), false)
  #1::var"#1#2"{DataType}
  tmp::Type{MArray{Tuple{4},Float64,1,4}}

Body::Array{_A,1} where _A
1 ─ %1  = Main.bar(foo)::Core.Compiler.Const(4, false)
│         (tmp = Core.apply_type(Main.MVector, %1, Main.Float64))
│   %3  = Main.:(var"#1#2")::Core.Compiler.Const(var"#1#2", false)
│   %4  = Core.typeof(tmp::Core.Compiler.Const(MArray{Tuple{4},Float64,1,4}, false))::Core.Compiler.Const(DataType, false)
│   %5  = Core.apply_type(%3, %4)::Core.Compiler.Const(var"#1#2"{DataType}, false)
│         (#1 = %new(%5, tmp::Core.Compiler.Const(MArray{Tuple{4},Float64,1,4}, false)))
│   %7  = #1::Core.Compiler.Const(var"#1#2"{DataType}(MArray{Tuple{4},Float64,1,4}), false)::Core.Compiler.Const(var"#1#2"{DataType}(MArray{Tuple{4},Float64,1,4}), false)
│   %8  = Main.bar(foo)::Core.Compiler.Const(4, false)
│   %9  = (1:%8)::Core.Compiler.Const(1:4, false)
│   %10 = Base.Generator(%7, %9)::Core.Compiler.Const(Base.Generator{UnitRange{Int64},var"#1#2"{DataType}}(var"#1#2"{DataType}(MArray{Tuple{4},Float64,1,4}), 1:4), false)
│   %11 = Base.collect(%10)::Array{_A,1} where _A
└──       return %11

baz1 and baz2 are just fine but baz3 results in a type instability that I didn’t expect. I can fix that by not using the “type alias” tmp = MVector{bar(foo), Float64} or adding tmp in front of the arrays

julia> function baz4(foo)
           [MVector{bar(foo), Float64}(undef) for _ in 1:bar(foo)]
       end
baz4 (generic function with 1 method)

julia> @code_warntype baz4(foo)
Variables
  #self#::Core.Compiler.Const(baz4, false)
  foo::Core.Compiler.Const(Foo{4}(), false)
  #3::var"#3#4"{Foo{4}}

Body::Array{MArray{Tuple{4},Float64,1,4},1}
1 ─ %1 = Main.:(var"#3#4")::Core.Compiler.Const(var"#3#4", false)
│   %2 = Core.typeof(foo)::Core.Compiler.Const(Foo{4}, false)
│   %3 = Core.apply_type(%1, %2)::Core.Compiler.Const(var"#3#4"{Foo{4}}, false)
│        (#3 = %new(%3, foo))
│   %5 = #3::Core.Compiler.Const(var"#3#4"{Foo{4}}(Foo{4}()), false)
│   %6 = Main.bar(foo)::Core.Compiler.Const(4, false)
│   %7 = (1:%6)::Core.Compiler.Const(1:4, false)
│   %8 = Base.Generator(%5, %7)::Core.Compiler.Const(Base.Generator{UnitRange{Int64},var"#3#4"{Foo{4}}}(var"#3#4"{Foo{4}}(Foo{4}()), 1:4), false)
│   %9 = Base.collect(%8)::Array{MArray{Tuple{4},Float64,1,4},1}
└──      return %9

julia> function baz5(foo)
           tmp = MVector{bar(foo), Float64}
           tmp[tmp(undef) for _ in 1:bar(foo)]
       end
baz5 (generic function with 1 method)

julia> @code_warntype baz5(foo)
Variables
  #self#::Core.Compiler.Const(baz5, false)
  foo::Core.Compiler.Const(Foo{4}(), false)
  tmp::Type{MArray{Tuple{4},Float64,1,4}}
  @_4::Array{MArray{Tuple{4},Float64,1,4},1}
  @_5::Int64
  @_6::Union{Nothing, Tuple{Int64,Int64}}

Body::Array{MArray{Tuple{4},Float64,1,4},1}
1 ─ %1  = Main.bar(foo)::Core.Compiler.Const(4, false)
│         (tmp = Core.apply_type(Main.MVector, %1, Main.Float64))
│         Core.NewvarNode(:(@_4))
│         Core.NewvarNode(:(@_5))
│         Core.NewvarNode(:(@_6))
│   %6  = Main.bar(foo)::Core.Compiler.Const(4, false)
│   %7  = (1:%6)::Core.Compiler.Const(1:4, false)
│   %8  = Base.IteratorSize(%7)::Core.Compiler.Const(Base.HasShape{1}(), false)
│   %9  = (%8 isa Base.SizeUnknown)::Core.Compiler.Const(false, false)
└──       goto #3 if not %9
2 ─       Core.Compiler.Const(:(Core.apply_type(Core.Array, tmp, 1)), false)
│         Core.Compiler.Const(:(@_4 = (%11)(Core.undef, 0)), false)
└──       Core.Compiler.Const(:(goto %15), false)
3 ┄       (@_4 = Base._array_for(tmp::Core.Compiler.Const(MArray{Tuple{4},Float64,1,4}, false), %7, %8))
│   %15 = Base.LinearIndices(@_4)::LinearIndices{1,Tuple{Base.OneTo{Int64}}}
│         (@_5 = Base.first(%15))
│   %17 = %7::Core.Compiler.Const(1:4, false)
│         (@_6 = Base.iterate(%17))
│   %19 = (@_6::Core.Compiler.Const((1, 1), false) === nothing)::Core.Compiler.Const(false, false)
│   %20 = Base.not_int(%19)::Core.Compiler.Const(true, false)
└──       goto #8 if not %20
4 ┄ %22 = @_6::Tuple{Int64,Int64}::Tuple{Int64,Int64}
│         Core.getfield(%22, 1)
│   %24 = Core.getfield(%22, 2)::Int64
│   %25 = (tmp::Core.Compiler.Const(MArray{Tuple{4},Float64,1,4}, false))(Main.undef)::MArray{Tuple{4},Float64,1,4}
│         $(Expr(:inbounds, true))
└──       goto #6 if not %9
5 ─       Core.Compiler.Const(:(Base.push!(@_4, %25)), false)
└──       Core.Compiler.Const(:(goto %31), false)
6 ┄       Base.setindex!(@_4, %25, @_5)
│         $(Expr(:inbounds, :pop))
│         (@_5 = Base.add_int(@_5, 1))
│         (@_6 = Base.iterate(%17, %24))
│   %34 = (@_6 === nothing)::Bool
│   %35 = Base.not_int(%34)::Bool
└──       goto #8 if not %35
7 ─       goto #4
8 ┄       return @_4

Quoting @rdeits in https://discourse.julialang.org/t/type-instability-with-staticarrays-jl-and-type-alias/39763

Yeah, that’s interesting. It looks like tmp is being handled as ::DataType rather than ::Type{MArray{...}}. That seems like something that would need to be handled in Julia itself, rather than StaticArrays.

[JeffBezanson]

It should be fairly easy to keep more type information inside the closure, but I'll have to see if it affects compiler performance at all.

[JeffBezanson]

Duplicate of #23618