inference: reinfer and track missing code for inlining (#59413)

When code is potentially needed for inlining, but missing for any
reason, be sure to regenerate it during inference with the correct
`ci_meets_requirement` flags (SOURCE_MODE_GET_SOURCE instead of
NOT_REQUIRED) so it is prepared for the optimizer if needed.

This was supposed to be the correct fix for someone else's bug with
missing inlining when expected, but I don't remember what bug it was
anymore. It should however improve handling and tracking of code (and
worlds of code) during inference, which should help for building upon
this base.

```
julia> @atomic Base.method_instance(promote, (Float64,Int)).cache.inferred = 0x22 # mark with inlining cost, deleting code

julia> @code_typed 1.0+1 # this should regenerate the code, but on master instead results in:
CodeInfo(
    @ promotion.jl:433 within `+`
1 ─ %1 =    invoke Base.promote(x::Float64, y::Int64)::Tuple{Float64, Float64}
│   %2 =   builtin Core.getfield(%1, 1)::Float64
│   %3 =   builtin Core.getfield(%1, 2)::Float64
│   @ promotion.jl:433 within `+` @ float.jl:492
│   %4 = intrinsic Base.add_float(%2, %3)::Float64
└──      return %4
) => Float64
```

Author: vtjnash

Compiler/src/abstractinterpretation.jl

@@ -1,26 +1,5 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
-"""
-    struct InternalCodeCache
-
-Internally, each `MethodInstance` keep a unique global cache of code instances
-that have been created for the given method instance, stratified by world age
-ranges. This struct abstracts over access to this cache.
-"""
-struct InternalCodeCache
-    owner::Any # `jl_egal` is used for comparison
-end
-
-function setindex!(cache::InternalCodeCache, ci::CodeInstance, mi::MethodInstance)
-    @assert ci.owner === cache.owner
-    m = mi.def
-    if isa(m, Method)
-        ccall(:jl_push_newly_inferred, Cvoid, (Any,), ci)
-    end
-    ccall(:jl_mi_cache_insert, Cvoid, (Any, Any), mi, ci)
-    return cache
-end
-
 struct WorldRange
     min_world::UInt
     max_world::UInt
@@ -49,39 +28,43 @@ function union(a::WorldRange, b::WorldRange)
 end
 
 """
-    struct WorldView
+    struct InternalCodeCache
 
-Takes a given cache and provides access to the cache contents for the given
-range of world ages, rather than defaulting to the current active world age.
+Internally, each `MethodInstance` keep a unique global cache of code instances
+that have been created for the given method instance, stratified by world age
+ranges. This struct abstracts over access to this cache.
 """
-struct WorldView{Cache}
-    cache::Cache
+struct InternalCodeCache
+    owner::Any # `jl_egal` is used for comparison
     worlds::WorldRange
-    WorldView(cache::Cache, range::WorldRange) where Cache = new{Cache}(cache, range)
+    InternalCodeCache(@nospecialize(owner), wr::WorldRange) = new(owner, wr)
+    InternalCodeCache(@nospecialize(owner), args...) = new(owner, WorldRange(args...))
 end
-WorldView(cache, args...) = WorldView(cache, WorldRange(args...))
-WorldView(wvc::WorldView, wr::WorldRange) = WorldView(wvc.cache, wr)
-WorldView(wvc::WorldView, args...) = WorldView(wvc.cache, args...)
 
-function haskey(wvc::WorldView{InternalCodeCache}, mi::MethodInstance)
-    return ccall(:jl_rettype_inferred, Any, (Any, Any, UInt, UInt), wvc.cache.owner, mi, first(wvc.worlds), last(wvc.worlds)) !== nothing
+function setindex!(cache::InternalCodeCache, ci::CodeInstance, mi::MethodInstance)
+    @assert ci.owner === cache.owner
+    m = mi.def
+    if isa(m, Method)
+        ccall(:jl_push_newly_inferred, Cvoid, (Any,), ci)
+    end
+    ccall(:jl_mi_cache_insert, Cvoid, (Any, Any), mi, ci)
+    return cache
 end
 
-function get(wvc::WorldView{InternalCodeCache}, mi::MethodInstance, default)
-    r = ccall(:jl_rettype_inferred, Any, (Any, Any, UInt, UInt), wvc.cache.owner, mi, first(wvc.worlds), last(wvc.worlds))
+function haskey(wvc::InternalCodeCache, mi::MethodInstance)
+    return ccall(:jl_rettype_inferred, Any, (Any, Any, UInt, UInt), wvc.owner, mi, first(wvc.worlds), last(wvc.worlds)) !== nothing
+end
+
+function get(wvc::InternalCodeCache, mi::MethodInstance, default)
+    r = ccall(:jl_rettype_inferred, Any, (Any, Any, UInt, UInt), wvc.owner, mi, first(wvc.worlds), last(wvc.worlds))
     if r === nothing
         return default
     end
     return r::CodeInstance
 end
 
-function getindex(wvc::WorldView{InternalCodeCache}, mi::MethodInstance)
+function getindex(wvc::InternalCodeCache, mi::MethodInstance)
     r = get(wvc, mi, nothing)
     r === nothing && throw(KeyError(mi))
     return r::CodeInstance
 end
-
-function setindex!(wvc::WorldView{InternalCodeCache}, ci::CodeInstance, mi::MethodInstance)
-    setindex!(wvc.cache, ci, mi)
-    return wvc
-end

Compiler/src/cicache.jl

@@ -89,7 +89,7 @@ function is_argtype_match(𝕃::AbstractLattice,
     end
 end
 
-function va_process_argtypes(𝕃::AbstractLattice, given_argtypes::Vector{Any}, nargs::UInt, isva::Bool)
+function va_process_argtypes(𝕃::AbstractLattice, given_argtypes::Vector{Any}, nargs::UInt, isva::Bool, mi::MethodInstance)
     nargs = Int(nargs)
     if isva || (!isempty(given_argtypes) && isvarargtype(given_argtypes[end]))
         isva_given_argtypes = Vector{Any}(undef, nargs)
@@ -120,7 +120,10 @@ function va_process_argtypes(𝕃::AbstractLattice, given_argtypes::Vector{Any},
         end
         return isva_given_argtypes
     end
-    @assert length(given_argtypes) == nargs "invalid `given_argtypes` for `mi`"
+    if length(given_argtypes) != nargs
+        println(given_argtypes, " != ", nargs, " for ", mi)
+        throw(AssertionError("invalid `given_argtypes` for `mi`"))
+    end
     return given_argtypes
 end
 
@@ -178,16 +181,17 @@ function elim_free_typevars(@nospecialize t)
     end
 end
 
-function cache_lookup(𝕃::AbstractLattice, mi::MethodInstance, given_argtypes::Vector{Any},
-                      cache::Vector{InferenceResult})
+function constprop_cache_lookup(𝕃::AbstractLattice, mi::MethodInstance, given_argtypes::Vector{Any}, cache::Vector{InferenceResult})
     method = mi.def::Method
     nargtypes = length(given_argtypes)
     for cached_result in cache
         cached_result.tombstone && continue # ignore deleted entries (due to LimitedAccuracy)
         cached_result.linfo === mi || continue
         cache_argtypes = cached_result.argtypes
         @assert length(cache_argtypes) == nargtypes "invalid `cache_argtypes` for `mi`"
-        cache_overridden_by_const = cached_result.overridden_by_const::BitVector
+        cache_overridden_by_const = cached_result.overridden_by_const
+        cache_overridden_by_const === nothing && continue
+        cache_overridden_by_const = cache_overridden_by_const::BitVector
         for i in 1:nargtypes
             if !is_argtype_match(𝕃, given_argtypes[i], cache_argtypes[i], cache_overridden_by_const[i])
                 @goto next_cache

Compiler/src/inferenceresult.jl

@@ -217,7 +217,6 @@ end
 const CACHE_MODE_NULL     = 0x00      # not cached, optimization optional
 const CACHE_MODE_GLOBAL   = 0x01 << 0 # cached globally, optimization required
 const CACHE_MODE_LOCAL    = 0x01 << 1 # cached locally, optimization required
-const CACHE_MODE_VOLATILE = 0x01 << 2 # not cached, optimization required
 
 abstract type Handler end
 get_enter_idx(handler::Handler) = get_enter_idx_impl(handler)::Int
@@ -262,7 +261,7 @@ intersect(world::WorldWithRange, valid_worlds::WorldRange) =
 mutable struct InferenceState
     #= information about this method instance =#
     linfo::MethodInstance
-    world::WorldWithRange
+    valid_worlds::WorldRange
     mod::Module
     sptypes::Vector{VarState}
     slottypes::Vector{Any}
@@ -349,7 +348,7 @@ mutable struct InferenceState
         bb_vartable1 = bb_vartables[1] = VarTable(undef, nslots)
         argtypes = result.argtypes
 
-        argtypes = va_process_argtypes(typeinf_lattice(interp), argtypes, src.nargs, src.isva)
+        argtypes = va_process_argtypes(typeinf_lattice(interp), argtypes, src.nargs, src.isva, mi)
 
         nargtypes = length(argtypes)
         for i = 1:nslots
@@ -392,7 +391,7 @@ mutable struct InferenceState
         parentid = frameid = cycleid = 0
 
         this = new(
-            mi, WorldWithRange(world, valid_worlds), mod, sptypes, slottypes, src, cfg, spec_info,
+            mi, valid_worlds, mod, sptypes, slottypes, src, cfg, spec_info,
             currbb, currpc, ip, handler_info, ssavalue_uses, bb_vartables, bb_saw_latestworld, ssavaluetypes, ssaflags, edges, stmt_info,
             tasks, pclimitations, limitations, cycle_backedges, callstack, parentid, frameid, cycleid,
             result, unreachable, bestguess, exc_bestguess, ipo_effects,
@@ -401,9 +400,6 @@ mutable struct InferenceState
             interp)
 
         # some more setups
-        if !iszero(cache_mode & CACHE_MODE_LOCAL)
-            push!(get_inference_cache(interp), result)
-        end
         if !iszero(cache_mode & CACHE_MODE_GLOBAL)
             push!(callstack, this)
             this.cycleid = this.frameid = length(callstack)
@@ -412,7 +408,7 @@ mutable struct InferenceState
         # Apply generated function restrictions
         if src.min_world != 1 || src.max_world != typemax(UInt)
             # From generated functions
-            update_valid_age!(this, WorldRange(src.min_world, src.max_world))
+            update_valid_age!(this, world, WorldRange(src.min_world, src.max_world))
         end
 
         return this
@@ -615,8 +611,6 @@ function convert_cache_mode(cache_mode::Symbol)
         return CACHE_MODE_GLOBAL
     elseif cache_mode === :local
         return CACHE_MODE_LOCAL
-    elseif cache_mode === :volatile
-        return CACHE_MODE_VOLATILE
     elseif cache_mode === :no
         return CACHE_MODE_NULL
     end
@@ -821,7 +815,7 @@ mutable struct IRInterpretationState
     const spec_info::SpecInfo
     const ir::IRCode
     const mi::MethodInstance
-    world::WorldWithRange
+    valid_worlds::WorldRange
     curridx::Int
     time_caches::Float64
     time_paused::UInt64
@@ -836,9 +830,10 @@ mutable struct IRInterpretationState
     frameid::Int
     parentid::Int
 
-    function IRInterpretationState(interp::AbstractInterpreter,
-        spec_info::SpecInfo, ir::IRCode, mi::MethodInstance, argtypes::Vector{Any},
-        world::UInt, min_world::UInt, max_world::UInt)
+    function IRInterpretationState(
+            interp::AbstractInterpreter, spec_info::SpecInfo, ir::IRCode,
+            mi::MethodInstance, argtypes::Vector{Any}, min_world::UInt, max_world::UInt
+        )
         curridx = 1
         given_argtypes = Vector{Any}(undef, length(argtypes))
         for i = 1:length(given_argtypes)
@@ -856,28 +851,32 @@ mutable struct IRInterpretationState
         ssa_refined = BitSet()
         lazyreachability = LazyCFGReachability(ir)
         valid_worlds = WorldRange(min_world, max_world == typemax(UInt) ? get_world_counter() : max_world)
+        if !(get_inference_world(interp) in valid_worlds)
+            error("invalid age range update")
+        end
         tasks = WorkThunk[]
         edges = Any[]
         callstack = AbsIntState[]
-        return new(spec_info, ir, mi, WorldWithRange(world, valid_worlds),
+        return new(spec_info, ir, mi, valid_worlds,
                 curridx, 0.0, 0, argtypes_refined, ir.sptypes, tpdum,
                 ssa_refined, lazyreachability, tasks, edges, callstack, 0, 0)
     end
 end
 
-function IRInterpretationState(interp::AbstractInterpreter,
-    codeinst::CodeInstance, mi::MethodInstance, argtypes::Vector{Any}, world::UInt)
+function IRInterpretationState(
+        interp::AbstractInterpreter, codeinst::CodeInstance, mi::MethodInstance,
+        argtypes::Vector{Any}, @nospecialize(src)
+    )
     @assert get_ci_mi(codeinst) === mi "method instance is not synced with code instance"
-    src = @atomic :monotonic codeinst.inferred
     if isa(src, String)
         src = _uncompressed_ir(codeinst, src)
     else
         isa(src, CodeInfo) || return nothing
     end
     spec_info = SpecInfo(src)
     ir = inflate_ir(src, mi)
-    argtypes = va_process_argtypes(optimizer_lattice(interp), argtypes, src.nargs, src.isva)
-    return IRInterpretationState(interp, spec_info, ir, mi, argtypes, world,
+    argtypes = va_process_argtypes(optimizer_lattice(interp), argtypes, src.nargs, src.isva, mi)
+    return IRInterpretationState(interp, spec_info, ir, mi, argtypes,
                                  codeinst.min_world, codeinst.max_world)
 end
 
@@ -900,7 +899,7 @@ function print_callstack(frame::AbsIntState)
         end
         print("] ")
         print(frame_instance(sv))
-        is_cached(sv) || print("  [uncached]")
+        is_cached(sv) || print("  [not globally cached]")
         sv.parentid == idx - 1 || print(" [parent=", sv.parentid, "]")
         isempty(callers_in_cycle(sv)) || print(" [cycle=", sv.cycleid, "]")
         println()
@@ -964,9 +963,6 @@ spec_info(sv::IRInterpretationState) = sv.spec_info
 propagate_inbounds(sv::AbsIntState) = spec_info(sv).propagate_inbounds
 method_for_inference_limit_heuristics(sv::AbsIntState) = spec_info(sv).method_for_inference_limit_heuristics
 
-frame_world(sv::InferenceState) = sv.world.this
-frame_world(sv::IRInterpretationState) = sv.world.this
-
 function is_effect_overridden(sv::AbsIntState, effect::Symbol)
     if is_effect_overridden(frame_instance(sv), effect)
         return true
@@ -986,9 +982,13 @@ has_conditional(𝕃::AbstractLattice, ::InferenceState) = has_conditional(𝕃)
 has_conditional(::AbstractLattice, ::IRInterpretationState) = false
 
 # work towards converging the valid age range for sv
-function update_valid_age!(sv::AbsIntState, valid_worlds::WorldRange)
-    sv.world = intersect(sv.world, valid_worlds)
-    return sv.world.valid_worlds
+function update_valid_age!(sv::AbsIntState, world, valid_worlds::WorldRange)
+    valid_worlds = intersect(sv.valid_worlds, valid_worlds)
+    if !(world in valid_worlds)
+        error("invalid age range update")
+    end
+    sv.valid_worlds = valid_worlds
+    return valid_worlds
 end
 
 """

Compiler/src/inferencestate.jl

@@ -99,14 +99,9 @@ end
 
 const TOP_TUPLE = GlobalRef(Core, :tuple)
 
-# This corresponds to the type of `CodeInfo`'s `inlining_cost` field
-const InlineCostType = UInt16
-const MAX_INLINE_COST = typemax(InlineCostType)
-const MIN_INLINE_COST = InlineCostType(10)
-const MaybeCompressed = Union{CodeInfo, String}
-
-is_inlineable(@nospecialize src::MaybeCompressed) =
-    ccall(:jl_ir_inlining_cost, InlineCostType, (Any,), src) != MAX_INLINE_COST
+inlining_cost(@nospecialize src) =
+    src isa Union{MaybeCompressed,UInt8} ? ccall(:jl_ir_inlining_cost, InlineCostType, (Any,), src) : MAX_INLINE_COST
+is_inlineable(@nospecialize src) = inlining_cost(src) != MAX_INLINE_COST
 set_inlineable!(src::CodeInfo, val::Bool) =
     src.inlining_cost = (val ? MIN_INLINE_COST : MAX_INLINE_COST)
 
@@ -158,47 +153,38 @@ end
 
 struct InliningState{Interp<:AbstractInterpreter}
     edges::Vector{Any}
-    world::UInt
     interp::Interp
     opt_cache::IdDict{MethodInstance,CodeInstance}
 end
 function InliningState(sv::InferenceState, interp::AbstractInterpreter,
                        opt_cache::IdDict{MethodInstance,CodeInstance}=IdDict{MethodInstance,CodeInstance}())
-    return InliningState(sv.edges, frame_world(sv), interp, opt_cache)
+    return InliningState(sv.edges, interp, opt_cache)
 end
 function InliningState(interp::AbstractInterpreter,
                        opt_cache::IdDict{MethodInstance,CodeInstance}=IdDict{MethodInstance,CodeInstance}())
-    return InliningState(Any[], get_inference_world(interp), interp, opt_cache)
+    return InliningState(Any[], interp, opt_cache)
 end
 
 struct OptimizerCache{CodeCache}
-    wvc::WorldView{CodeCache}
-    owner
+    cache::CodeCache
     opt_cache::IdDict{MethodInstance,CodeInstance}
     function OptimizerCache(
-        wvc::WorldView{CodeCache},
-        @nospecialize(owner),
+        cache::CodeCache,
         opt_cache::IdDict{MethodInstance,CodeInstance}) where CodeCache
-        new{CodeCache}(wvc, owner, opt_cache)
+        return new{CodeCache}(cache, opt_cache)
     end
 end
-function get((; wvc, owner, opt_cache)::OptimizerCache, mi::MethodInstance, default)
+function get((; cache, 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
+        return opt_cache[mi] # this is incomplete right now, but will be finished (by finish_cycle) before caching anything
     end
-    return get(wvc, mi, default)
+    return get(cache, mi, default)
 end
 
 # get `code_cache(::AbstractInterpreter)` from `state::InliningState`
 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)
+    cache = code_cache(state.interp)
+    return OptimizerCache(cache, state.opt_cache)
 end
 
 mutable struct OptimizationResult
@@ -678,15 +664,15 @@ GetNativeEscapeCache(interp::AbstractInterpreter) = GetNativeEscapeCache(code_ca
 function ((; code_cache)::GetNativeEscapeCache)(codeinst::Union{CodeInstance,MethodInstance})
     if codeinst isa MethodInstance
         codeinst = get(code_cache, codeinst, nothing)
-        codeinst isa CodeInstance || return false
+        codeinst === nothing && return false
     end
     argescapes = traverse_analysis_results(codeinst) do @nospecialize result
         return result isa EscapeAnalysis.ArgEscapeCache ? result : nothing
     end
     if argescapes !== nothing
         return argescapes
     end
-    effects = decode_effects(codeinst.ipo_purity_bits)
+    effects = codeinst isa CodeInstance ? decode_effects(codeinst.ipo_purity_bits) : codeinst.ipo_effects
     if is_effect_free(effects) && is_inaccessiblememonly(effects)
         # We might not have run EA on simple frames without any escapes (e.g. when optimization
         # is skipped when result is constant-folded by abstract interpretation). If those