Streamline bail path in inlining (#36771)

When inlining declines to inline something, it instead turns them into
:invoke statements. These are then turned into direct (non-inlined)
calls by codegen or otherwise receive a fast path at runtime. While
inlining has evolved quite a bit, this code has stayed much the same
since it was introduced four years ago and doesn't seem to make much
sense as is. In particular:

1. For the non-`invoke()` case we were doing an extra method look that
seems entirely superfluous, because we already had to do the very same
method lookup just to reach this point. The only thing this path was
doing at that point was creating a "compilable" specialization (which
might use a slightly different signature). We might as well do that
directly.

2. For the invoke case, we were pro-actively adding the specialization
to the `->invokes` dispatch cache. However, this doesn't make much sense
a priori either, because the bail path does not go through the runtime
`invoke()` code that uses that cache (it did many years ago when this
code was introduced, but hasn't in a long time). There does not seem
to be a good reason to believe that this signature will be any more
likely than any other to be invoked using the runtime mechanism.

This cleans up that path by getting rid of both the superfluous method
lookup and the superfluous addition to the `->invokes` cache. There
should be a slight performance improvement as well from avoiding
this superfluous work, but the bail path is less common than one
might expect (the vast majority of call sites are inlined) and in
measurements the effect seems to be in the noise. Nevertheless,
it seems like a nice simplification and is conceptually clearer.

Author: Keno

base/boot.jl

@@ -406,6 +406,8 @@ eval(Core, :(CodeInstance(mi::MethodInstance, @nospecialize(rettype), @nospecial
 eval(Core, :(Const(@nospecialize(v)) = $(Expr(:new, :Const, :v, false))))
 eval(Core, :(Const(@nospecialize(v), actual::Bool) = $(Expr(:new, :Const, :v, :actual))))
 eval(Core, :(PartialStruct(@nospecialize(typ), fields::Array{Any, 1}) = $(Expr(:new, :PartialStruct, :typ, :fields))))
+eval(Core, :(MethodMatch(@nospecialize(spec_types), sparams::SimpleVector, method::Method, fully_covers::Bool) =
+    $(Expr(:new, :MethodMatch, :spec_types, :sparams, :method, :fully_covers))))
 
 Module(name::Symbol=:anonymous, std_imports::Bool=true) = ccall(:jl_f_new_module, Ref{Module}, (Any, Bool), name, std_imports)
 

base/compiler/ssair/inlining.jl

@@ -578,23 +578,6 @@ function batch_inline!(todo::Vector{Any}, ir::IRCode, linetable::Vector{LineInfo
     return ir
 end
 
-function spec_lambda(@nospecialize(atype), sv::OptimizationState, @nospecialize(invoke_data))
-    min_valid = RefValue{UInt}(typemin(UInt))
-    max_valid = RefValue{UInt}(typemax(UInt))
-    if invoke_data === nothing
-        mi = ccall(:jl_get_spec_lambda, Any, (Any, UInt, Ptr{UInt}, Ptr{UInt}), atype, sv.world, min_valid, max_valid)
-    else
-        invoke_data = invoke_data::InvokeData
-        atype <: invoke_data.types0 || return nothing
-        mi = ccall(:jl_get_invoke_lambda, Any, (Any, Any), invoke_data.entry, atype)
-        min_valid[] = invoke_data.min_valid
-        max_valid[] = invoke_data.max_valid
-    end
-    mi !== nothing && add_backedge!(mi::MethodInstance, sv)
-    update_valid_age!(sv, WorldRange(min_valid[], max_valid[]))
-    return mi
-end
-
 # This assumes the caller has verified that all arguments to the _apply call are Tuples.
 function rewrite_apply_exprargs!(ir::IRCode, todo::Vector{Any}, idx::Int, argexprs::Vector{Any}, atypes::Vector{Any}, arginfos::Vector{Any}, arg_start::Int, sv::OptimizationState)
     new_argexprs = Any[argexprs[arg_start]]
@@ -688,16 +671,23 @@ function singleton_type(@nospecialize(ft))
     return nothing
 end
 
-function analyze_method!(idx::Int, sig::Signature, @nospecialize(metharg), methsp::SimpleVector,
-                         method::Method, stmt::Expr, sv::OptimizationState,
-                         isinvoke::Bool, invoke_data::Union{InvokeData,Nothing}, @nospecialize(stmttyp))
+function compileable_specialization(match::MethodMatch, sv::OptimizationState)
+    mi = specialize_method(match, false, true)
+    mi !== nothing && add_backedge!(mi::MethodInstance, sv)
+    return mi
+end
+
+function analyze_method!(idx::Int, sig::Signature, match::MethodMatch,
+                         stmt::Expr, sv::OptimizationState,
+                         isinvoke::Bool, @nospecialize(stmttyp))
     f, ft, atypes, atype_unlimited = sig.f, sig.ft, sig.atypes, sig.atype
+    method = match.method
     methsig = method.sig
 
     # Check whether this call just evaluates to a constant
     if isa(f, widenconst(ft)) &&
             isa(stmttyp, Const) && stmttyp.actual && is_inlineable_constant(stmttyp.val)
-        return ConstantCase(quoted(stmttyp.val), method, Any[methsp...], metharg)
+        return ConstantCase(quoted(stmttyp.val), method, Any[match.sparams...], match.spec_types)
     end
 
     # Check that we habe the correct number of arguments
@@ -713,34 +703,34 @@ function analyze_method!(idx::Int, sig::Signature, @nospecialize(metharg), meths
 
     # Bail out if any static parameters are left as TypeVar
     ok = true
-    for i = 1:length(methsp)
-        isa(methsp[i], TypeVar) && return nothing
+    for i = 1:length(match.sparams)
+        isa(match.sparams[i], TypeVar) && return nothing
     end
 
     # See if there exists a specialization for this method signature
-    mi = specialize_method(method, metharg, methsp, true) # Union{Nothing, MethodInstance}
+    mi = specialize_method(match, true) # Union{Nothing, MethodInstance}
     if !isa(mi, MethodInstance)
-        return spec_lambda(atype_unlimited, sv, invoke_data)
+        return compileable_specialization(match, sv)
     end
 
     isconst, src = find_inferred(mi, atypes, sv, stmttyp)
     if isconst
         if sv.params.inlining
             add_backedge!(mi, sv)
-            return ConstantCase(src, method, Any[methsp...], metharg)
+            return ConstantCase(src, method, Any[match.sparams...], match.spec_types)
         else
-            return spec_lambda(atype_unlimited, sv, invoke_data)
+            return compileable_specialization(match, sv)
         end
     end
     if src === nothing
-        return spec_lambda(atype_unlimited, sv, invoke_data)
+        return compileable_specialization(match, sv)
     end
 
     src_inferred = ccall(:jl_ir_flag_inferred, Bool, (Any,), src)
     src_inlineable = ccall(:jl_ir_flag_inlineable, Bool, (Any,), src)
 
     if !(src_inferred && src_inlineable && sv.params.inlining)
-        return spec_lambda(atype_unlimited, sv, invoke_data)
+        return compileable_specialization(match, sv)
     end
 
     # At this point we're committed to performing the inlining, add the backedge
@@ -767,7 +757,7 @@ function analyze_method!(idx::Int, sig::Signature, @nospecialize(metharg), meths
     return InliningTodo(idx,
         na > 0 && method.isva,
         isinvoke, na,
-        method, Any[methsp...], metharg,
+        method, Any[match.sparams...], match.spec_types,
         inline_linetable, ir2, linear_inline_eligible(ir2))
 end
 
@@ -985,7 +975,8 @@ function inline_invoke!(ir::IRCode, idx::Int, sig::Signature, invoke_data::Invok
     (metharg, methsp) = ccall(:jl_type_intersection_with_env, Any, (Any, Any),
             sig.atype, method.sig)::SimpleVector
     methsp = methsp::SimpleVector
-    result = analyze_method!(idx, sig, metharg, methsp, method, stmt, sv, true, invoke_data, calltype)
+    match = MethodMatch(metharg, methsp, method, true)
+    result = analyze_method!(idx, sig, match, stmt, sv, true, calltype)
     handle_single_case!(ir, stmt, idx, result, true, todo)
     update_valid_age!(sv, WorldRange(invoke_data.min_valid, invoke_data.max_valid))
     return nothing
@@ -1087,22 +1078,21 @@ function analyze_single_call!(ir::IRCode, todo::Vector{Any}, idx::Int, @nospecia
             only_method = false
         end
         for match in meth
-            (metharg, methsp, method) = (match.spec_types, match.sparams, match.method)
-            signature_union = Union{signature_union, metharg}
-            if !isdispatchtuple(metharg)
+            signature_union = Union{signature_union, match.spec_types}
+            if !isdispatchtuple(match.spec_types)
                 fully_covered = false
                 continue
             end
-            case_sig = Signature(sig.f, sig.ft, sig.atypes, metharg)
-            case = analyze_method!(idx, case_sig, metharg, methsp, method,
-                stmt, sv, false, nothing, calltype)
+            case_sig = Signature(sig.f, sig.ft, sig.atypes, match.spec_types)
+            case = analyze_method!(idx, case_sig, match,
+                stmt, sv, false, calltype)
             if case === nothing
                 fully_covered = false
                 continue
-            elseif _any(p->p[1] === metharg, cases)
+            elseif _any(p->p[1] === match.spec_types, cases)
                 continue
             end
-            push!(cases, Pair{Any,Any}(metharg, case))
+            push!(cases, Pair{Any,Any}(match.spec_types, case))
         end
     end
 
@@ -1113,18 +1103,17 @@ function analyze_single_call!(ir::IRCode, todo::Vector{Any}, idx::Int, @nospecia
     # we inline, even if the signature is not a dispatch tuple
     if signature_fully_covered && length(cases) == 0 && only_method isa Method
         if length(infos) > 1
-            method = only_method
             (metharg, methsp) = ccall(:jl_type_intersection_with_env, Any, (Any, Any),
-                sig.atype, method.sig)::SimpleVector
+                sig.atype, only_method.sig)::SimpleVector
+            match = MethodMatch(metharg, methsp, only_method, true)
         else
             @assert length(meth) == 1
-            (metharg, methsp, method) = (meth[1].spec_types, meth[1].sparams, meth[1].method)
+            match = meth[1]
         end
         fully_covered = true
-        case = analyze_method!(idx, sig, metharg, methsp, method,
-            stmt, sv, false, nothing, calltype)
+        case = analyze_method!(idx, sig, match, stmt, sv, false, calltype)
         case === nothing && return
-        push!(cases, Pair{Any,Any}(metharg, case))
+        push!(cases, Pair{Any,Any}(match.spec_types, case))
     end
     if !signature_fully_covered
         fully_covered = false

base/compiler/utilities.jl

@@ -118,8 +118,24 @@ function retrieve_code_info(linfo::MethodInstance)
     end
 end
 
+# Get at the nonfunction_mt, which happens to be the mt of SimpleVector
+const nonfunction_mt = typename(SimpleVector).mt
+
+function get_compileable_sig(method::Method, @nospecialize(atypes), sparams::SimpleVector)
+    isa(atypes, DataType) || return Nothing
+    mt = ccall(:jl_method_table_for, Any, (Any,), atypes)
+    mt === nothing && return nothing
+    return ccall(:jl_normalize_to_compilable_sig, Any, (Any, Any, Any, Any),
+        mt, atypes, sparams, method)
+end
+
 # get a handle to the unique specialization object representing a particular instantiation of a call
-function specialize_method(method::Method, @nospecialize(atypes), sparams::SimpleVector, preexisting::Bool=false)
+function specialize_method(method::Method, @nospecialize(atypes), sparams::SimpleVector, preexisting::Bool=false, compilesig::Bool=false)
+    if compilesig
+        new_atypes = get_compileable_sig(method, atypes, sparams)
+        new_atypes === nothing && return nothing
+        atypes = new_atypes
+    end
     if preexisting
         # check cached specializations
         # for an existing result stored there
@@ -128,8 +144,8 @@ function specialize_method(method::Method, @nospecialize(atypes), sparams::Simpl
     return ccall(:jl_specializations_get_linfo, Ref{MethodInstance}, (Any, Any, Any), method, atypes, sparams)
 end
 
-function specialize_method(match::MethodMatch, preexisting::Bool=false)
-    return specialize_method(match.method, match.spec_types, match.sparams, preexisting)
+function specialize_method(match::MethodMatch, preexisting::Bool=false, compilesig::Bool=false)
+    return specialize_method(match.method, match.spec_types, match.sparams, preexisting, compilesig)
 end
 
 # This function is used for computing alternate limit heuristics

src/gf.c

@@ -1945,6 +1945,20 @@ JL_DLLEXPORT int32_t jl_invoke_api(jl_code_instance_t *codeinst)
     return -1;
 }
 
+JL_DLLEXPORT jl_value_t *jl_normalize_to_compilable_sig(jl_methtable_t *mt, jl_tupletype_t *ti, jl_svec_t *env, jl_method_t *m)
+{
+    jl_tupletype_t *tt = NULL;
+    jl_svec_t *newparams = NULL;
+    JL_GC_PUSH2(&tt, &newparams);
+    intptr_t nspec = (mt == jl_type_type_mt || mt == jl_nonfunction_mt ? m->nargs + 1 : mt->max_args + 2);
+    jl_compilation_sig(ti, env, m, nspec, &newparams);
+    tt = (newparams ? jl_apply_tuple_type(newparams) : ti);
+    int is_compileable = ((jl_datatype_t*)ti)->isdispatchtuple ||
+        jl_isa_compileable_sig(tt, m);
+    JL_GC_POP();
+    return is_compileable ? (jl_value_t*)tt : jl_nothing;
+}
+
 // compile-time method lookup
 jl_method_instance_t *jl_get_specialization1(jl_tupletype_t *types JL_PROPAGATES_ROOT, size_t world, size_t *min_valid, size_t *max_valid, int mt_cache)
 {
@@ -1965,9 +1979,8 @@ jl_method_instance_t *jl_get_specialization1(jl_tupletype_t *types JL_PROPAGATES
         *max_valid = max_valid2;
     if (matches == jl_false || jl_array_len(matches) != 1 || ambig)
         return NULL;
-    jl_tupletype_t *tt = NULL;
-    jl_svec_t *newparams = NULL;
-    JL_GC_PUSH3(&matches, &tt, &newparams);
+    jl_value_t *tt = NULL;
+    JL_GC_PUSH2(&matches, &tt);
     jl_method_match_t *match = (jl_method_match_t*)jl_array_ptr_ref(matches, 0);
     jl_method_t *m = match->method;
     jl_svec_t *env = match->sparams;
@@ -1987,12 +2000,8 @@ jl_method_instance_t *jl_get_specialization1(jl_tupletype_t *types JL_PROPAGATES
                 JL_UNLOCK(&mt->writelock);
             }
             else {
-                intptr_t nspec = (mt == jl_type_type_mt || mt == jl_nonfunction_mt ? m->nargs + 1 : mt->max_args + 2);
-                jl_compilation_sig(ti, env, m, nspec, &newparams);
-                tt = (newparams ? jl_apply_tuple_type(newparams) : ti);
-                int is_compileable = ((jl_datatype_t*)ti)->isdispatchtuple ||
-                    jl_isa_compileable_sig(tt, m);
-                if (is_compileable) {
+                tt = jl_normalize_to_compilable_sig(mt, ti, env, m);
+                if (tt != jl_nothing) {
                     nf = jl_specializations_get_linfo(m, (jl_value_t*)tt, env);
                 }
             }
@@ -2075,14 +2084,6 @@ JL_DLLEXPORT int jl_compile_hint(jl_tupletype_t *types)
     return 1;
 }
 
-JL_DLLEXPORT jl_value_t *jl_get_spec_lambda(jl_tupletype_t *types, size_t world, size_t *min_valid, size_t *max_valid)
-{
-    jl_method_instance_t *mi = jl_get_specialization1(types, world, min_valid, max_valid, 0);
-    if (!mi)
-        return jl_nothing;
-    return (jl_value_t*)mi;
-}
-
 // add type of `f` to front of argument tuple type
 jl_value_t *jl_argtype_with_function(jl_function_t *f, jl_value_t *types0)
 {
@@ -2437,50 +2438,6 @@ static jl_value_t *jl_gf_invoke_by_method(jl_method_t *method, jl_value_t *gf, j
     return _jl_invoke(gf, args, nargs - 1, mfunc, world);
 }
 
-JL_DLLEXPORT jl_value_t *jl_get_invoke_lambda(jl_method_t *method, jl_value_t *tt)
-{
-    // TODO: refactor this method to be more like `jl_get_specialization1`
-    if (!jl_is_datatype(tt) || !((jl_datatype_t*)tt)->isdispatchtuple)
-        return jl_nothing;
-
-    jl_typemap_entry_t *tm = NULL;
-    struct jl_typemap_assoc search = {(jl_value_t*)tt, 1, NULL, 0, ~(size_t)0};
-    if (method->invokes != NULL) {
-        tm = jl_typemap_assoc_by_type(method->invokes, &search, /*offs*/1, /*subtype*/1);
-        if (tm) {
-            return (jl_value_t*)tm->func.linfo;
-        }
-    }
-
-    JL_LOCK(&method->writelock);
-    if (method->invokes != NULL) {
-        tm = jl_typemap_assoc_by_type(method->invokes, &search, /*offs*/1, /*subtype*/1);
-        if (tm) {
-            jl_method_instance_t *mfunc = tm->func.linfo;
-            JL_UNLOCK(&method->writelock);
-            return (jl_value_t*)mfunc;
-        }
-    }
-
-    jl_svec_t *tpenv = jl_emptysvec;
-    JL_GC_PUSH1(&tpenv);
-    if (jl_is_unionall(method->sig)) {
-        jl_value_t *ti =
-            jl_type_intersection_env(tt, (jl_value_t*)method->sig, &tpenv);
-        assert(ti != (jl_value_t*)jl_bottom_type);
-        (void)ti;
-    }
-
-    if (method->invokes == NULL)
-        method->invokes = jl_nothing;
-
-    jl_method_instance_t *mfunc = cache_method(NULL, &method->invokes, (jl_value_t*)method,
-                                               (jl_tupletype_t*)tt, method, 1, 1, ~(size_t)0, tpenv);
-    JL_GC_POP();
-    JL_UNLOCK(&method->writelock);
-    return (jl_value_t*)mfunc;
-}
-
 // Return value is rooted globally
 jl_function_t *jl_new_generic_function_with_supertype(jl_sym_t *name, jl_module_t *module, jl_datatype_t *st)
 {