re-infer source when it's not cached

Author: aviatesk

base/compiler/abstractinterpretation.jl

@@ -676,7 +676,7 @@ function const_prop_methodinstance_heuristic(interp::AbstractInterpreter, match:
         cache_inf = code.inferred
         if !(cache_inf === nothing)
             # TODO maybe we want to respect callsite `@inline`/`@noinline` annotations here ?
-            cache_inlineable = inlining_policy(interp)(cache_inf, nothing, match) !== nothing
+            cache_inlineable = inlining_policy(interp, cache_inf, 0x00, match) !== nothing
         end
     end
     if !cache_inlineable

base/compiler/optimize.jl

@@ -21,15 +21,14 @@ function push!(et::EdgeTracker, ci::CodeInstance)
     push!(et, ci.def)
 end
 
-struct InliningState{S <: Union{EdgeTracker, Nothing}, T, P, U}
+struct InliningState{S <: Union{EdgeTracker, Nothing}, T, I<:AbstractInterpreter}
     params::OptimizationParams
     et::S
     mi_cache::T
-    inf_cache::U
-    policy::P
+    interp::I
 end
 
-function default_inlining_policy(@nospecialize(src), stmt_flag::Union{Nothing,UInt8}, match::Union{MethodMatch,InferenceResult})
+function inlining_policy(interp::AbstractInterpreter, @nospecialize(src), stmt_flag::UInt8, match::Union{MethodMatch,InferenceResult})
     if isa(src, CodeInfo) || isa(src, Vector{UInt8})
         src_inferred = ccall(:jl_ir_flag_inferred, Bool, (Any,), src)
         src_inlineable = is_stmt_inline(stmt_flag) || ccall(:jl_ir_flag_inlineable, Bool, (Any,), src)
@@ -38,9 +37,12 @@ function default_inlining_policy(@nospecialize(src), stmt_flag::Union{Nothing,UI
         return (is_stmt_inline(stmt_flag) || src.src.inlineable) ? src.ir : nothing
     elseif src === nothing && is_stmt_inline(stmt_flag) && isa(match, MethodMatch)
         # when the source isn't available at this moment, try to re-infer and inline it
-        # HACK in order to avoid cycles here, we disable inlining and makes sure the following inference never comes here
-        # TODO sort out `AbstractInterpreter` interface to handle this well, and also inference should try to keep the source if the statement will be inlined
-        interp = NativeInterpreter(; opt_params = OptimizationParams(; inlining = false))
+        # NOTE we can make inference try to keep the source if the call is going to be inlined,
+        # but then inlining will depend on local state of inference and so the first entry
+        # and the succeeding ones may generate different code; rather we always re-infer
+        # the source to avoid the problem while it's obviously not most efficient
+        # HACK disable inlining for the re-inference to avoid cycles by making sure the following inference never comes here again
+        interp = NativeInterpreter(get_world_counter(interp); opt_params = OptimizationParams(; inlining = false))
         src, rt = typeinf_code(interp, match.method, match.spec_types, match.sparams, true)
         return src
     end
@@ -65,8 +67,7 @@ mutable struct OptimizationState
         inlining = InliningState(params,
             EdgeTracker(s_edges, frame.valid_worlds),
             WorldView(code_cache(interp), frame.world),
-            get_inference_cache(interp),
-            inlining_policy(interp))
+            interp)
         return new(frame.linfo,
                    frame.src, nothing, frame.stmt_info, frame.mod, frame.nargs,
                    frame.sptypes, frame.slottypes, false,
@@ -100,8 +101,7 @@ mutable struct OptimizationState
         inlining = InliningState(params,
             nothing,
             WorldView(code_cache(interp), get_world_counter()),
-            get_inference_cache(interp),
-            inlining_policy(interp))
+            interp)
         return new(linfo,
                    src, nothing, stmt_info, inmodule, nargs,
                    sptypes_from_meth_instance(linfo), slottypes, false,
@@ -192,10 +192,8 @@ function isinlineable(m::Method, me::OptimizationState, params::OptimizationPara
     return inlineable
 end
 
-is_stmt_inline(stmt_flag::UInt8)   = stmt_flag & IR_FLAG_INLINE != 0
-is_stmt_inline(::Nothing)          = false
+is_stmt_inline(stmt_flag::UInt8)   = stmt_flag & IR_FLAG_INLINE   != 0
 is_stmt_noinline(stmt_flag::UInt8) = stmt_flag & IR_FLAG_NOINLINE != 0
-is_stmt_noinline(::Nothing)        = false # not used for now
 
 # These affect control flow within the function (so may not be removed
 # if there is no usage within the function), but don't affect the purity

base/compiler/ssair/inlining.jl

@@ -722,7 +722,7 @@ function resolve_todo(todo::InliningTodo, state::InliningState, flag::UInt8)
     (; match) = todo.spec::DelayedInliningSpec
 
     #XXX: update_valid_age!(min_valid[1], max_valid[1], sv)
-    isconst, src, argtypes = false, nothing, nothing
+    isconst, src = false, nothing
     if isa(match, InferenceResult)
         let inferred_src = match.src
             if isa(inferred_src, Const)
@@ -734,9 +734,6 @@ function resolve_todo(todo::InliningTodo, state::InliningState, flag::UInt8)
                 isconst, src = false, inferred_src
             end
         end
-        if is_stmt_inline(flag)
-            argtypes = match.argtypes
-        end
     else
         linfo = get(state.mi_cache, todo.mi, nothing)
         if linfo isa CodeInstance
@@ -749,9 +746,6 @@ function resolve_todo(todo::InliningTodo, state::InliningState, flag::UInt8)
         else
             isconst, src = false, linfo
         end
-        if is_stmt_inline(flag)
-            argtypes = collect(todo.mi.specTypes.parameters)::Vector{Any}
-        end
     end
 
     et = state.et
@@ -761,18 +755,7 @@ function resolve_todo(todo::InliningTodo, state::InliningState, flag::UInt8)
         return ConstantCase(src)
     end
 
-    if argtypes !== nothing && src === nothing
-        inf_cache = state.inf_cache
-        inf_result = cache_lookup(todo.mi, argtypes, inf_cache)
-        if isa(inf_result, InferenceResult)
-            src = inf_result.src
-            if isa(src, OptimizationState)
-                src = src.src
-            end
-        end
-    end
-
-    src = state.policy(src, flag, match)
+    src = inlining_policy(state.interp, src, flag, match)
 
     if src === nothing
         return compileable_specialization(et, match)
@@ -1405,7 +1388,8 @@ end
 function late_inline_special_case!(ir::IRCode, sig::Signature, idx::Int, stmt::Expr, params::OptimizationParams)
     f, ft, atypes = sig.f, sig.ft, sig.atypes
     typ = ir.stmts[idx][:type]
-    if params.inlining && length(atypes) == 3 && istopfunction(f, :!==)
+    isinlining = params.inlining
+    if isinlining && length(atypes) == 3 && istopfunction(f, :!==)
         # special-case inliner for !== that precedes _methods_by_ftype union splitting
         # and that works, even though inference generally avoids inferring the `!==` Method
         if isa(typ, Const)
@@ -1417,7 +1401,7 @@ function late_inline_special_case!(ir::IRCode, sig::Signature, idx::Int, stmt::E
         not_call = Expr(:call, GlobalRef(Core.Intrinsics, :not_int), cmp_call_ssa)
         ir[SSAValue(idx)] = not_call
         return true
-    elseif params.inlining && length(atypes) == 3 && istopfunction(f, :(>:))
+    elseif isinlining && length(atypes) == 3 && istopfunction(f, :(>:))
         # special-case inliner for issupertype
         # that works, even though inference generally avoids inferring the `>:` Method
         if isa(typ, Const) && _builtin_nothrow(<:, Any[atypes[3], atypes[2]], typ)
@@ -1427,7 +1411,7 @@ function late_inline_special_case!(ir::IRCode, sig::Signature, idx::Int, stmt::E
         subtype_call = Expr(:call, GlobalRef(Core, :(<:)), stmt.args[3], stmt.args[2])
         ir[SSAValue(idx)] = subtype_call
         return true
-    elseif params.inlining && f === TypeVar && 2 <= length(atypes) <= 4 && (atypes[2] ⊑ Symbol)
+    elseif isinlining && f === TypeVar && 2 <= length(atypes) <= 4 && (atypes[2] ⊑ Symbol)
         ir[SSAValue(idx)] = Expr(:call, GlobalRef(Core, :_typevar), stmt.args[2],
             length(stmt.args) < 4 ? Bottom : stmt.args[3],
             length(stmt.args) == 2 ? Any : stmt.args[end])

base/compiler/types.jl

@@ -213,7 +213,6 @@ may_discard_trees(ni::NativeInterpreter) = true
 verbose_stmt_info(ni::NativeInterpreter) = false
 
 method_table(ai::AbstractInterpreter) = InternalMethodTable(get_world_counter(ai))
-inlining_policy(ai::AbstractInterpreter) = default_inlining_policy
 
 # define inference bail out logic
 # `NativeInterpreter` bails out from inference when

test/compiler/inline.jl

@@ -539,8 +539,19 @@ end
             end
         end
 
+        # test inlining of un-cached callsites
+
         import Core.Compiler: isType
+
         limited(a) = @noinline(isType(a)) ? @inline(limited(a.parameters[1])) : rand(a)
+
+        function multilimited(a)
+            if @noinline(isType(a))
+                return @inline(multilimited(a.parameters[1]))
+            else
+                return rand(Bool) ? rand(a) : @inline(multilimited(a))
+            end
+        end
     end
 
     let ci = code_typed1(m.force_inline_explicit, (Int,))
@@ -593,6 +604,11 @@ end
     let ci = code_typed1(m.limited, (Any,))
         @test count(x->isinvoke(x, :isType), ci.code) == 2
     end
+    # check that inlining for recursive callsites doesn't depend on inference local cache
+    let ci1 = code_typed1(m.multilimited, (Any,))
+        ci2 = code_typed1(m.multilimited, (Any,))
+        @test ci1.code == ci2.code
+    end
 end
 
 # Issue #41299 - inlining deletes error check in :>