Add inference_batch for sharing const-prop results across compilations

Project.toml

@@ -16,6 +16,7 @@ Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
+ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Tracy = "e689c965-62c8-4b79-b2c5-8359227902fd"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
@@ -47,6 +48,7 @@ NVPTX_LLVM_Backend_jll = "22"
 PrecompileTools = "1.0.2"
 Preferences = "1"
 REPL = "1"
+ScopedValues = "1.5"
 TOML = "1"
 Tracy = "0.1.4"
 UUIDs = "1"

src/GPUCompiler.jl

@@ -44,6 +44,8 @@ const HAS_INTEGRATED_CACHE = VERSION >= v"1.11.0-DEV.1552"
 # `CompilerCaching.`.
 import CompilerCaching
 
+using ScopedValues: ScopedValue, with
+
 # Optional callback invoked from `compile(...)` / `cached_compilation(...)` before
 # compilation runs. Set by `@device_code_*` reflection macros. Defined here (early)
 # so the legacy `cached_compilation` in deprecated.jl can reference it.

src/interface.jl

@@ -595,6 +595,7 @@ end # HAS_INTEGRATED_CACHE
 @public apply_relocations!, resolved_relocations, resolved_relocation_table
 @public supports_relocatable_ir
 @public GPUCompilerCacheToken, cache_owner, cached_results
+@public inference_batch
 
 # the method table to use
 #

src/jlgen.jl

@@ -150,6 +150,53 @@ get_method_table_view(world::UInt, mt::CC.MethodTable) = CC.OverlayMethodTable(w
 # VERSION >= v"1.14.0-DEV.1691"
 const INFERENCE_CACHE_TYPE = isdefined(CC, :InferenceCache) ? CC.InferenceCache : Vector{CC.InferenceResult}
 
+# Interpreter-local inference caches shared within an `inference_batch`, keyed by
+# task, cache owner and world: results are only interchangeable between jobs that
+# infer identically (same owner) in the same world, and a cache is only ever used
+# by one task at a time.
+struct InferenceBatch
+    caches::Dict{Tuple{Task,Any,UInt},INFERENCE_CACHE_TYPE}
+    lock::ReentrantLock
+    InferenceBatch() = new(Dict{Tuple{Task,Any,UInt},INFERENCE_CACHE_TYPE}(), ReentrantLock())
+end
+
+const current_inference_batch = ScopedValue{Union{Nothing,InferenceBatch}}(nothing)
+
+"""
+    inference_batch(f)
+
+Run `f()` such that the interpreters GPUCompiler constructs within share their local
+inference cache — per task, per [`cache_owner`](@ref) and per world age — instead of
+each starting from an empty one.
+
+The local cache holds the results of constant propagation into callees. Julia keeps
+them for the lifetime of the interpreter, and GPUCompiler constructs an interpreter
+per compilation, so a batch of compilations of the same method with different
+constants (an autotuning sweep) re-derives them every time. Within a batch, later
+compilations on the same task reuse what earlier ones inferred.
+
+```julia
+GPUCompiler.inference_batch() do
+    for config in configs
+        compile(config)
+    end
+end
+```
+
+Sharing is per owner and world because only jobs with the same owner, in the same
+world, produce interchangeable inference results; and per task so that a cache is
+never used concurrently — tasks spawned inside the scope get their own caches.
+"""
+inference_batch(f) = with(f, current_inference_batch => InferenceBatch())
+
+# The inference cache for a new interpreter partitioned by `owner` at `world`: fresh,
+# or the current task's cache for that owner and world within an `inference_batch`.
+function inference_cache(@nospecialize(owner), world::UInt)
+    batch = current_inference_batch[]
+    batch === nothing && return INFERENCE_CACHE_TYPE()
+    Base.@lock batch.lock get!(INFERENCE_CACHE_TYPE, batch.caches, (current_task(), owner, world))
+end
+
 """
     GPUInterpreter
 
@@ -187,7 +234,7 @@ function GPUInterpreter(world::UInt=Base.get_world_counter();
                         always_inline::Bool=false)
     @assert world <= Base.get_world_counter()
     return GPUInterpreter{typeof(method_table_view)}(
-        world, method_table_view, owner, INFERENCE_CACHE_TYPE(),
+        world, method_table_view, owner, inference_cache(owner, world),
         inf_params, opt_params, always_inline)
 end
 
@@ -216,7 +263,7 @@ function GPUInterpreter(world::UInt=Base.get_world_counter();
                         always_inline::Bool=false)
     @assert world <= Base.get_world_counter()
     return GPUInterpreter{typeof(method_table_view)}(
-        world, method_table_view, code_cache, Vector{CC.InferenceResult}(),
+        world, method_table_view, code_cache, inference_cache(code_cache, world),
         inf_params, opt_params, always_inline)
 end
 

test/native.jl

@@ -25,6 +25,38 @@
     end
 end
 
+@testset "inference batch" begin
+    mod = @eval module $(gensym())
+        f(x::Int) = nothing
+    end
+    job, _ = Native.create_job(mod.f, (Int,))
+    other_job, _ = Native.create_job(mod.f, (Int,); always_inline=true)   # other owner
+    cache_of(job) = GPUCompiler.get_interpreter(job).inf_cache
+
+    # outside a batch: a fresh cache per interpreter
+    @test cache_of(job) !== cache_of(job)
+
+    GPUCompiler.inference_batch() do
+        # same task, same owner: shared
+        @test cache_of(job) === cache_of(job)
+        # different owner: separate
+        @test cache_of(job) !== cache_of(other_job)
+        # different world: separate
+        older_job = GPUCompiler.CompilerJob(job.source, job.config, job.world - 1)
+        @test cache_of(older_job) !== cache_of(job)
+        # different task: separate, and stable within that task
+        spawned = fetch(Threads.@spawn (cache_of(job), cache_of(job)))
+        @test spawned[1] === spawned[2]
+        @test spawned[1] !== cache_of(job)
+        # compilation uses it
+        Native.code_execution(mod.f, (Int,))
+        @test cache_of(job) === cache_of(job)
+    end
+
+    # the batch is gone with its scope
+    @test cache_of(job) !== cache_of(job)
+end
+
 @testset "method instances for type-valued callees and arguments" begin
     # JuliaLang/julia#62001: closed type-valued callees and arguments
     # dispatch on Core.TypeEgal keys instead of Type{T}