Defer global caching of CodeInstance to post-optimization step (#58343)

This PR extracts the caching improvements from
#56687, implemented by @aviatesk.
It essentially defers global caching to the post-optimization step,
giving a temporary cache to the optimizer instead of relying on the
global cache.

The issue with caching globally before optimization is that any errors
occurring within the optimizer may leave a partially initialized
`CodeInstance` in the cache, which was meant to be updated
post-optimization. Exceptions should not be thrown in the native
compilation pipeline, but abstract interpreters extending optimization
routines will frequently encounter some during an iterative development
(see JuliaComputing/DAECompiler.jl#25).

An extra benefit from deferring global caching is that the optimizer can
then more safely update `CodeInstance`s with new inference information,
as is the intent of #56687.

---------

Co-authored-by: Shuhei Kadowaki <[email protected]>
Co-authored-by: Cody Tapscott <[email protected]>

Author: 3 people

Compiler/src/optimize.jl

@@ -147,16 +147,46 @@ struct InliningState{Interp<:AbstractInterpreter}
     edges::Vector{Any}
     world::UInt
     interp::Interp
+    opt_cache::IdDict{MethodInstance,CodeInstance}
 end
-function InliningState(sv::InferenceState, interp::AbstractInterpreter)
-    return InliningState(sv.edges, frame_world(sv), interp)
+function InliningState(sv::InferenceState, interp::AbstractInterpreter,
+                       opt_cache::IdDict{MethodInstance,CodeInstance}=IdDict{MethodInstance,CodeInstance}())
+    return InliningState(sv.edges, frame_world(sv), interp, opt_cache)
 end
-function InliningState(interp::AbstractInterpreter)
-    return InliningState(Any[], get_inference_world(interp), interp)
+function InliningState(interp::AbstractInterpreter,
+                       opt_cache::IdDict{MethodInstance,CodeInstance}=IdDict{MethodInstance,CodeInstance}())
+    return InliningState(Any[], get_inference_world(interp), interp, opt_cache)
+end
+
+struct OptimizerCache{CodeCache}
+    wvc::WorldView{CodeCache}
+    owner
+    opt_cache::IdDict{MethodInstance,CodeInstance}
+    function OptimizerCache(
+        wvc::WorldView{CodeCache},
+        @nospecialize(owner),
+        opt_cache::IdDict{MethodInstance,CodeInstance}) where CodeCache
+        new{CodeCache}(wvc, owner, opt_cache)
+    end
+end
+function get((; wvc, owner, opt_cache)::OptimizerCache, mi::MethodInstance, default)
+    if haskey(opt_cache, mi)
+        codeinst = opt_cache[mi]
+        @assert codeinst.min_world ≤ wvc.worlds.min_world &&
+                wvc.worlds.max_world ≤ codeinst.max_world &&
+                codeinst.owner === owner
+        @assert isdefined(codeinst, :inferred) && codeinst.inferred === nothing
+        return codeinst
+    end
+    return get(wvc, mi, default)
 end
 
 # get `code_cache(::AbstractInterpreter)` from `state::InliningState`
-code_cache(state::InliningState) = WorldView(code_cache(state.interp), state.world)
+function code_cache(state::InliningState)
+    cache = WorldView(code_cache(state.interp), state.world)
+    owner = cache_owner(state.interp)
+    return OptimizerCache(cache, owner, state.opt_cache)
+end
 
 mutable struct OptimizationResult
     ir::IRCode
@@ -183,13 +213,15 @@ mutable struct OptimizationState{Interp<:AbstractInterpreter}
     bb_vartables::Vector{Union{Nothing,VarTable}}
     insert_coverage::Bool
 end
-function OptimizationState(sv::InferenceState, interp::AbstractInterpreter)
-    inlining = InliningState(sv, interp)
+function OptimizationState(sv::InferenceState, interp::AbstractInterpreter,
+                           opt_cache::IdDict{MethodInstance,CodeInstance}=IdDict{MethodInstance,CodeInstance}())
+    inlining = InliningState(sv, interp, opt_cache)
     return OptimizationState(sv.linfo, sv.src, nothing, sv.stmt_info, sv.mod,
                              sv.sptypes, sv.slottypes, inlining, sv.cfg,
                              sv.unreachable, sv.bb_vartables, sv.insert_coverage)
 end
-function OptimizationState(mi::MethodInstance, src::CodeInfo, interp::AbstractInterpreter)
+function OptimizationState(mi::MethodInstance, src::CodeInfo, interp::AbstractInterpreter,
+                           opt_cache::IdDict{MethodInstance,CodeInstance}=IdDict{MethodInstance,CodeInstance}())
     # prepare src for running optimization passes if it isn't already
     nssavalues = src.ssavaluetypes
     if nssavalues isa Int
@@ -209,7 +241,7 @@ function OptimizationState(mi::MethodInstance, src::CodeInfo, interp::AbstractIn
     mod = isa(def, Method) ? def.module : def
     # Allow using the global MI cache, but don't track edges.
     # This method is mostly used for unit testing the optimizer
-    inlining = InliningState(interp)
+    inlining = InliningState(interp, opt_cache)
     cfg = compute_basic_blocks(src.code)
     unreachable = BitSet()
     bb_vartables = Union{VarTable,Nothing}[]

Compiler/src/typeinfer.jl

@@ -110,6 +110,7 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState, validation
     elseif isdefined(result, :ci)
         edges = result_edges(interp, caller)
         ci = result.ci
+        mi = result.linfo
         # if we aren't cached, we don't need this edge
         # but our caller might, so let's just make it anyways
         if last(result.valid_worlds) >= validation_world
@@ -143,22 +144,27 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState, validation
                     resize!(inferred_result.slottypes::Vector{Any}, nslots)
                     resize!(inferred_result.slotnames, nslots)
                 end
-                inferred_result = maybe_compress_codeinfo(interp, result.linfo, inferred_result)
+                inferred_result = maybe_compress_codeinfo(interp, mi, inferred_result)
                 result.is_src_volatile = false
             elseif ci.owner === nothing
                 # The global cache can only handle objects that codegen understands
                 inferred_result = nothing
             end
         end
         if debuginfo === nothing
-            debuginfo = DebugInfo(result.linfo)
+            debuginfo = DebugInfo(mi)
         end
+        min_world, max_world = first(result.valid_worlds), last(result.valid_worlds)
+        ipo_effects = encode_effects(result.ipo_effects)
         time_now = _time_ns()
         time_self_ns = caller.time_self_ns + (time_now - time_before)
         time_total = (time_now - caller.time_start - caller.time_paused) * 1e-9
         ccall(:jl_update_codeinst, Cvoid, (Any, Any, Int32, UInt, UInt, UInt32, Any, Float64, Float64, Float64, Any, Any),
-            ci, inferred_result, const_flag, first(result.valid_worlds), last(result.valid_worlds), encode_effects(result.ipo_effects),
+            ci, inferred_result, const_flag, min_world, max_world, ipo_effects,
             result.analysis_results, time_total, caller.time_caches, time_self_ns * 1e-9, debuginfo, edges)
+        if is_cached(caller) # CACHE_MODE_GLOBAL
+            cache_result!(interp, result, ci)
+        end
         engine_reject(interp, ci)
         codegen = codegen_cache(interp)
         if !discard_src && codegen !== nothing && (isa(uncompressed, CodeInfo) || isa(uncompressed, OptimizationState))
@@ -171,7 +177,6 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState, validation
                 # This is necessary to get decent bootstrapping performance
                 # when compiling the compiler to inject everything eagerly
                 # where codegen can start finding and using it right away
-                mi = result.linfo
                 if mi.def isa Method && isa_compileable_sig(mi) && is_cached(caller)
                     ccall(:jl_add_codeinst_to_jit, Cvoid, (Any, Any), ci, uncompressed)
                 end
@@ -181,6 +186,11 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState, validation
     return nothing
 end
 
+function cache_result!(interp::AbstractInterpreter, result::InferenceResult, ci::CodeInstance)
+    mi = result.linfo
+    code_cache(interp)[mi] = ci
+end
+
 function finish!(interp::AbstractInterpreter, mi::MethodInstance, ci::CodeInstance, src::CodeInfo)
     user_edges = src.edges
     edges = user_edges isa SimpleVector ? user_edges : user_edges === nothing ? Core.svec() : Core.svec(user_edges...)
@@ -215,11 +225,13 @@ function finish!(interp::AbstractInterpreter, mi::MethodInstance, ci::CodeInstan
 end
 
 function finish_nocycle(::AbstractInterpreter, frame::InferenceState, time_before::UInt64)
-    finishinfer!(frame, frame.interp, frame.cycleid)
+    opt_cache = IdDict{MethodInstance,CodeInstance}()
+    finishinfer!(frame, frame.interp, frame.cycleid, opt_cache)
     opt = frame.result.src
     if opt isa OptimizationState # implies `may_optimize(caller.interp) === true`
         optimize(frame.interp, opt, frame.result)
     end
+    empty!(opt_cache)
     validation_world = get_world_counter()
     finish!(frame.interp, frame, validation_world, time_before)
     if isdefined(frame.result, :ci)
@@ -249,10 +261,11 @@ function finish_cycle(::AbstractInterpreter, frames::Vector{AbsIntState}, cyclei
         cycle_valid_worlds = intersect(cycle_valid_worlds, caller.world.valid_worlds)
         cycle_valid_effects = merge_effects(cycle_valid_effects, caller.ipo_effects)
     end
+    opt_cache = IdDict{MethodInstance,CodeInstance}()
     for frameid = cycleid:length(frames)
         caller = frames[frameid]::InferenceState
         adjust_cycle_frame!(caller, cycle_valid_worlds, cycle_valid_effects)
-        finishinfer!(caller, caller.interp, cycleid)
+        finishinfer!(caller, caller.interp, cycleid, opt_cache)
         time_now = _time_ns()
         caller.time_self_ns += (time_now - time_before)
         time_before = time_now
@@ -273,6 +286,7 @@ function finish_cycle(::AbstractInterpreter, frames::Vector{AbsIntState}, cyclei
         caller.time_paused = UInt64(0)
         caller.time_caches = 0.0
     end
+    empty!(opt_cache)
     cycletop = frames[cycleid]::InferenceState
     time_start = cycletop.time_start
     validation_world = get_world_counter()
@@ -434,22 +448,6 @@ function maybe_compress_codeinfo(interp::AbstractInterpreter, mi::MethodInstance
     return ci
 end
 
-function cache_result!(interp::AbstractInterpreter, result::InferenceResult, ci::CodeInstance)
-    @assert isdefined(ci, :inferred)
-    # check if the existing linfo metadata is also sufficient to describe the current inference result
-    # to decide if it is worth caching this right now
-    mi = result.linfo
-    cache = WorldView(code_cache(interp), result.valid_worlds)
-    if haskey(cache, mi)
-        ci = cache[mi]
-        # n.b.: accurate edge representation might cause the CodeInstance for this to be constructed later
-        @assert isdefined(ci, :inferred)
-        return false
-    end
-    code_cache(interp)[mi] = ci
-    return true
-end
-
 function cycle_fix_limited(@nospecialize(typ), sv::InferenceState, cycleid::Int)
     if typ isa LimitedAccuracy
         frames = sv.callstack::Vector{AbsIntState}
@@ -579,7 +577,8 @@ const empty_edges = Core.svec()
 
 # inference completed on `me`
 # update the MethodInstance
-function finishinfer!(me::InferenceState, interp::AbstractInterpreter, cycleid::Int)
+function finishinfer!(me::InferenceState, interp::AbstractInterpreter, cycleid::Int,
+                      opt_cache::IdDict{MethodInstance, CodeInstance})
     # prepare to run optimization passes on fulltree
     @assert isempty(me.ip)
     # inspect whether our inference had a limited result accuracy,
@@ -635,7 +634,7 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter, cycleid::
                 # disable optimization if we've already obtained very accurate result
                 !result_is_constabi(interp, result)
         if doopt
-            result.src = OptimizationState(me, interp)
+            result.src = OptimizationState(me, interp, opt_cache)
         else
             result.src = me.src # for reflection etc.
         end
@@ -670,23 +669,40 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter, cycleid::
             rettype_const = nothing
             const_flags = 0x0
         end
+
         di = nothing
         edges = empty_edges # `edges` will be updated within `finish!`
         ci = result.ci
+        min_world, max_world = first(result.valid_worlds), last(result.valid_worlds)
         ccall(:jl_fill_codeinst, Cvoid, (Any, Any, Any, Any, Int32, UInt, UInt, UInt32, Any, Any, Any),
             ci, widenconst(result_type), widenconst(result.exc_result), rettype_const, const_flags,
-            first(result.valid_worlds), last(result.valid_worlds),
+            min_world, max_world,
             encode_effects(result.ipo_effects), result.analysis_results, di, edges)
         if is_cached(me) # CACHE_MODE_GLOBAL
-            cached_result = cache_result!(me.interp, result, ci)
-            if !cached_result
+            already_cached = is_already_cached(me.interp, result, ci)
+            if already_cached
                 me.cache_mode = CACHE_MODE_VOLATILE
+            else
+                opt_cache[result.linfo] = ci
             end
         end
     end
     nothing
 end
 
+function is_already_cached(interp::AbstractInterpreter, result::InferenceResult, ci::CodeInstance)
+    # check if the existing linfo metadata is also sufficient to describe the current inference result
+    # to decide if it is worth caching this right now
+    mi = result.linfo
+    cache = WorldView(code_cache(interp), result.valid_worlds)
+    if haskey(cache, mi)
+        # n.b.: accurate edge representation might cause the CodeInstance for this to be constructed later
+        @assert isdefined(cache[mi], :inferred)
+        return true
+    end
+    return false
+end
+
 # Iterate a series of back-edges that need registering, based on the provided forward edge list.
 # Back-edges are returned as (invokesig, item), where the item is a Binding, MethodInstance, or
 # MethodTable.