Refactor generated functions to use an empty CodeInfo. (#2663)

Author: vchuravy

src/analyses/activity.jl

@@ -451,48 +451,20 @@ Base.@nospecializeinfer @inline function active_reg(@nospecialize(ST::Type), wor
     return result
 end
 
+function active_reg_nothrow end
+
 function active_reg_nothrow_generator(world::UInt, source::Union{Method, LineNumberNode}, T, self, _)
     @nospecialize
     result = active_reg(T, world)
-
-    # create an empty CodeInfo to return the result
-    ci = ccall(:jl_new_code_info_uninit, Ref{Core.CodeInfo}, ())
-    
-    @static if isdefined(Core, :DebugInfo)
-        # TODO: Add proper debug info
-        ci.debuginfo = Core.DebugInfo(:none)
-    else
-        ci.codelocs = Int32[]
-        ci.linetable = [
-            Core.Compiler.LineInfoNode(@__MODULE__, :active_reg_nothrow, source.file, Int32(source.line), Int32(0))
-        ]
-    end
     check_activity_cache_invalidations(world)
-    ci.min_world = world
-    ci.max_world = typemax(UInt)
 
-    edges = Any[]
-    # Create the edge for the "query"
-    # TODO: Check if we can use `Tuple{typeof(EnzymeRules.inactive_type), T}` directly
+    slotnames = Core.svec(Symbol("#self#"), :T)
+    code = Any[Core.Compiler.ReturnNode(result)]
+    ci = create_fresh_codeinfo(active_reg_nothrow, source, world, slotnames, code)
+    
+    ci.edges = Any[]
     inactive_type_sig = Tuple{typeof(EnzymeRules.inactive_type), Type}
-    push!(edges, ccall(:jl_method_table_for, Any, (Any,), inactive_type_sig)::Core.MethodTable)
-    push!(edges, inactive_type_sig)
-
-    ci.edges = edges
-
-    # prepare the slots
-    ci.slotnames = Symbol[Symbol("#self#"), :t]
-    ci.slotflags = UInt8[0x00 for i = 1:2]
-
-    # return the result
-    ci.code = Any[Core.Compiler.ReturnNode(result)]
-    ci.ssaflags = UInt32[0x00]   # Julia's native compilation pipeline (and its verifier) expects `ssaflags` to be the same length as `code`
-    @static if isdefined(Core, :DebugInfo)
-    else
-        push!(ci.codelocs, 1)
-    end
-
-    ci.ssavaluetypes = 1
+    add_edge!(ci.edges, inactive_type_sig)
 
     return ci
 end

src/compiler.jl

@@ -27,7 +27,9 @@ import Enzyme:
     FnTypeInfo,
     Logic,
     allocatedinline,
-    ismutabletype
+    ismutabletype,
+    create_fresh_codeinfo,
+    add_edge!
 using Enzyme
 
 import EnzymeCore
@@ -6227,88 +6229,37 @@ end
     end
 end
 
+function thunk end
+
 function thunk_generator(world::UInt, source::Union{Method, LineNumberNode}, @nospecialize(FA::Type), @nospecialize(A::Type), @nospecialize(TT::Type), Mode::Enzyme.API.CDerivativeMode, Width::Int, @nospecialize(ModifiedBetween::(NTuple{N, Bool} where N)), ReturnPrimal::Bool, ShadowInit::Bool, @nospecialize(ABI::Type), ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, @nospecialize(self), @nospecialize(fakeworld), @nospecialize(fa::Type), @nospecialize(a::Type), @nospecialize(tt::Type), @nospecialize(mode::Type), @nospecialize(width::Type), @nospecialize(modifiedbetween::Type), @nospecialize(returnprimal::Type), @nospecialize(shadowinit::Type), @nospecialize(abi::Type), @nospecialize(erriffuncwritten::Type), @nospecialize(runtimeactivity::Type), @nospecialize(strongzero::Type))
     @nospecialize
 
-    parmnames = (:fakeworld, :fa, :a, :tt, :mode, :width, :modifiedbetween, :returnprimal, :shadowinit, :abi, :erriffuncwritten, :runtimeactivity, :strongzero)
-    stub = Core.GeneratedFunctionStub(identity, Core.svec(:methodinstance, parmnames...), Core.svec())
+    slotnames = Core.svec(Symbol("#self#"), 
+                    :fakeworld, :fa, :a, :tt, :mode, :width,
+                    :modifiedbetween, :returnprimal, :shadowinit,
+                    :abi, :erriffuncwritten, :runtimeactivity, :strongzero)
+    stub = Core.GeneratedFunctionStub(thunk, slotnames, Core.svec())
 
     ft = eltype(FA)
     primal_tt = Tuple{map(eltype, TT.parameters)...}
     # look up the method match
-    method_error = :(throw(MethodError($ft, $primal_tt, $world)))
 
     min_world = Ref{UInt}(typemin(UInt))
     max_world = Ref{UInt}(typemax(UInt))
 
     mi = my_methodinstance(Mode == API.DEM_ForwardMode ? Forward : Reverse, ft, primal_tt, world, min_world, max_world)
 
-    mi === nothing && return stub(world, source, method_error)
+    mi === nothing && return stub(world, source, :(throw(MethodError($ft, $primal_tt, $world))))
 
     check_activity_cache_invalidations(world)
 
-    min_world2 = Ref{UInt}(typemin(UInt))
-    max_world2 = Ref{UInt}(typemax(UInt))
-   
-    mi2 = my_methodinstance(Mode == API.DEM_ForwardMode ? Forward : Reverse, typeof(Base.identity), Tuple{Nothing}, world, min_world2, max_world2)
-
-    ci = Core.Compiler.retrieve_code_info(mi2, world)::Core.Compiler.CodeInfo
-
-    # prepare a new code info
-    new_ci = copy(ci)
-    empty!(new_ci.code)
-    @static if isdefined(Core, :DebugInfo)
-      new_ci.debuginfo = Core.DebugInfo(:none)
-    else
-      empty!(new_ci.codelocs)
-      resize!(new_ci.linetable, 1)                # see note below
-    end
-    empty!(new_ci.ssaflags)
-    new_ci.ssavaluetypes = 0
-    # new_ci.min_world = min_world[]
-    new_ci.min_world = world
-    new_ci.max_world = max_world[]
-
-    edges = Any[mi]
-
-    if Mode == API.DEM_ForwardMode
-        fwd_sig = Tuple{typeof(EnzymeRules.forward), <:EnzymeRules.FwdConfig, <:Enzyme.EnzymeCore.Annotation, Type{<:Enzyme.EnzymeCore.Annotation},Vararg{Enzyme.EnzymeCore.Annotation}}
-        push!(edges, ccall(:jl_method_table_for, Any, (Any,), fwd_sig)::Core.MethodTable)
-        push!(edges, fwd_sig)
-    else
-        rev_sig = Tuple{typeof(EnzymeRules.augmented_primal), <:EnzymeRules.RevConfig, <:Enzyme.EnzymeCore.Annotation, Type{<:Enzyme.EnzymeCore.Annotation},Vararg{Enzyme.EnzymeCore.Annotation}}
-        push!(edges, ccall(:jl_method_table_for, Any, (Any,), rev_sig)::Core.MethodTable)
-        push!(edges, rev_sig)
-        
-        rev_sig = Tuple{typeof(EnzymeRules.reverse), <:EnzymeRules.RevConfig, <:Enzyme.EnzymeCore.Annotation, Union{Type{<:Enzyme.EnzymeCore.Annotation}, Enzyme.EnzymeCore.Active}, Any, Vararg{Enzyme.EnzymeCore.Annotation}}
-        push!(edges, ccall(:jl_method_table_for, Any, (Any,), rev_sig)::Core.MethodTable)
-        push!(edges, rev_sig)
-    end
-    
-    ina_sig = Tuple{typeof(EnzymeRules.inactive), Vararg{Any}}
-    push!(edges, ccall(:jl_method_table_for, Any, (Any,), ina_sig)::Core.MethodTable)
-    push!(edges, ina_sig)
-    
-    for gen_sig in (
-        Tuple{typeof(EnzymeRules.inactive_noinl), Vararg{Any}},
-        Tuple{typeof(EnzymeRules.noalias), Vararg{Any}},
-        Tuple{typeof(EnzymeRules.inactive_type), Type},
-    )
-        push!(edges, ccall(:jl_method_table_for, Any, (Any,), gen_sig)::Core.MethodTable)
-        push!(edges, gen_sig)
-    end
-
-    new_ci.edges = edges
-
-    # XXX: setting this edge does not give us proper method invalidation, see
-    #      JuliaLang/julia#34962 which demonstrates we also need to "call" the kernel.
-    #      invoking `code_llvm` also does the necessary codegen, as does calling the
-    #      underlying C methods -- which GPUCompiler does, so everything Just Works.
+    edges = Any[]
+    add_edge!(edges, mi)
 
     ts_ctx = JuliaContext()
     ctx = context(ts_ctx)
     activate(ctx)
-    res = try
+    result = try
         thunkbase(
             mi,
             world,
@@ -6331,25 +6282,35 @@ function thunk_generator(world::UInt, source::Union{Method, LineNumberNode}, @no
         dispose(ts_ctx)
     end
 
-    # prepare the slots
-    new_ci.slotnames = Symbol[Symbol("#self#"), parmnames...]
-    new_ci.slotflags = UInt8[0x00 for i = 1:length(new_ci.slotnames)]
+    code = Any[Core.Compiler.ReturnNode(result)]
+    ci = create_fresh_codeinfo(thunk, source, world, slotnames, code)
 
-    # return the codegen world age
-    push!(new_ci.code, Core.Compiler.ReturnNode(res))
-    push!(new_ci.ssaflags, 0x00)   # Julia's native compilation pipeline (and its verifier) expects `ssaflags` to be the same length as `code`
-    @static if isdefined(Core, :DebugInfo)
+
+
+    if Mode == API.DEM_ForwardMode
+        fwd_sig = Tuple{typeof(EnzymeRules.forward), <:EnzymeRules.FwdConfig, <:Enzyme.EnzymeCore.Annotation, Type{<:Enzyme.EnzymeCore.Annotation},Vararg{Enzyme.EnzymeCore.Annotation}}
+        add_edge!(edges, fwd_sig)
     else
-      push!(new_ci.codelocs, 1)   # see note below
+        rev_sig = Tuple{typeof(EnzymeRules.augmented_primal), <:EnzymeRules.RevConfig, <:Enzyme.EnzymeCore.Annotation, Type{<:Enzyme.EnzymeCore.Annotation},Vararg{Enzyme.EnzymeCore.Annotation}}
+        add_edge!(edges, rev_sig)
+        
+        rev_sig = Tuple{typeof(EnzymeRules.reverse), <:EnzymeRules.RevConfig, <:Enzyme.EnzymeCore.Annotation, Union{Type{<:Enzyme.EnzymeCore.Annotation}, Enzyme.EnzymeCore.Active}, Any, Vararg{Enzyme.EnzymeCore.Annotation}}
+        add_edge!(edges, rev_sig)
+    end
+    
+    ina_sig = Tuple{typeof(EnzymeRules.inactive), Vararg{Any}}
+    add_edge!(edges, ina_sig)
+    
+    for gen_sig in (
+        Tuple{typeof(EnzymeRules.inactive_noinl), Vararg{Any}},
+        Tuple{typeof(EnzymeRules.noalias), Vararg{Any}},
+        Tuple{typeof(EnzymeRules.inactive_type), Type},
+    )
+        add_edge!(edges, gen_sig)
     end
-    new_ci.ssavaluetypes += 1
-
-    # NOTE: we keep the first entry of the original linetable, and use it for location info
-    #       on the call to check_cache. we can't not have a codeloc (using 0 causes
-    #       corruption of the back trace), and reusing the target function's info
-    #       has as advantage that we see the name of the kernel in the backtraces.
 
-    return new_ci
+    ci.edges = edges
+    return ci
 end
 
 @eval @inline function thunk(
@@ -6386,48 +6347,30 @@ end
 
 import GPUCompiler: deferred_codegen_jobs
 
+function deferred_id_codegen end
+
 function deferred_id_generator(world::UInt, source::Union{Method, LineNumberNode}, @nospecialize(FA::Type), @nospecialize(A::Type), @nospecialize(TT::Type), Mode::Enzyme.API.CDerivativeMode, Width::Int, @nospecialize(ModifiedBetween::(NTuple{N, Bool} where N)), ReturnPrimal::Bool, ShadowInit::Bool, @nospecialize(ExpectedTapeType::Type), ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, @nospecialize(self), @nospecialize(fa::Type), @nospecialize(a::Type), @nospecialize(tt::Type), @nospecialize(mode::Type), @nospecialize(width::Type), @nospecialize(modifiedbetween::Type), @nospecialize(returnprimal::Type), @nospecialize(shadowinit::Type), @nospecialize(expectedtapetype::Type), @nospecialize(erriffuncwritten::Type), @nospecialize(runtimeactivity::Type), @nospecialize(strongzero::Type))
     @nospecialize
 
-    parmnames = (:fa, :a, :tt, :mode, :width, :modifiedbetween, :returnprimal, :shadowinit, :expectedtapetype, :erriffuncwritten, :runtimeactivity, :strongzero)
-    stub = Core.GeneratedFunctionStub(identity, Core.svec(:methodinstance, parmnames...), Core.svec())
+    slotnames = Core.svec(Symbol("#self#"),
+                          :fa, :a, :tt, :mode, :width, :modifiedbetween,
+                          :returnprimal, :shadowinit, :expectedtapetype,
+                          :erriffuncwritten, :runtimeactivity, :strongzero)
+
+    stub = Core.GeneratedFunctionStub(deferred_id_generator, slotnames, Core.svec())
 
     ft = eltype(FA)
     primal_tt = Tuple{map(eltype, TT.parameters)...}
     # look up the method match
-    method_error = :(throw(MethodError($ft, $primal_tt, $world)))
 
     min_world = Ref{UInt}(typemin(UInt))
     max_world = Ref{UInt}(typemax(UInt))
 
     mi = my_methodinstance(Mode == API.DEM_ForwardMode ? Forward : Reverse, ft, primal_tt, world, min_world, max_world)
 
-    mi === nothing && return stub(world, source, method_error)
+    mi === nothing && return stub(world, source, :(throw(MethodError($ft, $primal_tt, $world))))
 
-    ci = Core.Compiler.retrieve_code_info(mi, world)::Core.Compiler.CodeInfo
-
-    # prepare a new code info
-    new_ci = copy(ci)
-    empty!(new_ci.code)
-    @static if isdefined(Core, :DebugInfo)
-      new_ci.debuginfo = Core.DebugInfo(:none)
-    else
-      empty!(new_ci.codelocs)
-      resize!(new_ci.linetable, 1)                # see note below
-    end
-    empty!(new_ci.ssaflags)
-    new_ci.ssavaluetypes = 0
-    # new_ci.min_world = min_world[]
-    new_ci.min_world = world
-    new_ci.max_world = max_world[]
-    new_ci.edges = Core.MethodInstance[mi]
-    # XXX: setting this edge does not give us proper method invalidation, see
-    #      JuliaLang/julia#34962 which demonstrates we also need to "call" the kernel.
-    #      invoking `code_llvm` also does the necessary codegen, as does calling the
-    #      underlying C methods -- which GPUCompiler does, so everything Just Works.
-        
     target = EnzymeTarget()
-
     rt2 = if A isa UnionAll
         rrt = primal_return_type_world(Mode == API.DEM_ForwardMode ? Forward : Reverse, world, mi)
 
@@ -6472,25 +6415,12 @@ function deferred_id_generator(world::UInt, source::Union{Method, LineNumberNode
     id = Base.reinterpret(Int, pointer(addr))
     deferred_codegen_jobs[id] = job
 
-    # prepare the slots
-    new_ci.slotnames = Symbol[Symbol("#self#"), parmnames...]
-    new_ci.slotflags = UInt8[0x00 for i = 1:length(new_ci.slotnames)]
-
-    # return the codegen world age
-    push!(new_ci.code, Core.Compiler.ReturnNode(reinterpret(Ptr{Cvoid}, id)))
-    push!(new_ci.ssaflags, 0x00)   # Julia's native compilation pipeline (and its verifier) expects `ssaflags` to be the same length as `code`
-    @static if isdefined(Core, :DebugInfo)
-    else
-      push!(new_ci.codelocs, 1)   # see note below
-    end
-    new_ci.ssavaluetypes += 1
+    code = Any[Core.Compiler.ReturnNode(reinterpret(Ptr{Cvoid}, id))]
+    ci = create_fresh_codeinfo(deferred_id_codegen, source, world, slotnames, code)
 
-    # NOTE: we keep the first entry of the original linetable, and use it for location info
-    #       on the call to check_cache. we can't not have a codeloc (using 0 causes
-    #       corruption of the back trace), and reusing the target function's info
-    #       has as advantage that we see the name of the kernel in the backtraces.
+    ci.edges = Any[mi]
 
-    return new_ci
+    return ci
 end
 
 @eval @inline function deferred_id_codegen(

src/compiler/interpreter.jl

@@ -24,85 +24,28 @@ else
     import Core.Compiler: get_world_counter, get_world_counter as get_inference_world
 end
 
+function rule_backedge_holder end
+
 function rule_backedge_holder_generator(world::UInt, source, self, ft::Type)
     @nospecialize
-    sig = Tuple{typeof(Base.identity), Int}
-    min_world = Ref{UInt}(typemin(UInt))
-    max_world = Ref{UInt}(typemax(UInt))
-    has_ambig = Ptr{Int32}(C_NULL)
-    mthds = Base._methods_by_ftype(
-        sig,
-        nothing,
-        -1, #=lim=#
-        world,
-        false, #=ambig=#
-        min_world,
-        max_world,
-        has_ambig,
-    )
-    mtypes, msp, m = mthds[1]
-    mi = ccall(
-        :jl_specializations_get_linfo,
-        Ref{Core.MethodInstance},
-        (Any, Any, Any),
-        m,
-        mtypes,
-        msp,
-    )
-    ci = Core.Compiler.retrieve_code_info(mi, world)::Core.Compiler.CodeInfo
 
-    # prepare a new code info
-    new_ci = copy(ci)
-    empty!(new_ci.code)
-    @static if isdefined(Core, :DebugInfo)
-      new_ci.debuginfo = Core.DebugInfo(:none)
-    else
-      empty!(new_ci.codelocs)
-      resize!(new_ci.linetable, 1)                # see note below
-    end
-    empty!(new_ci.ssaflags)
-    new_ci.ssavaluetypes = 0
-    new_ci.min_world = min_world[]
-    new_ci.max_world = max_world[]
+    code = Any[Core.Compiler.ReturnNode(world)]
+    ci = Core.Compiler.create_fresh_codeinfo(rule_backedge_holder, source, world, Core.svec(Symbol("#self#"), :ft), code)
 
-    ### TODO: backedge from inactive, augmented_primal, forward, reverse
     edges = Any[]
 
     if ft == typeof(EnzymeRules.augmented_primal)
         sig = Tuple{typeof(EnzymeRules.augmented_primal), <:RevConfig, <:Annotation, Type{<:Annotation},Vararg{Annotation}}
-        push!(edges, ccall(:jl_method_table_for, Any, (Any,), sig))
-        push!(edges, sig)
     elseif ft == typeof(EnzymeRules.forward)
         sig = Tuple{typeof(EnzymeRules.forward), <:FwdConfig, <:Annotation, Type{<:Annotation},Vararg{Annotation}}
-        push!(edges, ccall(:jl_method_table_for, Any, (Any,), sig))
-        push!(edges, sig)
     else
         sig = Tuple{typeof(EnzymeRules.inactive), Vararg{Annotation}}
-        push!(edges, ccall(:jl_method_table_for, Any, (Any,), sig))
-        push!(edges, sig)
     end
+    add_edge!(edges, sig)
 
-    new_ci.edges = edges
-
-    # XXX: setting this edge does not give us proper method invalidation, see
-    #      JuliaLang/julia#34962 which demonstrates we also need to "call" the kernel.
-    #      invoking `code_llvm` also does the necessary codegen, as does calling the
-    #      underlying C methods -- which GPUCompiler does, so everything Just Works.
-
-    # prepare the slots
-    new_ci.slotnames = Symbol[Symbol("#self#"), :ft]
-    new_ci.slotflags = UInt8[0x00 for i = 1:2]
-
-    # return the codegen world age
-    push!(new_ci.code, Core.Compiler.ReturnNode(world))
-    push!(new_ci.ssaflags, 0x00)   # Julia's native compilation pipeline (and its verifier) expects `ssaflags` to be the same length as `code`
-    @static if isdefined(Core, :DebugInfo)
-    else
-      push!(new_ci.codelocs, 1)   # see note below
-    end
-    new_ci.ssavaluetypes += 1
+    ci.edges = edges
 
-    return new_ci
+    return ci
 end
 
 @eval Base.@assume_effects :removable :foldable :nothrow @inline function rule_backedge_holder(ft)