Allocations when invoking a `CodeInstance` from a generated function

[MasonProtter]

I was playing with the mechanism introduced in #56660, and was hoping that I could produce a CodeInstance in a generated function, and then invoke it

Here's a MWE constructing an interpreter and creating a codeinstance with it in a generated function, and then invoking that code instance:

definitions `contextual_mwe.jl`

using Core: MethodMatch, MethodInstance, CodeInstance, Compiler
using Core.Compiler:
    _methods_by_ftype,
    InferenceParams,
    get_world_counter,
    tls_world_age,
    InferenceResult,
    typeinf,
    InternalMethodTable,
    InferenceState,
    NativeInterpreter,
    AbstractInterpreter,
    OptimizationParams,
    MethodTableView,
    OverlayMethodTable,
    WorldRange,
    finish,
    InferenceResult,
    OptimizationState,
    OptimizationParams,
    optimize,
    store_backedges,
    finish!,
    CodeInfo,
    isexpr,
    argextype,
    singleton_type,
    Builtin,
    GlobalRef,
    CodeInstance,
    ir_to_codeinf!,
    typeinf_ext_toplevel,
    IRCode,
    argextype,
    widenconst

const CC = Core.Compiler

using Base:
    specialize_method,
    isexpr

function with_ctx_interpreter(f, args...)
    cinst = generated_ci_in_absint(f, args)
    invoke(f, cinst, args...)
end

struct MyCtx end

struct ContextualInterpreter <: AbstractInterpreter
    inf_cache::Vector{InferenceResult}
    world::UInt
    inf_params::InferenceParams
    opt_params::OptimizationParams
    codegen_cache::IdDict{CodeInstance, CodeInfo}
    function ContextualInterpreter(world::UInt,
                                   ip::InferenceParams,
                                   op::OptimizationParams)
        @assert world <= Base.get_world_counter()
        return new(
            InferenceResult[],
            world,
            ip,
            op,
            IdDict{CodeInstance, CodeInfo}()
        )
    end
end
function ContextualInterpreter(;world=Base.get_world_counter(),
                               inf_params=InferenceParams(),
                               opt_params=OptimizationParams()) 
    ContextualInterpreter(world, inf_params, opt_params)
end

Core.Compiler.InferenceParams(interp::ContextualInterpreter) = interp.inf_params
Core.Compiler.OptimizationParams(interp::ContextualInterpreter) = interp.opt_params
Core.Compiler.get_inference_world(interp::ContextualInterpreter) = interp.world
Core.Compiler.get_inference_cache(interp::ContextualInterpreter) = interp.inf_cache
Core.Compiler.cache_owner(interp::ContextualInterpreter) = MyCtx()
Core.Compiler.codegen_cache(interp::ContextualInterpreter) = interp.codegen_cache


@noinline function Core.OptimizedGenerics.CompilerPlugins.typeinf(::MyCtx, mi::MethodInstance, source_mode::UInt8)
    # Base.invoke_in_world(which(Core.OptimizedGenerics.CompilerPlugins.typeinf, Tuple{ContextOwner, MethodInstance, UInt8}).primary_world,
    Compiler.typeinf_ext_toplevel(ContextualInterpreter(; world=Base.tls_world_age()),
                                  mi, source_mode)
end

@noinline function Core.OptimizedGenerics.CompilerPlugins.typeinf_edge(::MyCtx, mi::MethodInstance, parent_frame::Compiler.InferenceState, world::UInt, source_mode::UInt8)
    interp = ContextualInterpreter{Owner}(; world)
    Compiler.typeinf_edge(interp, mi.def, mi.specTypes, Core.svec(),
                          parent_frame, false, false)
end


function expr_to_codeinfo(m::Module, argnames, spnames, sp, e::Expr)
    lam = Expr(:lambda, argnames,
               Expr(Symbol("scope-block"),
                    Expr(:block,
                         Expr(:return,
                              Expr(:block,
                                   e,
                                   )))))
    ex = if spnames === nothing || isempty(spnames)
        lam
    else
        Expr(Symbol("with-static-parameters"), lam, spnames...)
    end
    ci = Base.generated_body_to_codeinfo(ex, @__MODULE__(), false)
    @assert ci isa Core.CodeInfo "Failed to create a CodeInfo from the given expression. This might mean it contains a closure or comprehension?\n Offending expression: $e"
    ci
end

function generated_ci_in_absint_body(world::UInt, lnn, this, f, args)
    sig = Type{Tuple{f, args.parameters...}}
    sig isa Type{<:Type{<:Tuple}} || error()
    tt = sig.parameters[1]
    interp = ContextualInterpreter(; world)

    match, valid_worlds = Core.Compiler.findsup(tt, Core.Compiler.method_table(interp))
    if match === nothing
        error(lazy"Unable to find matching $tt")
    end
    mi = specialize_method(match.method, match.spec_types, match.sparams)::MethodInstance
    
    cinst = Core.OptimizedGenerics.CompilerPlugins.typeinf(MyCtx(), mi, Compiler.SOURCE_MODE_ABI)
    
    ci = expr_to_codeinfo(@__MODULE__(), [Symbol("#self#"), :f, :args], [], (), :(return $cinst))
    
    matches = Base._methods_by_ftype(sig, -1, world)
    if !isnothing(matches)
        ci.edges = Core.MethodInstance[]
        for match in Base._methods_by_ftype(sig, -1, world)
            mi = Base.specialize_method(match) 
            push!(ci.edges, mi)
        end
    end
    return ci
end

function refresh_generated_ci_in_absint()
    @eval function generated_ci_in_absint(f, args)
        $(Expr(:meta, :generated_only))
        $(Expr(:meta, :generated, generated_ci_in_absint_body))
    end
end
refresh_generated_ci_in_absint()

This works

julia> includet("contextual_mwe.jl")

julia> with_ctx_interpreter(1.0) do x
           sin(x) + 1
       end
1.8414709848078965

but it allocates when the compiler can't constant fold the whole thing away:

julia> @btime with_ctx_interpreter(y) do x
           sin(x) + 1
       end setup=y=rand()
  31.316 ns (3 allocations: 48 bytes)
1.0876985847957277

I see this allocation behaviour on both 1.12 and nightly.

[vtjnash]

Accessing any copy of Compiler from a generated function is strict UB (can cause arbitrary corruption of the code caches and incorrect results or arbitrary memory corruption), so it isn't too surprising that it goes poorly

[MasonProtter]

Isn't the point of this mechanism that it lets us do it from a separate compiler though? I guess your point is that even copies of the compiler are not allowed? Unfortunate.

If this thing shouldn't be allowed to work, then does that mean that #56650 should be revived?

[vtjnash]

The point of that work is you can call invoke(f, get_codeinst_at_runtime(f, args...), args...) and avoid the UB of @generated

[Keno]

Accessing any copy of Compiler from a generated function is strict UB (can cause arbitrary corruption of the code caches and incorrect results or arbitrary memory corruption), so it isn't too surprising that it goes poorly

Is your concern that the Compiler is not guaranteed to produce stable answers or is is about the caching of non-native owners?

[Keno]

In any case, I suspect there's probably some codegen issue here that can be fixed, so we can separate the larger philosophical point.

[vtjnash]

Accessing the cache will cause memory corruption since it is not able to access the get_world_counter value, as being able to do so would violate the constraints on having a stable answer, as required to avoid memory corruption.

[Keno]

I gets the invocation world age though and can use that for its computations. I agree that in practice there is a possibility of violating that invariant due to cache ordering artifact, but wouldn't it be fine in theory? Isn't the reason that we pass the world to the generated function precisely so that it can introspect the properties in a well defined world age?