newinterp: add external cache mode (#54074)

When working on debugging inference behavior or performance tuning,
having direct access to inspect the global cache can be really handy.
Right now, `@newinterp` utilizes the `Core.Compiler.InternalCodeCache`
for its global caching, which doesn’t allow for such inspections. This
commit introduces the `@newinterp NewInterp cache_externally=true`
option. This enhancement allows `NewInterp` to store `CodeInstance` in
an external cache whose type is `IdDict{MethodInstance, CodeInstance}`.

Author: aviatesk

test/compiler/AbstractInterpreter.jl

@@ -514,3 +514,20 @@ let src = code_typed((Int,); interp=CustomDataInterp()) do x
     @test count(isinvoke(:cos), src.code) == 1
     @test count(isinvoke(:+), src.code) == 0
 end
+
+# ephemeral cache mode
+@newinterp DebugInterp #=ephemeral_cache=#true
+func_ext_cache1(a) = func_ext_cache2(a) * cos(a)
+func_ext_cache2(a) = sin(a)
+let interp = DebugInterp()
+    @test Base.infer_return_type(func_ext_cache1, (Float64,); interp) === Float64
+    @test isdefined(interp, :code_cache)
+    found = false
+    for (mi, codeinst) in interp.code_cache.dict
+        if mi.def.name === :func_ext_cache2
+            found = true
+            break
+        end
+    end
+    @test found
+end

test/compiler/newinterp.jl

@@ -3,35 +3,62 @@
 # TODO set up a version who defines new interpreter with persistent cache?
 
 """
-    @newinterp NewInterpreter
+    @newinterp NewInterpreter [ephemeral_cache::Bool=false]
 
 Defines new `NewInterpreter <: AbstractInterpreter` whose cache is separated
 from the native code cache, satisfying the minimum interface requirements.
+
+When the `ephemeral_cache=true` option is specified, `NewInterpreter` will hold
+`CodeInstance` in an ephemeral non-integrated cache, rather than in the integrated
+`Core.Compiler.InternalCodeCache`.
+Keep in mind that ephemeral cache lacks support for invalidation and doesn't persist across
+sessions. However it is an usual Julia object of the type `code_cache::IdDict{MethodInstance,CodeInstance}`,
+making it easier for debugging and inspecting the compiler behavior.
 """
-macro newinterp(InterpName)
-    cache_token = QuoteNode(gensym(string(InterpName, "Cache")))
+macro newinterp(InterpName, ephemeral_cache::Bool=false)
+    cache_token = QuoteNode(gensym(string(InterpName, "CacheToken")))
+    InterpCacheName = esc(Symbol(string(InterpName, "Cache")))
     InterpName = esc(InterpName)
     C = Core
     CC = Core.Compiler
     quote
+        $(ephemeral_cache && quote
+        struct $InterpCacheName
+            dict::IdDict{$C.MethodInstance,$C.CodeInstance}
+        end
+        $InterpCacheName() = $InterpCacheName(IdDict{$C.MethodInstance,$C.CodeInstance}())
+        end)
         struct $InterpName <: $CC.AbstractInterpreter
             meta # additional information
             world::UInt
             inf_params::$CC.InferenceParams
             opt_params::$CC.OptimizationParams
             inf_cache::Vector{$CC.InferenceResult}
+            $(ephemeral_cache && :(code_cache::$InterpCacheName))
             function $InterpName(meta = nothing;
                                  world::UInt = Base.get_world_counter(),
                                  inf_params::$CC.InferenceParams = $CC.InferenceParams(),
                                  opt_params::$CC.OptimizationParams = $CC.OptimizationParams(),
-                                 inf_cache::Vector{$CC.InferenceResult} = $CC.InferenceResult[])
-                return new(meta, world, inf_params, opt_params, inf_cache)
+                                 inf_cache::Vector{$CC.InferenceResult} = $CC.InferenceResult[],
+                                 $(ephemeral_cache ?
+                                    Expr(:kw, :(code_cache::$InterpCacheName), :($InterpCacheName())) :
+                                    Expr(:kw, :_, :nothing)))
+                return $(ephemeral_cache ?
+                    :(new(meta, world, inf_params, opt_params, inf_cache, code_cache)) :
+                    :(new(meta, world, inf_params, opt_params, inf_cache)))
             end
         end
         $CC.InferenceParams(interp::$InterpName) = interp.inf_params
         $CC.OptimizationParams(interp::$InterpName) = interp.opt_params
         $CC.get_inference_world(interp::$InterpName) = interp.world
         $CC.get_inference_cache(interp::$InterpName) = interp.inf_cache
         $CC.cache_owner(::$InterpName) = $cache_token
+        $(ephemeral_cache && quote
+        $CC.code_cache(interp::$InterpName) = $CC.WorldView(interp.code_cache, $CC.WorldRange(interp.world))
+        $CC.get(wvc::$CC.WorldView{$InterpCacheName}, mi::$C.MethodInstance, default) = get(wvc.cache.dict, mi, default)
+        $CC.getindex(wvc::$CC.WorldView{$InterpCacheName}, mi::$C.MethodInstance) = getindex(wvc.cache.dict, mi)
+        $CC.haskey(wvc::$CC.WorldView{$InterpCacheName}, mi::$C.MethodInstance) = haskey(wvc.cache.dict, mi)
+        $CC.setindex!(wvc::$CC.WorldView{$InterpCacheName}, ci::$C.CodeInstance, mi::$C.MethodInstance) = setindex!(wvc.cache.dict, ci, mi)
+        end)
     end
 end