EA: use embedded CodeInstance directly for escape cache lookup

Compiler/src/optimize.jl

@@ -647,9 +647,11 @@ struct GetNativeEscapeCache{CodeCache}
     GetNativeEscapeCache(code_cache::CodeCache) where CodeCache = new{CodeCache}(code_cache)
 end
 GetNativeEscapeCache(interp::AbstractInterpreter) = GetNativeEscapeCache(code_cache(interp))
-function ((; code_cache)::GetNativeEscapeCache)(mi::MethodInstance)
-    codeinst = get(code_cache, mi, nothing)
-    codeinst isa CodeInstance || return false
+function ((; code_cache)::GetNativeEscapeCache)(codeinst::Union{CodeInstance,MethodInstance})
+    if codeinst isa MethodInstance
+        codeinst = get(code_cache, codeinst, nothing)
+        codeinst isa CodeInstance || return false
+    end
     argescapes = traverse_analysis_results(codeinst) do @nospecialize result
         return result isa EscapeAnalysis.ArgEscapeCache ? result : nothing
     end

Compiler/src/ssair/EscapeAnalysis.jl

@@ -944,14 +944,16 @@ end
 
 # escape statically-resolved call, i.e. `Expr(:invoke, ::MethodInstance, ...)`
 function escape_invoke!(astate::AnalysisState, pc::Int, args::Vector{Any})
-    mi = first(args)
-    if !(mi isa MethodInstance)
-        mi = (mi::CodeInstance).def # COMBAK get escape info directly from CI instead?
+    codeinst = first(args)
+    if codeinst isa MethodInstance
+        mi = codeinst
+    else
+        mi = (codeinst::CodeInstance).def
     end
     first_idx, last_idx = 2, length(args)
     add_liveness_changes!(astate, pc, args, first_idx, last_idx)
     # TODO inspect `astate.ir.stmts[pc][:info]` and use const-prop'ed `InferenceResult` if available
-    cache = astate.get_escape_cache(mi)
+    cache = astate.get_escape_cache(codeinst)
     ret = SSAValue(pc)
     if cache isa Bool
         if cache

Compiler/test/EAUtils.jl

@@ -14,27 +14,11 @@ import .Compiler:
     AbstractInterpreter, NativeInterpreter, WorldView, WorldRange, InferenceParams,
     OptimizationParams, get_world_counter, get_inference_cache, ipo_dataflow_analysis!
 # usings
-using Core:
-    CodeInstance, MethodInstance, CodeInfo
-using .Compiler:
-    InferenceResult, InferenceState, OptimizationState, IRCode
+using Core.IR
+using .Compiler: InferenceResult, InferenceState, OptimizationState, IRCode
 using .EA: analyze_escapes, ArgEscapeCache, ArgEscapeInfo, EscapeInfo, EscapeState
 
-struct EAToken end
-
-# when working outside of CC,
-# cache entire escape state for later inspection and debugging
-struct EscapeCacheInfo
-    argescapes::ArgEscapeCache
-    state::EscapeState # preserved just for debugging purpose
-    ir::IRCode         # preserved just for debugging purpose
-end
-
-struct EscapeCache
-    cache::IdDict{MethodInstance,EscapeCacheInfo} # TODO(aviatesk) Should this be CodeInstance to EscapeCacheInfo?
-end
-EscapeCache() = EscapeCache(IdDict{MethodInstance,EscapeCacheInfo}())
-const GLOBAL_ESCAPE_CACHE = EscapeCache()
+mutable struct EscapeAnalyzerCacheToken end
 
 struct EscapeResultForEntry
     ir::IRCode
@@ -47,78 +31,74 @@ mutable struct EscapeAnalyzer <: AbstractInterpreter
     const inf_params::InferenceParams
     const opt_params::OptimizationParams
     const inf_cache::Vector{InferenceResult}
-    const escape_cache::EscapeCache
+    const token::EscapeAnalyzerCacheToken
     const entry_mi::Union{Nothing,MethodInstance}
     result::EscapeResultForEntry
-    function EscapeAnalyzer(world::UInt, escape_cache::EscapeCache;
+    function EscapeAnalyzer(world::UInt, cache_token::EscapeAnalyzerCacheToken;
                             entry_mi::Union{Nothing,MethodInstance}=nothing)
         inf_params = InferenceParams()
         opt_params = OptimizationParams()
         inf_cache = InferenceResult[]
-        return new(world, inf_params, opt_params, inf_cache, escape_cache, entry_mi)
+        return new(world, inf_params, opt_params, inf_cache, cache_token, entry_mi)
     end
 end
 
 Compiler.InferenceParams(interp::EscapeAnalyzer) = interp.inf_params
 Compiler.OptimizationParams(interp::EscapeAnalyzer) = interp.opt_params
 Compiler.get_inference_world(interp::EscapeAnalyzer) = interp.world
 Compiler.get_inference_cache(interp::EscapeAnalyzer) = interp.inf_cache
-Compiler.cache_owner(::EscapeAnalyzer) = EAToken()
-Compiler.get_escape_cache(interp::EscapeAnalyzer) = GetEscapeCache(interp)
+Compiler.cache_owner(interp::EscapeAnalyzer) = interp.token
+Compiler.get_escape_cache(::EscapeAnalyzer) = GetEscapeCache()
 
 function Compiler.ipo_dataflow_analysis!(interp::EscapeAnalyzer, opt::OptimizationState,
-                                   ir::IRCode, caller::InferenceResult)
+                                         ir::IRCode, caller::InferenceResult)
     # run EA on all frames that have been optimized
     nargs = Int(opt.src.nargs)
     𝕃ₒ = Compiler.optimizer_lattice(interp)
-    get_escape_cache = GetEscapeCache(interp)
     estate = try
-        analyze_escapes(ir, nargs, 𝕃ₒ, get_escape_cache)
+        analyze_escapes(ir, nargs, 𝕃ₒ, GetEscapeCache())
     catch err
         @error "error happened within EA, inspect `Main.failedanalysis`"
-        failedanalysis = FailedAnalysis(caller, ir, nargs, get_escape_cache)
+        failedanalysis = FailedAnalysis(caller, ir, nargs)
         Core.eval(Main, :(failedanalysis = $failedanalysis))
         rethrow(err)
     end
     if caller.linfo === interp.entry_mi
         # return back the result
         interp.result = EscapeResultForEntry(Compiler.copy(ir), estate, caller.linfo)
     end
-    record_escapes!(interp, caller, estate, ir)
+    record_escapes!(caller, estate, ir)
 
     @invoke Compiler.ipo_dataflow_analysis!(interp::AbstractInterpreter, opt::OptimizationState,
-                                      ir::IRCode, caller::InferenceResult)
+                                            ir::IRCode, caller::InferenceResult)
 end
 
-function record_escapes!(interp::EscapeAnalyzer,
-    caller::InferenceResult, estate::EscapeState, ir::IRCode)
+# cache entire escape state for inspection and debugging
+struct EscapeCacheInfo
+    argescapes::ArgEscapeCache
+    state::EscapeState # preserved just for debugging purpose
+    ir::IRCode         # preserved just for debugging purpose
+end
+
+function record_escapes!(caller::InferenceResult, estate::EscapeState, ir::IRCode)
     argescapes = ArgEscapeCache(estate)
     ecacheinfo = EscapeCacheInfo(argescapes, estate, ir)
     return Compiler.stack_analysis_result!(caller, ecacheinfo)
 end
 
-struct GetEscapeCache
-    escape_cache::EscapeCache
-    GetEscapeCache(interp::EscapeAnalyzer) = new(interp.escape_cache)
-end
-function ((; escape_cache)::GetEscapeCache)(mi::MethodInstance)
-    ecacheinfo = get(escape_cache.cache, mi, nothing)
+struct GetEscapeCache end
+function (::GetEscapeCache)(codeinst::Union{CodeInstance,MethodInstance})
+    codeinst isa CodeInstance || return false
+    ecacheinfo = Compiler.traverse_analysis_results(codeinst) do @nospecialize result
+        return result isa EscapeCacheInfo ? result : nothing
+    end
     return ecacheinfo === nothing ? false : ecacheinfo.argescapes
 end
 
 struct FailedAnalysis
     caller::InferenceResult
     ir::IRCode
     nargs::Int
-    get_escape_cache::GetEscapeCache
-end
-
-function Compiler.finish!(interp::EscapeAnalyzer, state::InferenceState; can_discard_trees::Bool=Compiler.may_discard_trees(interp))
-    ecacheinfo = Compiler.traverse_analysis_results(state.result) do @nospecialize result
-        return result isa EscapeCacheInfo ? result : nothing
-    end
-    ecacheinfo isa EscapeCacheInfo && (interp.escape_cache.cache[state.linfo] = ecacheinfo)
-    return @invoke Compiler.finish!(interp::AbstractInterpreter, state::InferenceState; can_discard_trees)
 end
 
 # printing
@@ -313,23 +293,29 @@ while caching the analysis results.
 
 - `world::UInt = Base.get_world_counter()`:
   controls the world age to use when looking up methods, use current world age if not specified.
-- `interp::EscapeAnalyzer = EscapeAnalyzer(world)`:
-  specifies the escape analyzer to use, by default a new analyzer with the global cache is created.
+- `cache_token::EscapeAnalyzerCacheToken = EscapeAnalyzerCacheToken()`:
+  specifies the cache token to use, by default a new token is generated to ensure
+  that `code_escapes` uses a fresh cache and performs a new analysis on each invocation.
+  If you with to perform analysis with the global cache enabled, specify a particular token instance.
+- `interp::EscapeAnalyzer = EscapeAnalyzer(world, cache_token)`:
+  specifies the escape analyzer to use.
 - `debuginfo::Symbol = :none`:
   controls the amount of code metadata present in the output, possible options are `:none` or `:source`.
 """
 function code_escapes(@nospecialize(f), @nospecialize(types=Base.default_tt(f));
                       world::UInt = get_world_counter(),
+                      cache_token::EscapeAnalyzerCacheToken = EscapeAnalyzerCacheToken(),
                       debuginfo::Symbol = :none)
     tt = Base.signature_type(f, types)
     match = Base._which(tt; world, raise=true)
     mi = Compiler.specialize_method(match)
-    return code_escapes(mi; world, debuginfo)
+    return code_escapes(mi; world, cache_token, debuginfo)
 end
 
 function code_escapes(mi::MethodInstance;
                       world::UInt = get_world_counter(),
-                      interp::EscapeAnalyzer=EscapeAnalyzer(world, GLOBAL_ESCAPE_CACHE; entry_mi=mi),
+                      cache_token::EscapeAnalyzerCacheToken = EscapeAnalyzerCacheToken(),
+                      interp::EscapeAnalyzer=EscapeAnalyzer(world, cache_token; entry_mi=mi),
                       debuginfo::Symbol = :none)
     frame = Compiler.typeinf_frame(interp, mi, #=run_optimizer=#true)
     isdefined(interp, :result) || error("optimization didn't happen: maybe everything has been constant folded?")
@@ -351,12 +337,17 @@ Note that this version does not cache the analysis results.
 
 - `world::UInt = Base.get_world_counter()`:
   controls the world age to use when looking up methods, use current world age if not specified.
-- `interp::AbstractInterpreter = EscapeAnalyzer(world, EscapeCache())`:
+- `cache_token::EscapeAnalyzerCacheToken = EscapeAnalyzerCacheToken()`:
+  specifies the cache token to use, by default a new token is generated to ensure
+  that `code_escapes` uses a fresh cache and performs a new analysis on each invocation.
+  If you with to perform analysis with the global cache enabled, specify a particular token instance.
+- `interp::AbstractInterpreter = EscapeAnalyzer(world, cache_token)`:
   specifies the abstract interpreter to use, by default a new `EscapeAnalyzer` with an empty cache is created.
 """
 function code_escapes(ir::IRCode, nargs::Int;
                       world::UInt = get_world_counter(),
-                      interp::AbstractInterpreter=EscapeAnalyzer(world, EscapeCache()))
+                      cache_token::EscapeAnalyzerCacheToken = EscapeAnalyzerCacheToken(),
+                      interp::AbstractInterpreter=EscapeAnalyzer(world, cache_token))
     estate = analyze_escapes(ir, nargs, Compiler.optimizer_lattice(interp), Compiler.get_escape_cache(interp))
     return EscapeResult(ir, estate) # return back the result
 end

Compiler/test/EscapeAnalysis.jl

@@ -1527,7 +1527,17 @@ let result = code_escapes() do
     end
     i = last(findall(isnew, result.ir.stmts.stmt))
     @test_broken !has_return_escape(result.state[SSAValue(i)]) # TODO interprocedural alias analysis
-    @test !has_thrown_escape(result.state[SSAValue(i)])
+    @test_broken !has_thrown_escape(result.state[SSAValue(i)]) # IDEA embed const-prop'ed `CodeInstance` for `:invoke`?
+end
+let result = code_escapes((Base.RefValue{Base.RefValue{String}},)) do x
+        out1 = broadcast_noescape2(Ref(Ref("Hi")))
+        out2 = broadcast_noescape2(x)
+        return out1, out2
+    end
+    i = last(findall(isnew, result.ir.stmts.stmt))
+    @test_broken !has_return_escape(result.state[SSAValue(i)]) # TODO interprocedural alias analysis
+    @test_broken !has_thrown_escape(result.state[SSAValue(i)]) # IDEA embed const-prop'ed `CodeInstance` for `:invoke`?
+    @test has_thrown_escape(result.state[Argument(2)])
 end
 @noinline allescape_argument(a) = (global GV = a) # obvious escape
 let result = code_escapes() do