Type instability with StaticArrays.jl and "type alias"

ranocha · May 19, 2020, 4:29pm

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

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} but

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

Since I want to create several arrays like [MVector{bar(foo), Float64}(undef) for _ in 1:bar(foo)], I didn’t want to duplicate all the code. That’s why I used the assignment to tmp.

Is this an expected type instability?
Is there another way to circumvent it?
Would it be helpful to report this somewhere else (Julia or StaticArrays.jl issue tracker) to improve inference in the future?

Thanks!

rdeits · May 19, 2020, 4:45pm

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.

On the plus side, you can do a typed list comprehension to restore type stability in this case:

julia> function baz3(foo)
           tmp = MVector{bar(foo), Float64}
           tmp[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)
  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}

The tmp[...] syntax is just as if you were doing something like:

julia> Float64[i for i in 1:10]
10-element Array{Float64,1}:
...

ranocha · May 19, 2020, 4:58pm

Thanks, @rdeits! At least that solves a bit of code duplication for me. I’ll leave this open and wait a bit longer in case there might be more answers.

pablosanjose · May 19, 2020, 5:05pm

Well, after waiting a bit for feedback it might be good to open an issue in GitHub with this. It looks like an interesting example for the Julia compiler devs

ranocha · May 19, 2020, 5:07pm

That’s what I plan to do. It might just happen that someone on discourse knows “Oh, wait, that’s issue XYZ in Julia.” and I want to give them a bit more time before opening an issue there

ranocha · May 20, 2020, 4:01am

I’ve opened an issue: Type instability with StaticArrays.jl and “type alias” · Issue #35949 · JuliaLang/julia · GitHub

ranocha · May 20, 2020, 4:03am

Hmm, it seems that I can’t select an “answer”. I have only the options “like, share, flag, bookmark, and reply”.