Scope compile_hook to the installing task

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,10 +44,14 @@ const HAS_INTEGRATED_CACHE = VERSION >= v"1.11.0-DEV.1552"
 # `CompilerCaching.`.
 import CompilerCaching
 
-# 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.
-const compile_hook = Ref{Union{Nothing,Function}}(nothing)
+using ScopedValues: ScopedValue, with
+
+# Optional callback invoked with the `CompilerJob` from `compile(...)` /
+# `cached_compilation(...)` before compilation runs; installed by the `@device_code_*`
+# reflection macros with `with(compile_hook => hook)`, so it applies to the current
+# task and tasks spawned within the scope. Defined here (early) so the legacy
+# `cached_compilation` in deprecated.jl can reference it.
+const compile_hook = ScopedValue{Union{Nothing,Function}}(nothing)
 
 include("utils.jl")
 include("mangling.jl")

src/reflection.jl

@@ -312,28 +312,23 @@ function emit_hooked_compilation(inner_hook, ex...)
     user_code = ex[end]
     user_kwargs = ex[1:end-1]
     quote
-        # we only want to invoke the hook once for every compilation job
+        # invoke the hook once per compilation job; compilations may run concurrently
         jobs = Set()
+        jobs_lock = ReentrantLock()
         function outer_hook(job)
-            if !in(job, jobs)
-                # the user hook might invoke the compiler again, so disable the hook
-                old_hook = $compile_hook[]
-                try
-                    $compile_hook[] = nothing
-                    $inner_hook(job; $(map(esc, user_kwargs)...))
-                finally
-                    $compile_hook[] = old_hook
-                end
+            Base.@lock jobs_lock begin
+                job in jobs && return
                 push!(jobs, job)
             end
+            # the user hook might invoke the compiler again, so disable the hook
+            $with($compile_hook => nothing) do
+                $inner_hook(job; $(map(esc, user_kwargs)...))
+            end
         end
 
         # now invoke the user code with this hook in place
-        try
-            $compile_hook[] = outer_hook
+        $with($compile_hook => outer_hook) do
             $(esc(user_code))
-        finally
-            $compile_hook[] = nothing
         end
 
         if isempty(jobs)

test/Project.toml

@@ -16,6 +16,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SPIRV_LLVM_Backend_jll = "4376b9bf-cff8-51b6-bb48-39421dff0d0c"
 SPIRV_LLVM_Translator_jll = "4a5d46fc-d8cf-5151-a261-86b458210efb"
 SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4"
+ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 demumble_jll = "1e29f10c-031c-5a83-9565-69cddfc27673"

test/native.jl

@@ -25,6 +25,42 @@
     end
 end
 
+@testset "compile hook" begin
+    mod = @eval module $(gensym())
+        f(x::Int) = nothing
+        g(x::Int) = nothing
+    end
+
+    # the hook sees every job compiled within its scope, and nothing outside it
+    seen = []
+    with(GPUCompiler.compile_hook => job -> push!(seen, job)) do
+        Native.code_execution(mod.f, (Int,))
+    end
+    @test length(seen) == 1 && only(seen).source.def.name === :f
+    @test GPUCompiler.compile_hook[] === nothing
+    Native.code_execution(mod.g, (Int,))
+    @test length(seen) == 1
+
+    # the reflection macros go through the hook once per job, even when the user
+    # code compiles the same job repeatedly
+    lowered = GPUCompiler.@device_code_lowered begin
+        Native.code_execution(mod.f, (Int,))
+        Native.code_execution(mod.f, (Int,))
+    end
+    @test length(lowered) == 1
+    @test_throws "no kernels executed" GPUCompiler.@device_code_native mod.f(1)
+
+    # scoped: tasks spawned inside the scope inherit the hook, others don't
+    hook = job -> nothing
+    inherited = Ref{Any}(nothing)
+    with(GPUCompiler.compile_hook => hook) do
+        wait(Threads.@spawn inherited[] = GPUCompiler.compile_hook[])
+        @test GPUCompiler.compile_hook[] === hook
+    end
+    @test inherited[] === hook
+    @test fetch(Threads.@spawn GPUCompiler.compile_hook[]) === nothing
+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}

test/runtests.jl

@@ -11,6 +11,7 @@ const init_code = quote
     using LLVMDowngrader_jll
     using NVPTX_LLVM_Backend_jll
     using AMDGPU_LLVM_Backend_jll
+    using ScopedValues: with
 
     # include all helpers
     include(joinpath(@__DIR__, "helpers", "runtime.jl"))