Merge pull request #39792 from JuliaLang/kf/inlinemtable

Remove inlining's dependence on method table lookups

base/compiler/abstractinterpretation.jl

@@ -108,7 +108,7 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
         val = pure_eval_call(f, argtypes)
         if val !== false
             # TODO: add some sort of edge(s)
-            return CallMeta(val, MethodResultPure())
+            return CallMeta(val, MethodResultPure(info))
         end
     end
 
@@ -875,8 +875,10 @@ function abstract_apply(interp::AbstractInterpreter, @nospecialize(itft), @nospe
         push!(retinfos, ApplyCallInfo(call.info, arginfo))
         res = tmerge(res, call.rt)
         if bail_out_apply(interp, res, sv)
-            # No point carrying forward the info, we're not gonna inline it anyway
-            retinfo = nothing
+            if i != length(ctypes)
+                # No point carrying forward the info, we're not gonna inline it anyway
+                retinfo = false
+            end
             break
         end
     end
@@ -1074,6 +1076,26 @@ function abstract_call_unionall(argtypes::Vector{Any})
     return Any
 end
 
+function abstract_invoke(interp::AbstractInterpreter, @nospecialize(ft), @nospecialize(types), @nospecialize(argtype), sv::InferenceState)
+    nargtype = typeintersect(types, argtype)
+    nargtype === Bottom && return CallMeta(Bottom, false)
+    nargtype isa DataType || return CallMeta(Any, false) # other cases are not implemented below
+    isdispatchelem(ft) || return CallMeta(Any, false) # check that we might not have a subtype of `ft` at runtime, before doing supertype lookup below
+    types = rewrap_unionall(Tuple{ft, unwrap_unionall(types).parameters...}, types)
+    nargtype = Tuple{ft, nargtype.parameters...}
+    argtype = Tuple{ft, argtype.parameters...}
+    result = findsup(types, method_table(interp))
+    if result === nothing
+        return CallMeta(Any, false)
+    end
+    method, valid_worlds = result
+    update_valid_age!(sv, valid_worlds)
+    (ti, env) = ccall(:jl_type_intersection_with_env, Any, (Any, Any), nargtype, method.sig)::SimpleVector
+    rt, edge = typeinf_edge(interp, method, ti, env, sv)
+    edge !== nothing && add_backedge!(edge::MethodInstance, sv)
+    return CallMeta(rt, InvokeCallInfo(MethodMatch(ti, env, method, argtype <: method.sig)))
+end
+
 # call where the function is known exactly
 function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
         fargs::Union{Nothing,Vector{Any}}, argtypes::Vector{Any},
@@ -1088,8 +1110,16 @@ function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
             ft = argtype_by_index(argtypes, 3)
             (itft === Bottom || ft === Bottom) && return CallMeta(Bottom, false)
             return abstract_apply(interp, itft, ft, argtype_tail(argtypes, 4), sv, max_methods)
+        elseif f === invoke
+            ft = widenconst(argtype_by_index(argtypes, 2))
+            (sigty, isexact, isconcrete, istype) = instanceof_tfunc(argtype_by_index(argtypes, 3))
+            (ft === Bottom || sigty === Bottom) && return CallMeta(Bottom, false)
+            if isexact
+                return abstract_invoke(interp, ft, sigty, argtypes_to_type(argtype_tail(argtypes, 4)), sv)
+            end
+            return CallMeta(Any, false)
         end
-        return CallMeta(abstract_call_builtin(interp, f, fargs, argtypes, sv, max_methods), nothing)
+        return CallMeta(abstract_call_builtin(interp, f, fargs, argtypes, sv, max_methods), false)
     elseif f === Core.kwfunc
         if la == 2
             ft = widenconst(argtypes[2])
@@ -1111,7 +1141,7 @@ function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
         elseif la == 3
             ub_var = argtypes[3]
         end
-        return CallMeta(typevar_tfunc(n, lb_var, ub_var), nothing)
+        return CallMeta(typevar_tfunc(n, lb_var, ub_var), false)
     elseif f === UnionAll
         return CallMeta(abstract_call_unionall(argtypes), false)
     elseif f === Tuple && la == 2 && !isconcretetype(widenconst(argtypes[2]))
@@ -1129,11 +1159,11 @@ function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
         # mark !== as exactly a negated call to ===
         rty = abstract_call_known(interp, (===), fargs, argtypes, sv).rt
         if isa(rty, Conditional)
-            return CallMeta(Conditional(rty.var, rty.elsetype, rty.vtype), nothing) # swap if-else
+            return CallMeta(Conditional(rty.var, rty.elsetype, rty.vtype), false) # swap if-else
         elseif isa(rty, Const)
-            return CallMeta(Const(rty.val === false), nothing)
+            return CallMeta(Const(rty.val === false), MethodResultPure())
         end
-        return CallMeta(rty, nothing)
+        return CallMeta(rty, false)
     elseif la == 3 && istopfunction(f, :(>:))
         # mark issupertype as a exact alias for issubtype
         # swap T1 and T2 arguments and call <:

base/compiler/optimize.jl

@@ -26,11 +26,10 @@ struct InferenceCaches{T, S}
     mi_cache::S
 end
 
-struct InliningState{S <: Union{EdgeTracker, Nothing}, T <: Union{InferenceCaches, Nothing}, V <: Union{Nothing, MethodTableView}}
+struct InliningState{S <: Union{EdgeTracker, Nothing}, T <: Union{InferenceCaches, Nothing}}
     params::OptimizationParams
     et::S
     caches::T
-    method_table::V
 end
 
 mutable struct OptimizationState
@@ -49,8 +48,7 @@ mutable struct OptimizationState
             EdgeTracker(s_edges, frame.valid_worlds),
             InferenceCaches(
                 get_inference_cache(interp),
-                WorldView(code_cache(interp), frame.world)),
-            method_table(interp))
+                WorldView(code_cache(interp), frame.world)))
         return new(frame.linfo,
                    frame.src, frame.stmt_info, frame.mod, frame.nargs,
                    frame.sptypes, frame.slottypes, false,
@@ -85,8 +83,7 @@ mutable struct OptimizationState
             nothing,
             InferenceCaches(
                 get_inference_cache(interp),
-                WorldView(code_cache(interp), get_world_counter())),
-            method_table(interp))
+                WorldView(code_cache(interp), get_world_counter())))
         return new(linfo,
                    src, stmt_info, inmodule, nargs,
                    sptypes_from_meth_instance(linfo), slottypes, false,

base/compiler/ssair/inlining.jl

@@ -950,14 +950,14 @@ function inline_apply!(ir::IRCode, todo::Vector{Pair{Int, Any}}, idx::Int, sig::
         if isa(info, UnionSplitApplyCallInfo)
             if length(info.infos) != 1
                 # TODO: Handle union split applies?
-                new_info = info = nothing
+                new_info = info = false
             else
                 info = info.infos[1]
                 new_info = info.call
             end
         else
             @assert info === nothing || info === false
-            new_info = info = nothing
+            new_info = info = false
         end
         arg_start = 3
         atypes = sig.atypes
@@ -1016,20 +1016,23 @@ is_builtin(s::Signature) =
     isa(s.f, Builtin) ||
     s.ft ⊑ Builtin
 
-function inline_invoke!(ir::IRCode, idx::Int, sig::Signature, invoke_data::InvokeData, state::InliningState, todo::Vector{Pair{Int, Any}})
+function inline_invoke!(ir::IRCode, idx::Int, sig::Signature, info::InvokeCallInfo,
+        state::InliningState, todo::Vector{Pair{Int, Any}})
     stmt = ir.stmts[idx][:inst]
     calltype = ir.stmts[idx][:type]
-    method = invoke_data.entry
-    (metharg, methsp) = ccall(:jl_type_intersection_with_env, Any, (Any, Any),
-            sig.atype, method.sig)::SimpleVector
-    methsp = methsp::SimpleVector
-    match = MethodMatch(metharg, methsp, method, true)
-    et = state.et
-    result = analyze_method!(match, sig.atypes, et, state.caches, state.params, calltype)
-    handle_single_case!(ir, stmt, idx, result, true, todo)
-    if et !== nothing
-        intersect!(et, WorldRange(invoke_data.min_valid, invoke_data.max_valid))
+
+    if !info.match.fully_covers
+        # XXX: We could union split this
+        return nothing
     end
+
+    atypes = sig.atypes
+    atype0 = atypes[2]
+    atypes = atypes[4:end]
+    pushfirst!(atypes, atype0)
+
+    result = analyze_method!(info.match, atypes, state.et, state.caches, state.params, calltype)
+    handle_single_case!(ir, stmt, idx, result, true, todo)
     return nothing
 end
 
@@ -1061,43 +1064,18 @@ function process_simple!(ir::IRCode, todo::Vector{Pair{Int, Any}}, idx::Int, sta
         return nothing
     end
 
-    # Handle invoke
-    invoke_data = nothing
-    if sig.f === Core.invoke && length(sig.atypes) >= 3
-        res = compute_invoke_data(sig.atypes, state.method_table)
-        res === nothing && return nothing
-        (sig, invoke_data) = res
-    elseif is_builtin(sig)
-        # No inlining for builtins (other than what was previously handled)
+    if sig.f !== Core.invoke && is_builtin(sig)
+        # No inlining for builtins (other invoke/apply)
         return nothing
     end
 
     sig = with_atype(sig)
 
-    # In :invoke, make sure that the arguments we're passing are a subtype of the
-    # signature we're invoking.
-    (invoke_data === nothing || sig.atype <: invoke_data.types0) || return nothing
-
     # Special case inliners for regular functions
     if late_inline_special_case!(ir, sig, idx, stmt, state.params) || is_return_type(sig.f)
         return nothing
     end
-    return (sig, invoke_data)
-end
-
-# This is not currently called in the regular course, but may be needed
-# if we ever want to re-run inlining again later in the pass pipeline after
-# additional type information was discovered.
-function recompute_method_matches(@nospecialize(atype), params::OptimizationParams, et::EdgeTracker, method_table::MethodTableView)
-    # Regular case: Retrieve matching methods from cache (or compute them)
-    # World age does not need to be taken into account in the cache
-    # because it is forwarded from type inference through `sv.params`
-    # in the case that the cache is nonempty, so it should be unchanged
-    # The max number of methods should be the same as in inference most
-    # of the time, and should not affect correctness otherwise.
-    results = findall(atype, method_table; limit=params.MAX_METHODS)
-    results !== missing && intersect!(et, results.valid_worlds)
-    MethodMatchInfo(results)
+    return sig
 end
 
 function analyze_single_call!(ir::IRCode, todo::Vector{Pair{Int, Any}}, idx::Int, @nospecialize(stmt),
@@ -1208,63 +1186,52 @@ function assemble_inline_todo!(ir::IRCode, state::InliningState)
     # todo = (inline_idx, (isva, isinvoke, na), method, spvals, inline_linetable, inline_ir, lie)
     todo = Pair{Int, Any}[]
     et = state.et
-    method_table = state.method_table
     for idx in 1:length(ir.stmts)
-        r = process_simple!(ir, todo, idx, state)
-        r === nothing && continue
+        sig = process_simple!(ir, todo, idx, state)
+        sig === nothing && continue
 
         stmt = ir.stmts[idx][:inst]
         calltype = ir.stmts[idx][:type]
         info = ir.stmts[idx][:info]
-        # Inference determined this couldn't be analyzed. Don't question it.
-        if info === false
-            continue
-        end
-
-        (sig, invoke_data) = r
 
         # Check whether this call was @pure and evaluates to a constant
-        if calltype isa Const && info isa MethodResultPure
-            if is_inlineable_constant(calltype.val)
+        if info isa MethodResultPure
+            if calltype isa Const && is_inlineable_constant(calltype.val)
                 ir.stmts[idx][:inst] = quoted(calltype.val)
                 continue
             end
+            info = info.info
+        end
+
+        # Inference determined this couldn't be analyzed. Don't question it.
+        if info === false
+            continue
         end
 
         # If inference arrived at this result by using constant propagation,
-        # it'll performed a specialized analysis for just this case. Use its
+        # it'll have performed a specialized analysis for just this case. Use its
         # result.
         if isa(info, ConstCallInfo)
-            handle_const_call!(ir, idx, stmt, info, sig, calltype, state.et,
-                state.caches, invoke_data !== nothing, todo)
+            handle_const_call!(ir, idx, stmt, info, sig, calltype, state.et, state.caches,
+                sig.f === Core.invoke, todo)
             continue
         end
 
-        # Ok, now figure out what method to call
-        if invoke_data !== nothing
-            inline_invoke!(ir, idx, sig, invoke_data, state, todo)
+        # Handle invoke
+        if sig.f === Core.invoke
+            if isa(info, InvokeCallInfo)
+                inline_invoke!(ir, idx, sig, info, state, todo)
+            end
             continue
         end
 
-        nu = unionsplitcost(sig.atypes)
-        if nu == 1 || nu > state.params.MAX_UNION_SPLITTING
-            if !isa(info, MethodMatchInfo)
-                method_table === nothing && continue
-                et === nothing && continue
-                info = recompute_method_matches(sig.atype, state.params, et, method_table)
-            end
+        # Ok, now figure out what method to call
+        if isa(info, MethodMatchInfo)
             infos = MethodMatchInfo[info]
+        elseif isa(info, UnionSplitInfo)
+            infos = info.matches
         else
-            if !isa(info, UnionSplitInfo)
-                method_table === nothing && continue
-                et === nothing && continue
-                infos = MethodMatchInfo[]
-                for union_sig in UnionSplitSignature(sig.atypes)
-                    push!(infos, recompute_method_matches(argtypes_to_type(union_sig), state.params, et, method_table))
-                end
-            else
-                infos = info.matches
-            end
+            continue
         end
 
         analyze_single_call!(ir, todo, idx, stmt, sig, calltype, infos, state.et, state.caches, state.params)
@@ -1286,38 +1253,6 @@ function linear_inline_eligible(ir::IRCode)
     return true
 end
 
-function compute_invoke_data(@nospecialize(atypes), method_table)
-    ft = widenconst(atypes[2])
-    if !isdispatchelem(ft) || has_free_typevars(ft) || (ft <: Builtin)
-        # TODO: this can be rather aggressive at preventing inlining of closures
-        # but we need to check that `ft` can't have a subtype at runtime before using the supertype lookup below
-        return nothing
-    end
-    invoke_tt = widenconst(atypes[3])
-    if !isType(invoke_tt) || has_free_typevars(invoke_tt)
-        return nothing
-    end
-    invoke_tt = invoke_tt.parameters[1]
-    if !(isa(unwrap_unionall(invoke_tt), DataType) && invoke_tt <: Tuple)
-        return nothing
-    end
-    if method_table === nothing
-        # TODO: These should be forwarded in stmt_info, just like regular
-        # method lookup results
-        return nothing
-    end
-    invoke_types = rewrap_unionall(Tuple{ft, unwrap_unionall(invoke_tt).parameters...}, invoke_tt)
-    invoke_entry = findsup(invoke_types, method_table)
-    invoke_entry === nothing && return nothing
-    method, valid_worlds = invoke_entry
-    invoke_data = InvokeData(method, invoke_types, first(valid_worlds), last(valid_worlds))
-    atype0 = atypes[2]
-    atypes = atypes[4:end]
-    pushfirst!(atypes, atype0)
-    f = singleton_type(ft)
-    return (Signature(f, ft, atypes), invoke_data)
-end
-
 # Check for a number of functions known to be pure
 function ispuretopfunction(@nospecialize(f))
     return istopfunction(f, :typejoin) ||
@@ -1386,6 +1321,11 @@ 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)
+        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])
+        return true
     elseif is_return_type(f)
         if isconstType(typ)
             ir[SSAValue(idx)] = quoted(typ.parameters[1])

base/compiler/stmtinfo.jl

@@ -15,11 +15,17 @@ end
 """
     struct MethodResultPure
 
-This singleton represents a method result constant was proven to be
+This struct represents a method result constant was proven to be
 effect-free, including being no-throw (typically because the value was computed
 by calling an `@pure` function).
 """
-struct MethodResultPure end
+struct MethodResultPure
+    info::Any
+end
+let instance = MethodResultPure(false)
+    global MethodResultPure
+    MethodResultPure() = instance
+end
 
 """
     struct UnionSplitInfo
@@ -94,6 +100,16 @@ struct ConstCallInfo
     result::InferenceResult
 end
 
+"""
+    struct InvokeCallInfo
+
+Represents a resolved call to `invoke`, carrying the Method match of the
+method being processed.
+"""
+struct InvokeCallInfo
+    match::MethodMatch
+end
+
 # Stmt infos that are used by external consumers, but not by optimization.
 # These are not produced by default and must be explicitly opted into by
 # the AbstractInterpreter.