CompilerDevTools: add proof of concept for caching runtime calls

Compiler/extras/CompilerDevTools/src/CompilerDevTools.jl

@@ -2,43 +2,71 @@ module CompilerDevTools
 
 using Compiler
 using Core.IR
+using Base: isexpr
+
+struct SplitCacheOwner
+    version::Int # use a non-zero value when you don't want to reuse the default cache, e.g. for testing purposes.
+end
+
+SplitCacheOwner() = SplitCacheOwner(0)
 
-struct SplitCacheOwner; end
 struct SplitCacheInterp <: Compiler.AbstractInterpreter
     world::UInt
+    owner::SplitCacheOwner
     inf_params::Compiler.InferenceParams
     opt_params::Compiler.OptimizationParams
     inf_cache::Vector{Compiler.InferenceResult}
     function SplitCacheInterp(;
         world::UInt = Base.get_world_counter(),
+        owner::SplitCacheOwner = SplitCacheOwner(),
         inf_params::Compiler.InferenceParams = Compiler.InferenceParams(),
         opt_params::Compiler.OptimizationParams = Compiler.OptimizationParams(),
         inf_cache::Vector{Compiler.InferenceResult} = Compiler.InferenceResult[])
-        new(world, inf_params, opt_params, inf_cache)
+        new(world, owner, inf_params, opt_params, inf_cache)
     end
 end
 
 Compiler.InferenceParams(interp::SplitCacheInterp) = interp.inf_params
 Compiler.OptimizationParams(interp::SplitCacheInterp) = interp.opt_params
 Compiler.get_inference_world(interp::SplitCacheInterp) = interp.world
 Compiler.get_inference_cache(interp::SplitCacheInterp) = interp.inf_cache
-Compiler.cache_owner(::SplitCacheInterp) = SplitCacheOwner()
+Compiler.cache_owner(interp::SplitCacheInterp) = interp.owner
 
 import Core.OptimizedGenerics.CompilerPlugins: typeinf, typeinf_edge
-@eval @noinline typeinf(::SplitCacheOwner, mi::MethodInstance, source_mode::UInt8) =
-    Base.invoke_in_world(which(typeinf, Tuple{SplitCacheOwner, MethodInstance, UInt8}).primary_world, Compiler.typeinf_ext, SplitCacheInterp(; world=Base.tls_world_age()), mi, source_mode)
+@eval @noinline typeinf(owner::SplitCacheOwner, mi::MethodInstance, source_mode::UInt8) =
+    Base.invoke_in_world(which(typeinf, Tuple{SplitCacheOwner, MethodInstance, UInt8}).primary_world, Compiler.typeinf_ext, SplitCacheInterp(; world=Base.tls_world_age(), owner), mi, source_mode)
 
-@eval @noinline function typeinf_edge(::SplitCacheOwner, mi::MethodInstance, parent_frame::Compiler.InferenceState, world::UInt, source_mode::UInt8)
+@eval @noinline function typeinf_edge(owner::SplitCacheOwner, mi::MethodInstance, parent_frame::Compiler.InferenceState, world::UInt, source_mode::UInt8)
     # TODO: This isn't quite right, we're just sketching things for now
-    interp = SplitCacheInterp(; world)
+    interp = SplitCacheInterp(; world, owner)
     Compiler.typeinf_edge(interp, mi.def, mi.specTypes, Core.svec(), parent_frame, false, false)
 end
 
-function with_new_compiler(f, args...)
-    mi = @ccall jl_method_lookup(Any[f, args...]::Ptr{Any}, (1+length(args))::Csize_t, Base.tls_world_age()::Csize_t)::Ref{Core.MethodInstance}
-    world = Base.tls_world_age()
+function lookup_method_instance(f, args...)
+    @ccall jl_method_lookup(Any[f, args...]::Ptr{Any}, (1+length(args))::Csize_t, Base.tls_world_age()::Csize_t)::Ref{Core.MethodInstance}
+end
+
+function Compiler.optimize(interp::SplitCacheInterp, opt::Compiler.OptimizationState, caller::Compiler.InferenceResult)
+    @invoke Compiler.optimize(interp::Compiler.AbstractInterpreter, opt::Compiler.OptimizationState, caller::Compiler.InferenceResult)
+    ir = opt.ir::Compiler.IRCode
+    override = GlobalRef(@__MODULE__(), :with_new_compiler)
+    for inst in ir.stmts
+        stmt = inst[:stmt]
+        isexpr(stmt, :call) || continue
+        stmt.args[1] === override && continue
+        insert!(stmt.args, 1, override)
+        insert!(stmt.args, 3, interp.owner)
+    end
+end
+
+with_new_compiler(f, args...; owner::SplitCacheOwner = SplitCacheOwner()) = with_new_compiler(f, owner, args...)
+
+with_new_compiler(f::Core.Builtin, ::SplitCacheOwner, args...) = f(args...)
+
+function with_new_compiler(f, owner::SplitCacheOwner, args...)
+    mi = lookup_method_instance(f, args...)
     new_compiler_ci = Core.OptimizedGenerics.CompilerPlugins.typeinf(
-        SplitCacheOwner(), mi, Compiler.SOURCE_MODE_ABI
+        owner, mi, Compiler.SOURCE_MODE_ABI
     )
     invoke(f, new_compiler_ci, args...)
 end

Compiler/extras/CompilerDevTools/test/runtests.jl

@@ -0,0 +1,26 @@
+using Test
+using Compiler: Compiler, code_cache
+using Base: get_world_counter, inferencebarrier
+using CompilerDevTools
+using CompilerDevTools: lookup_method_instance, SplitCacheOwner, SplitCacheInterp
+
+!@isdefined(cache_version) && (cache_version = 1)
+new_cache_version() = SplitCacheOwner(global cache_version += 1)
+
+do_work(x, y) = x + y
+f1() = do_work(inferencebarrier(1), inferencebarrier(2))
+
+@testset "CompilerDevTools" begin
+  owner = new_cache_version()
+  world = get_world_counter()
+  interp = SplitCacheInterp(; world, owner)
+  cache = code_cache(interp)
+  mi = lookup_method_instance(f1)
+  @test !haskey(cache, mi)
+  @test with_new_compiler(f1; owner) === 3
+  @test haskey(cache, mi)
+  # Here `do_work` is compiled at runtime, and so must have
+  # required extra work to be cached under the same cache owner.
+  mi = lookup_method_instance(do_work, 1, 2)
+  @test haskey(cache, mi)
+end;