implement Jameson's suggestion

Author: aviatesk

Compiler/src/optimize.jl

@@ -143,16 +143,48 @@ 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},
+        owner,
+        opt_cache::IdDict{MethodInstance,CodeInstance}) where CodeCache
+        @nospecialize owner
+        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]
+        if (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
+    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 OptimizationState{Interp<:AbstractInterpreter}
     linfo::MethodInstance
@@ -168,13 +200,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
@@ -194,7 +228,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

@@ -106,6 +106,7 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState)
     #@assert last(result.valid_worlds) <= get_world_counter() || isempty(caller.edges)
     if isdefined(result, :ci)
         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) >= get_world_counter()
@@ -132,15 +133,15 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState)
                 end
                 di = inferred_result.debuginfo
                 uncompressed = inferred_result
-                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 !@isdefined di
-            di = DebugInfo(result.linfo)
+            di = DebugInfo(mi)
         end
         min_world, max_world = first(result.valid_worlds), last(result.valid_worlds)
         ipo_effects = encode_effects(result.ipo_effects)
@@ -149,6 +150,9 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState)
             UInt32, Any, Any, Any),
             ci, inferred_result, rettype, exctype, rettype_const, const_flags, min_world, max_world,
             ipo_effects, result.analysis_results, di, edges)
+        if is_cached(caller) # CACHE_MODE_GLOBAL
+            cache_result!(interp, mi, ci)
+        end
         engine_reject(interp, ci)
         if !discard_src && isdefined(interp, :codegen) && uncompressed isa CodeInfo
             # record that the caller could use this result to generate code when required, if desired, to avoid repeating n^2 work
@@ -157,7 +161,6 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState)
                 # 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)
                     ccall(:jl_add_codeinst_to_jit, Cvoid, (Any, Any), ci, uncompressed)
                 end
@@ -167,6 +170,10 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState)
     return nothing
 end
 
+function cache_result!(interp::AbstractInterpreter, mi::MethodInstance, ci::CodeInstance)
+    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...)
@@ -200,11 +207,13 @@ function finish!(interp::AbstractInterpreter, mi::MethodInstance, ci::CodeInstan
 end
 
 function finish_nocycle(::AbstractInterpreter, frame::InferenceState)
-    finishinfer!(frame, frame.interp)
+    opt_cache = IdDict{MethodInstance,CodeInstance}()
+    finishinfer!(frame, frame.interp, 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)
     finish!(frame.interp, frame)
     if frame.cycleid != 0
         frames = frame.callstack::Vector{AbsIntState}
@@ -227,10 +236,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)
+        finishinfer!(caller, caller.interp, opt_cache)
     end
     for frameid = cycleid:length(frames)
         caller = frames[frameid]::InferenceState
@@ -239,6 +249,7 @@ function finish_cycle(::AbstractInterpreter, frames::Vector{AbsIntState}, cyclei
             optimize(caller.interp, opt, caller.result)
         end
     end
+    empty!(opt_cache)
     for frameid = cycleid:length(frames)
         caller = frames[frameid]::InferenceState
         finish!(caller.interp, caller)
@@ -285,22 +296,6 @@ function maybe_compress_codeinfo(interp::AbstractInterpreter, mi::MethodInstance
     end
 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)
     if typ isa LimitedAccuracy
         if sv.parentid === 0
@@ -428,7 +423,8 @@ const empty_edges = Core.svec()
 
 # inference completed on `me`
 # update the MethodInstance
-function finishinfer!(me::InferenceState, interp::AbstractInterpreter)
+function finishinfer!(me::InferenceState, interp::AbstractInterpreter,
+                      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,
@@ -481,7 +477,7 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter)
                 # 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
@@ -502,9 +498,11 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter)
             ci, rettype, exctype, rettype_const, const_flags, 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
@@ -551,6 +549,19 @@ end
     return ResultForCache(rettype, exctype, rettype_const, const_flags)
 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
+
 # record the backedges
 function store_backedges(caller::CodeInstance, edges::SimpleVector)
     isa(caller.def.def, Method) || return # don't add backedges to toplevel method instance

Compiler/src/types.jl

@@ -1,5 +1,4 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
-#
 
 const WorkThunk = Any
 # #@eval struct WorkThunk

test/compileall.jl

@@ -1,3 +1,5 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+
 # This test builds a full system image, so it can take a little while.
 # We make it a separate test target here, so that it can run in parallel
 # with the rest of the tests.