use local cache for uncached inlining callsite

JuliaLang · Aug 13, 2021 · ea67453 · ea67453
1 parent 66d549b
commit ea67453
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 2 deletions.
diff --git a/base/compiler/optimize.jl b/base/compiler/optimize.jl
@@ -21,10 +21,11 @@ function push!(et::EdgeTracker, ci::CodeInstance)
     push!(et, ci.def)
 end
 
-struct InliningState{S <: Union{EdgeTracker, Nothing}, T, P}
+struct InliningState{S <: Union{EdgeTracker, Nothing}, T, P, U}
     params::OptimizationParams
     et::S
     mi_cache::T
+    inf_cache::U
     policy::P
 end
 
@@ -63,6 +64,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))
         return new(frame.linfo,
                    frame.src, nothing, frame.stmt_info, frame.mod,
@@ -92,6 +94,7 @@ mutable struct OptimizationState
         inlining = InliningState(params,
             nothing,
             WorldView(code_cache(interp), get_world_counter()),
+            get_inference_cache(interp),
             inlining_policy(interp))
         return new(linfo,
                    src, nothing, stmt_info, mod,

diff --git a/base/compiler/ssair/inlining.jl b/base/compiler/ssair/inlining.jl
@@ -725,7 +725,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 = false, nothing
+    isconst, src, argtypes = false, nothing, nothing
     if isa(match, InferenceResult)
         let inferred_src = match.src
             if isa(inferred_src, Const)
@@ -737,6 +737,9 @@ 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,6 +752,9 @@ 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
@@ -758,6 +764,17 @@ 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)
 
     if src === nothing

diff --git a/test/compiler/inline.jl b/test/compiler/inline.jl
@@ -538,6 +538,9 @@ end
                 return a, b
             end
         end
+
+        import Core.Compiler: isType
+        limited(a) = @noinline(isType(a)) ? @inline(limited(a.parameters[1])) : rand(a)
     end
 
     let ci = code_typed1(m.force_inline_explicit, (Int,))
@@ -586,6 +589,10 @@ end
     let ci = code_typed1(m.nested, (Int,Int))
         @test count(x->isinvoke(x, :notinlined), ci.code) == 1
     end
+
+    let ci = code_typed1(m.limited, (Any,))
+        @test count(x->isinvoke(x, :isType), ci.code) == 2
+    end
 end
 
 # Issue #41299 - inlining deletes error check in :>