Cache lookup of matching methods between inference and inlining (#34339)

The results of calling `_methods_by_ftype` in inference are now kept in
`InferenceState` objects, forwarded to `OptimizationState`s, and used in
inlining.

base/compiler/abstractinterpretation.jl

@@ -42,12 +42,22 @@ function abstract_call_gf_by_type(@nospecialize(f), argtypes::Vector{Any}, @nosp
         splitsigs = switchtupleunion(atype)
         applicable = Any[]
         for sig_n in splitsigs
-            xapplicable = _methods_by_ftype(sig_n, max_methods, sv.params.world, min_valid, max_valid)
+            (xapplicable, min_valid[1], max_valid[1]) =
+                get!(sv.matching_methods_cache, sig_n) do
+                    ms = _methods_by_ftype(sig_n, max_methods, sv.params.world,
+                                           min_valid, max_valid)
+                    return (ms, min_valid[1], max_valid[1])
+                end
             xapplicable === false && return Any
             append!(applicable, xapplicable)
         end
     else
-        applicable = _methods_by_ftype(atype, max_methods, sv.params.world, min_valid, max_valid)
+        (applicable, min_valid[1], max_valid[1]) =
+            get!(sv.matching_methods_cache, atype) do
+                ms = _methods_by_ftype(atype, max_methods, sv.params.world,
+                                       min_valid, max_valid)
+                return (ms, min_valid[1], max_valid[1])
+            end
         if applicable === false
             # this means too many methods matched
             # (assume this will always be true, so we don't compute / update valid age in this case)

base/compiler/inferencestate.jl

@@ -41,6 +41,10 @@ mutable struct InferenceState
     inferred::Bool
     dont_work_on_me::Bool
 
+    # cached results of calling `_methods_by_ftype`, including `min_valid` and
+    # `max_valid`, to be used in inlining
+    matching_methods_cache::IdDict{Any, Tuple{Any, UInt, UInt}}
+
     # src is assumed to be a newly-allocated CodeInfo, that can be modified in-place to contain intermediate results
     function InferenceState(result::InferenceResult, src::CodeInfo,
                             cached::Bool, params::Params)
@@ -101,7 +105,8 @@ mutable struct InferenceState
             Vector{Tuple{InferenceState,LineNum}}(), # cycle_backedges
             Vector{InferenceState}(), # callers_in_cycle
             #=parent=#nothing,
-            cached, false, false, false)
+            cached, false, false, false,
+            IdDict{Any, Tuple{Any, UInt, UInt}}())
         result.result = frame
         cached && push!(params.cache, result)
         return frame

base/compiler/optimize.jl

@@ -16,6 +16,9 @@ mutable struct OptimizationState
     sptypes::Vector{Any} # static parameters
     slottypes::Vector{Any}
     const_api::Bool
+    # cached results of calling `_methods_by_ftype` from inference, including
+    # `min_valid` and `max_valid`
+    matching_methods_cache::IdDict{Any, Tuple{Any, UInt, UInt}}
     function OptimizationState(frame::InferenceState)
         s_edges = frame.stmt_edges[1]
         if s_edges === nothing
@@ -27,7 +30,8 @@ mutable struct OptimizationState
                    s_edges::Vector{Any},
                    src, frame.mod, frame.nargs,
                    frame.min_valid, frame.max_valid,
-                   frame.params, frame.sptypes, frame.slottypes, false)
+                   frame.params, frame.sptypes, frame.slottypes, false,
+                   frame.matching_methods_cache)
     end
     function OptimizationState(linfo::MethodInstance, src::CodeInfo,
                                params::Params)
@@ -57,7 +61,8 @@ mutable struct OptimizationState
                    s_edges::Vector{Any},
                    src, inmodule, nargs,
                    UInt(1), get_world_counter(),
-                   params, sptypes_from_meth_instance(linfo), slottypes, false)
+                   params, sptypes_from_meth_instance(linfo), slottypes, false,
+                   IdDict{Any, Tuple{Any, UInt, UInt}}())
         end
 end
 

base/compiler/ssair/inlining.jl

@@ -1005,15 +1005,24 @@ function assemble_inline_todo!(ir::IRCode, sv::OptimizationState)
             continue
         end
 
-        # Regular case: Perform method matching
-        min_valid = UInt[typemin(UInt)]
-        max_valid = UInt[typemax(UInt)]
-        meth = _methods_by_ftype(sig.atype, sv.params.MAX_METHODS, sv.params.world, min_valid, max_valid)
+        # Regular case: Retrieve matching methods from cache (or compute them)
+        (meth, min_valid, max_valid) = get(sv.matching_methods_cache, sig.atype) do
+            # 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.
+            min_val = UInt[typemin(UInt)]
+            max_val = UInt[typemax(UInt)]
+            ms = _methods_by_ftype(sig.atype, sv.params.MAX_METHODS,
+                                   sv.params.world, min_val, max_val)
+            return (ms, min_val[1], max_val[1])
+        end
         if meth === false || length(meth) == 0
             # No applicable method, or too many applicable methods
             continue
         end
-        update_valid_age!(min_valid[1], max_valid[1], sv)
+        update_valid_age!(min_valid, max_valid, sv)
 
         cases = Pair{Any, Any}[]
         # TODO: This could be better