Const-prop through tuple

[maleadt]

The following fails to generate efficient code:

inner(::Type{Bool}, i) = nothing
outer(I...) = inner(Bool, I...)

using Cassette
Cassette.@context Ctx

f = (args...) -> Cassette.overdub(Ctx(), outer, args...)
code_warntype(f, Tuple{Int})

Body::ANY
10 1 ─ %1 = Cassette.overdub::Core.Compiler.Const(Cassette.overdub, false)
   │   %2 = (getfield)(args, 1)::Int64
   │   %3 = invoke %1($(QuoteNode(Cassette.Context{nametype(Ctx),Nothing,Cassette.NoPass,Nothing,Nothing}(nametype(Ctx)(), nothing, Cassette.NoPass(), nothing, nothing)))::Cassette.Context{nametype(Ctx),Nothing,Cassette.NoPass,Nothing,Nothing}, outer::typeof(outer), %2::Int64)::ANY
   └──      return %3

... whereas without overdubbing, we simply get:

Body::Nothing
2 1 ─     return

Checking the code for overdub reveals why we get the invoke:

code_warntype(Cassette.overdub, Tuple{typeof(Ctx()), typeof(outer), Int})

Body::ANY
2 1 ─ %1  = (Core.getfield)(##overdub_arguments#369, 2)::Int64
  │   %2  = (Core.tuple)(%1)::Tuple{Int64}
  └──       goto #3 if not true
  2 ─ %4  = Core.tuple::Core.Compiler.Const(tuple, false)
  │   %5  = Main.Bool::Core.Compiler.Const(Bool, false)
  └── %6  = (%4)(%5)::Tuple{DataType}
  3 ─ %7  = φ (#2 => %6, #1 => $(QuoteNode(Cassette.OverdubInstead())))::UNION{OVERDUBINSTEAD, TUPLE{DATATYPE}}
  │   %8  = π (%7, Core.Compiler.Const((Bool,), false))
  └──       goto #5 if not true
  4 ─ %10 = Core._apply::typeof(Core._apply)
  │   %11 = (%10)(tuple, %8, %2)::Tuple{DataType,Int64}
  │   %12 = (getfield)(%11, 1)::DATATYPE
  │   %13 = (getfield)(%11, 2)::Int64
  └── %14 = (Cassette.overdub)($(QuoteNode(Cassette.Context{nametype(Ctx),Nothing,Cassette.NoPass,Nothing,Nothing}(nametype(Ctx)(), nothing, Cassette.NoPass(), nothing, nothing))), inner, %12, %13)::ANY
  5 ─ %15 = φ (#4 => %14, #3 => $(QuoteNode(Cassette.OverdubInstead())))::ANY
  └──       return %15

Looks like a failure to const-prop the Bool argument through the argument tuple. But then again, the second element of the tuple isn't const, and I'm not sure whether that can be expressed with Compiler.Const.

This specific pattern arises with checkbounds(::Array, ::Integer) which calls to checkbounds(Bool, ...) (ie. passing a Type{Bool}), so this probably penalizes quite a lot of code as soon as it indexes arrays.

[jrevels]

I've had to work around a couple of things like this by just making them primitives to force a specialization but this is starting to get bothersome.

It looks like in this case _apply should see the const (Bool,), so maybe abstract interpretation can propagate this better? Especially if the applied function is tuple, is it worth special casing abstract_apply such that e.g. Core._apply(tuple, (T,), 1)::Tuple{Type{T},Int} if (T,) is const? I'm too ignorant to know whether this a good idea or not - seems like inference might not like being fed more kind types like this...I'll play around with it. EDIT: I half-implemented this idea and it worked well for this use case, but after discussions with Jameson, it's now clear that an actual correct implementation will require a bit more work than just changing widenconst for const tuple values....we'll have to add a tuple type lattice element to inference.

[jrevels]

This appears to be fixed by JuliaLang/julia#28955:

julia> code_typed(f, Tuple{Int})
1-element Array{Any,1}:
 CodeInfo(
1 1 ─ %1 = (getfield)(args, 1)::Int64                       │
  │   %2 = (Core.tuple)(%1)::Tuple{Int64}                   │╻   outer
  │   %3 = Core.tuple::typeof(tuple)                        ││
  │   %4 = Main.Bool::Type{Bool}                            ││
  │   %5 = invoke fallback($(QuoteNode(Cassette.Context{nametype(Ctx),Nothing,Cassette.NoPass,Nothing,Nothing}(nametype(Ctx)(), nothing, Cassette.NoPass(), nothing, nothing)))::Cassette.Context{nametype(Ctx),Nothing,Cassette.NoPass,Nothing,Nothing}, %3::Function, %4::Vararg{Any,N} where N)::Tuple{DataType}
  │   %6 = Core._apply::typeof(_apply)                      │││╻   apply_args
  │        (%6)(tuple, %5, %2)::Tuple{Type{Bool},Int64}     ││││
  └──      return                                           │
) => Nothing

[Keno]

The extra _apply here is probably fixed by JuliaLang/julia#29622

[maleadt]

Initial repro seems fixed indeed, but zooming out the general pattern as used by checkbounds doesn't. MWE:

outer(A::AbstractArray, i...) = inner(Bool, A, i...)
inner(::Type{Bool}, A::AbstractArray, i...) = nothing


using Cassette
Cassette.@context Ctx

overdubbed = (args...) -> Cassette.overdub(Ctx(), outer, args...)


using InteractiveUtils

code_warntype(outer, Tuple{Vector{Int}, Int})
code_warntype(overdubbed, Tuple{Vector{Int}, Int})

Body::Nothing
1 1 ─     return
Body::ANY
8 1 ─ %1 = Cassette.overdub::Core.Compiler.Const(Cassette.overdub, false)
  │   %2 = (getfield)(args, 1)::Array{Int64,1}
  │   %3 = (getfield)(args, 2)::Int64
  │   %4 = invoke %1($(QuoteNode(Cassette.Context{nametype(Ctx),Nothing,Cassette.NoPass,Nothing,Nothing}(nametype(Ctx)(), nothing, Cassette.NoPass(), nothing, nothing)))::Cassette.Context{nametype(Ctx),Nothing,Cassette.NoPass,Nothing,Nothing}, outer::typeof(outer), %2::Array{Int64,1}, %3::Int64)::ANY
  └──      return %4

Latest master of both Julia and Cassette.

[jrevels]

Looks like there's still some _apply/varargs weirdness to shake out in inference:

julia> code_typed(overdubbed, Tuple{Vector{Int}, Int}, optimize=false)
1-element Array{Any,1}:
 CodeInfo(
1 1 ─ %1 = Cassette.overdub::Const(overdub, false)                          │
  │   %2 = (Main.Ctx)()::Const(Context{nametype(Ctx),Nothing,NoPass,Nothing,Nothing}(nametype(Ctx)(), nothing, NoPass(), nothing, nothing), false)
  │   %3 = (Core.tuple)(%2, Main.outer)::Const((Context{nametype(Ctx),Nothing,NoPass,Nothing,Nothing}(nametype(Ctx)(), nothing, NoPass(), nothing, nothing), outer), false)
  │   %4 = (Core._apply)(%1, %3, args)::Any                                 │
  └──      return %4                                                        │
) => Any

at least this should be easier to handle now that we have PartialTuple.

[jrevels]

Did some more digging here today through the InferenceResult cache for this invocation (see here, for repro of the cache construction set mymethod to overdubbed and mysig to Tuple{typeof(mymethod),Vector{Int},Int}).

If we look at the result for the overdubbed outer call:

julia> r = p.cache[11]; # the one with the `Any` result

julia> r.linfo
MethodInstance for overdub(::Cassette.Context{nametype(Ctx),Nothing,Cassette.NoPass,Nothing,Nothing}, ::typeof(outer), ::Array{Int64,1}, ::Int64)

julia> r.src
CodeInfo(
1 1 ─ %1 = (Core.getfield)(##overdub_arguments#359, 2)::Array{Int64,1}                                   │
  │   %2 = (Core.getfield)(##overdub_arguments#359, 3)::Int64                                            │
  └──      goto #3 if not true                                                                           │
  2 ─ %4 = Main.Bool::Type{Bool}                                                                         │
  3 ─ %5 = φ (#2 => %4)::DataType                                                                        │
  │   %6 = φ (#2 => %1)::Array{Int64,1}                                                                  │
  └──      goto #5 if not true                                                                           │
  4 ─ %8 = (Cassette.overdub)($(QuoteNode(Cassette.Context{nametype(Ctx),Nothing,Cassette.NoPass,Nothing,Nothing}(nametype(Ctx)(), nothing, Cassette.NoPass(), nothing, nothing))), inner, %5, %6, %2)::Any
  5 ─ %9 = φ (#4 => %8, #3 => $(QuoteNode(Cassette.OverdubInstead())))::Any                              │
  └──      return %9                                                                                     │
)

Looks like there's a phi node there that's normalizing Type{Bool} to DataType? @Keno could this be the culprit? I had optimize set to false but @vtjnash told me that "inference lies and runs optimization regardless" (I couldn't squeeze the reason out of him, though 😛).

Plugging away further (on the branch from JuliaLang/julia#29893):

julia> ctx = Ctx();

julia> tup(args...) = tuple(args...)
tup (generic function with 1 method)

julia> code_typed(tup, Tuple{Type{Bool},Int}, optimize = false)
1-element Array{Any,1}:
 CodeInfo(
    @ REPL[29]:1 within `tup'
1 ─ %1 = (Core._apply)(Main.tuple, args)::PartialTuple(Tuple{DataType,Int64}, Any[Const(Bool, false), Int64])
└──      return %1
) => Tuple{DataType,Int64}

julia> code_typed(Cassette.overdub, Tuple{typeof(ctx), typeof(tup), Type{Bool},Int}, optimize = false)
1-element Array{Any,1}:
 CodeInfo(
    @ REPL[29]:1 within `tup'
1 ─       (#self# = (Core.getfield)(##overdub_arguments#359, 1))::Const(tup, false)
│   %2  = (Core.getfield)(##overdub_arguments#359, 2)::Const(Bool, false)
│   %3  = (Core.getfield)(##overdub_arguments#359, 3)::Int64
│         (args = (Core.tuple)(%2, %3))::PartialTuple(Tuple{DataType,Int64}, Any[Const(Bool, false), Int64])
│         (Cassette.prehook)(##overdub_context#358, Core._apply, Main.tuple, args::PartialTuple(Tuple{DataType,Int64}, Any[Const(Bool, false), Int64]))::Any
│         (##overdub_tmp#360 = (Cassette.execute)(##overdub_context#358, Core._apply, Main.tuple, args::PartialTuple(Tuple{DataType,Int64}, Any[Const(Bool, false), Int64])))::Const(OverdubInstead(), false)
│   %7  = (##overdub_tmp#360::Const(OverdubInstead(), false) isa Cassette.OverdubInstead)::Const(true, false)
└──       goto #3 if not %7
2 ─       (##overdub_tmp#360 = (Cassette.overdub)(##overdub_context#358, Core._apply, Main.tuple, args::PartialTuple(Tuple{DataType,Int64}, Any[Const(Bool, false), Int64])))::Tuple{DataType,Int64}
3 ┄       (Cassette.posthook)(##overdub_context#358, ##overdub_tmp#360::Tuple{DataType,Int64}, Core._apply, Main.tuple, args::PartialTuple(Tuple{DataType,Int64}, Any[Const(Bool, false), Int64]))::Any
│   %11 = ##overdub_tmp#360::Tuple{DataType,Int64}::Tuple{DataType,Int64}
└──       return %11
) => Tuple{DataType,Int64}

Looks like the returned SSA values from overdubbed _apply calls on tuple aren't annotated as PartialTuple like I would expect (since the non-overdubbed _apply calls to tuple are annotated PartialTuple, as you can see above).

[Keno]

I don't think phi nodes are doing anything special here. They just tmerge together all the incoming values. In any case, you should be able to see this without the effects of optimization if the function you're interested in is the most top level function you pass to code_typed.

[jrevels]

@maleadt looks like your updated MWE is resolved now via JuliaLang/julia#30385.

julia> outer(A::AbstractArray, i...) = inner(Bool, A, i...)
outer (generic function with 1 method)

julia> inner(::Type{Bool}, A::AbstractArray, i...) = nothing
inner (generic function with 1 method)

julia> using Cassette

julia> Cassette.@context Ctx
Cassette.Context{nametype(Ctx),M,P,T,B} where B<:Union{Nothing, IdDict{Module,Dict{Symbol,BindingMeta}}} where P<:Cassette.AbstractPass where T<:Union{Nothing, Tag} where M

julia> overdubbed = (args...) -> Cassette.overdub(Ctx(), outer, args...)
#4 (generic function with 1 method)

julia> using InteractiveUtils

julia> code_warntype(outer, Tuple{Vector{Int}, Int})
Body::Nothing
1 ─     return

julia> code_warntype(overdubbed, Tuple{Vector{Int}, Int})
Body::Nothing
1 ─     return

julia> versioninfo()
Julia Version 1.2.0-DEV.33
Commit 92ac90e5d7 (2018-12-18 20:54 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin17.5.0)
  CPU: Intel(R) Core(TM) i7-7920HQ CPU @ 3.10GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.1 (ORCJIT, skylake)

That improvement should hopefully be backported upcoming Julia patch releases.

Closing this, but I suspect your real use case could run into different troubles due to JuliaLang/julia#28003, which still needs fixing upstream.

Definitely open new issues if/when you find more MWEs, it's super appreciated! Thanks.