Making vcat of splatted comprehension type-stable

HenriDeh · July 29, 2020, 11:55am

Hi,

I am looking for a way to make this kind of operation type-stable. There are many uses to structs with lists of callables but I can’t seem to find a way to make the concatenation of their outputs type-stable

julia> struct Foo{T}
       fs::T
       end

julia> (foo::Foo)(tup) = vcat([foo.fs[1](tup[1]) for i in 1:lenght(tup)]...)

julia> foo = Foo((x->diff(x), x-> round.(x)))
Foo{Tuple{var"#15#17",var"#16#18"}}((var"#15#17"(), var"#16#18"()))

julia> @code_warntype foo((rand(5), rand(5)))
Variables
  foo::Core.Compiler.Const(Foo{Tuple{var"#15#17",var"#16#18"}}((var"#15#17"(), var"#16#18"())), false)
  tup::Tuple{Array{Float64,1},Array{Float64,1}}
  #9::var"#9#10"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}

Body::Any
1 ─ %1  = Main.:(var"#9#10")::Core.Compiler.Const(var"#9#10", false)
│   %2  = Core.typeof(foo)::Core.Compiler.Const(Foo{Tuple{var"#15#17",var"#16#18"}}, false)
│   %3  = Core.typeof(tup)::Core.Compiler.Const(Tuple{Array{Float64,1},Array{Float64,1}}, false)
│   %4  = Core.apply_type(%1, %2, %3)::Core.Compiler.Const(var"#9#10"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}, false)
│         (#9 = %new(%4, foo, tup))
│   %6  = #9::var"#9#10"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}
│   %7  = Main.lenght(tup)::Any
│   %8  = (1:%7)::Any
│   %9  = Base.Generator(%6, %8)::Base.Generator{_A,var"#9#10"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}} where _A
│   %10 = Base.collect(%9)::Array{T,N} where T where N
│   %11 = Core._apply_iterate(Base.iterate, Main.vcat, %10)::Any
└──       return %11

I also tried with reduce:

julia> (foo::Foo)(tup) = reduce(vcat, [foo.fs[1](tup[1]) for i in 1:lenght(tup)])

julia> @code_warntype foo((rand(5), rand(5)))
Variables
  foo::Core.Compiler.Const(Foo{Tuple{var"#15#17",var"#16#18"}}((var"#15#17"(), var"#16#18"())), false)
  tup::Tuple{Array{Float64,1},Array{Float64,1}}
  #19::var"#19#20"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}

Body::Any
1 ─ %1  = Main.:(var"#19#20")::Core.Compiler.Const(var"#19#20", false)
│   %2  = Core.typeof(foo)::Core.Compiler.Const(Foo{Tuple{var"#15#17",var"#16#18"}}, false)
│   %3  = Core.typeof(tup)::Core.Compiler.Const(Tuple{Array{Float64,1},Array{Float64,1}}, false)
│   %4  = Core.apply_type(%1, %2, %3)::Core.Compiler.Const(var"#19#20"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}, false)
│         (#19 = %new(%4, foo, tup))
│   %6  = #19::var"#19#20"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}
│   %7  = Main.lenght(tup)::Any
│   %8  = (1:%7)::Any
│   %9  = Base.Generator(%6, %8)::Base.Generator{_A,var"#19#20"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}} where _A
│   %10 = Base.collect(%9)::Array{T,N} where T where N
│   %11 = Main.reduce(Main.vcat, %10)::Any
└──       return %11

For some reason, the compiler doesn’t seem to catch that lenght(tup) is an Int. If I replace by a constant (which is not desirable) I get a type-stable operation but the vcat still is type-stable:

julia> (foo::Foo)(tup) = vcat([foo.fs[1](tup[1]) for i in 1:2]...)

julia> @code_warntype foo((rand(5), rand(5)))
Variables
  foo::Core.Compiler.Const(Foo{Tuple{var"#15#17",var"#16#18"}}((var"#15#17"(), var"#16#18"())), false)
  tup::Tuple{Array{Float64,1},Array{Float64,1}}
  #25::var"#25#26"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}

Body::Union{Array{Any,1}, Array{Float64,1}}
1 ─ %1  = Main.:(var"#25#26")::Core.Compiler.Const(var"#25#26", false)
│   %2  = Core.typeof(foo)::Core.Compiler.Const(Foo{Tuple{var"#15#17",var"#16#18"}}, false)
│   %3  = Core.typeof(tup)::Core.Compiler.Const(Tuple{Array{Float64,1},Array{Float64,1}}, false)
│   %4  = Core.apply_type(%1, %2, %3)::Core.Compiler.Const(var"#25#26"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}, false)
│         (#25 = %new(%4, foo, tup))
│   %6  = #25::var"#25#26"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}
│   %7  = (1:2)::Core.Compiler.Const(1:2, false)
│   %8  = Base.Generator(%6, %7)::Core.Compiler.PartialStruct(Base.Generator{UnitRange{Int64},var"#25#26"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}}, Any[var"#25#26"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}, Core.Compiler.Const(1:2, false)])
│   %9  = Base.collect(%8)::Array{Array{Float64,1},1}
│   %10 = Core._apply_iterate(Base.iterate, Main.vcat, %9)::Union{Array{Any,1}, Array{Float64,1}}
└──       return %10

With reduce, I get a type-stable function but reduce applies vcat several times which seems unefficient:

julia> (foo::Foo)(tup) = reduce(vcat, [foo.fs[1](tup[1]) for i in 1:2])

julia> @code_warntype foo((rand(5), rand(5)))
Variables
  foo::Core.Compiler.Const(Foo{Tuple{var"#15#17",var"#16#18"}}((var"#15#17"(), var"#16#18"())), false)
  tup::Tuple{Array{Float64,1},Array{Float64,1}}
  #23::var"#23#24"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}

Body::Array{Float64,1}
1 ─ %1  = Main.:(var"#23#24")::Core.Compiler.Const(var"#23#24", false)
│   %2  = Core.typeof(foo)::Core.Compiler.Const(Foo{Tuple{var"#15#17",var"#16#18"}}, false)
│   %3  = Core.typeof(tup)::Core.Compiler.Const(Tuple{Array{Float64,1},Array{Float64,1}}, false)
│   %4  = Core.apply_type(%1, %2, %3)::Core.Compiler.Const(var"#23#24"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}, false)
│         (#23 = %new(%4, foo, tup))
│   %6  = #23::var"#23#24"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}
│   %7  = (1:2)::Core.Compiler.Const(1:2, false)
│   %8  = Base.Generator(%6, %7)::Core.Compiler.PartialStruct(Base.Generator{UnitRange{Int64},var"#23#24"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}}, Any[var"#23#24"{Foo{Tuple{var"#15#17",var"#16#18"}},Tuple{Array{Float64,1},Array{Float64,1}}}, Core.Compiler.Const(1:2, false)])
│   %9  = Base.collect(%8)::Array{Array{Float64,1},1}
│   %10 = Main.reduce(Main.vcat, %9)::Array{Float64,1}
└──       return %10

So, I see two problems here. 1), vcat(::Array{Array}...) is not type-stable (::Union{Array{Any,1}, Array{Float64,1}}). 2) length(tup) is ::Any.

How would you do this efficiently ?

rdeits · July 30, 2020, 3:15pm

That’s because you’ve spelled length as lenght

It’s always a good idea to run your function before using @code_warntype because it can be tricky to tell the difference between real type instabilities and issues caused by plain old errors in your code.

HenriDeh · July 30, 2020, 3:19pm

Ha silly me, thanks for that.
Though there is still the problem that vcat of a splatted array of arrays is a union-type.

julia> @code_warntype foo((rand(5), rand(5)))
Variables
  foo::Core.Compiler.Const(Foo{Tuple{var"#13#15",var"#14#16"}}((var"#13#15"(), var"#14#16"())), false)
  tup::Tuple{Array{Float64,1},Array{Float64,1}}
  #11::var"#11#12"{Foo{Tuple{var"#13#15",var"#14#16"}},Tuple{Array{Float64,1},Array{Float64,1}}}

Body::Union{Array{Any,1}, Array{Float64,1}}
1 ─ %1  = Main.:(var"#11#12")::Core.Compiler.Const(var"#11#12", false)
│   %2  = Core.typeof(foo)::Core.Compiler.Const(Foo{Tuple{var"#13#15",var"#14#16"}}, false)
│   %3  = Core.typeof(tup)::Core.Compiler.Const(Tuple{Array{Float64,1},Array{Float64,1}}, false)
│   %4  = Core.apply_type(%1, %2, %3)::Core.Compiler.Const(var"#11#12"{Foo{Tuple{var"#13#15",var"#14#16"}},Tuple{Array{Float64,1},Array{Float64,1}}}, false)
│         (#11 = %new(%4, foo, tup))
│   %6  = #11::var"#11#12"{Foo{Tuple{var"#13#15",var"#14#16"}},Tuple{Array{Float64,1},Array{Float64,1}}}
│   %7  = Main.length(tup)::Core.Compiler.Const(2, false)
│   %8  = (1:%7)::Core.Compiler.Const(1:2, false)
│   %9  = Base.Generator(%6, %8)::Core.Compiler.PartialStruct(Base.Generator{UnitRange{Int64},var"#11#12"{Foo{Tuple{var"#13#15",var"#14#16"}},Tuple{Array{Float64,1},Array{Float64,1}}}}, Any[var"#11#12"{Foo{Tuple{var"#13#15",var"#14#16"}},Tuple{Array{Float64,1},Array{Float64,1}}}, Core.Compiler.Const(1:2, false)])
│   %10 = Base.collect(%9)::Array{Array{Float64,1},1}
│   %11 = Core._apply(Main.vcat, %10)::Union{Array{Any,1}, Array{Float64,1}}
└──       return %11

rdeits · July 30, 2020, 3:49pm

Is using a splatting operation here actually desirable? reduce(vcat, x) is generally recommended over vcat(x...) anyway.

In general, splatting an object whose length is not trivially known by the compiler (so, basically anything except a Tuple or something wrapping a Tuple) will cause type-instabilities because it means calling a function with a number of arguments unknown to the compiler.

HenriDeh · July 30, 2020, 3:58pm

Haaa yes okay I understand. Consequently, if I use a tuple comprehension, the function becomes type-stable

julia> (foo::Foo)(tup) = vcat((foo.fs[1](tup[1]) for i in 1:length(tup))...)

julia> @code_warntype foo((rand(5), rand(5)))
Variables
  foo::Core.Compiler.Const(Foo{Tuple{var"#13#15",var"#14#16"}}((var"#13#15"(), var"#14#16"())), false)
  tup::Tuple{Array{Float64,1},Array{Float64,1}}
  #19::var"#19#20"{Foo{Tuple{var"#13#15",var"#14#16"}},Tuple{Array{Float64,1},Array{Float64,1}}}

Body::Array{Float64,1}
1 ─ %1  = Main.:(var"#19#20")::Core.Compiler.Const(var"#19#20", false)
│   %2  = Core.typeof(foo)::Core.Compiler.Const(Foo{Tuple{var"#13#15",var"#14#16"}}, false)
│   %3  = Core.typeof(tup)::Core.Compiler.Const(Tuple{Array{Float64,1},Array{Float64,1}}, false)
│   %4  = Core.apply_type(%1, %2, %3)::Core.Compiler.Const(var"#19#20"{Foo{Tuple{var"#13#15",var"#14#16"}},Tuple{Array{Float64,1},Array{Float64,1}}}, false)
│         (#19 = %new(%4, foo, tup))
│   %6  = #19::var"#19#20"{Foo{Tuple{var"#13#15",var"#14#16"}},Tuple{Array{Float64,1},Array{Float64,1}}}
│   %7  = Main.length(tup)::Core.Compiler.Const(2, false)
│   %8  = (1:%7)::Core.Compiler.Const(1:2, false)
│   %9  = Base.Generator(%6, %8)::Core.Compiler.PartialStruct(Base.Generator{UnitRange{Int64},var"#19#20"{Foo{Tuple{var"#13#15",var"#14#16"}},Tuple{Array{Float64,1},Array{Float64,1}}}}, Any[var"#19#20"{Foo{Tuple{var"#13#15",var"#14#16"}},Tuple{Array{Float64,1},Array{Float64,1}}}, Core.Compiler.Const(1:2, false)])
│   %10 = Core._apply(Main.vcat, %9)::Array{Float64,1}
└──       return %10

jishnub · August 16, 2020, 12:56pm

reduce(vcat, x) returns a different result from vcat(x...) if x is a single-element collection. This might be an edge case worth considering in case the length of x is not known a-priori to be >1.

julia> reduce(vcat, [1])
1

julia> vcat([1]...)
1-element Array{Int64,1}:
 1

One way to avoid this is by specifying an empty initial value of the correct type.

julia> reduce(vcat, [1], init=Int[])
1-element Array{Int64,1}:
 1