diff --git a/Project.toml b/Project.toml index 92cb05d0..4cf15d93 100644 --- a/Project.toml +++ b/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" diff --git a/src/GPUCompiler.jl b/src/GPUCompiler.jl index a093b9cc..5721c871 100644 --- a/src/GPUCompiler.jl +++ b/src/GPUCompiler.jl @@ -44,10 +44,12 @@ 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 + +# Hook used by the `@device_code_*` macros. Scope it to the current task and its +# children so concurrent reflection calls do not interfere. Defined here 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") diff --git a/src/reflection.jl b/src/reflection.jl index 2b498a44..e84da2a9 100644 --- a/src/reflection.jl +++ b/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 + # The job set and hook output are shared by child tasks, so update them together. jobs = Set() + jobs_lock = ReentrantLock() function outer_hook(job) - if !in(job, jobs) + Base.@lock jobs_lock begin + job in jobs && return + push!(jobs, job) # the user hook might invoke the compiler again, so disable the hook - old_hook = $compile_hook[] - try - $compile_hook[] = nothing + $with($compile_hook => nothing) do $inner_hook(job; $(map(esc, user_kwargs)...)) - finally - $compile_hook[] = old_hook end - push!(jobs, job) 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) diff --git a/test/Project.toml b/test/Project.toml index bfcc59e0..8d2eba27 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/native.jl b/test/native.jl index 57cf19ae..5729430a 100644 --- a/test/native.jl +++ b/test/native.jl @@ -25,6 +25,73 @@ 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) + + # concurrent compilations within the scope are all observed, exactly once each + mod2 = @eval module $(gensym()) + $((:($(Symbol(:f, i))(x::Int) = nothing) for i in 1:8)...) + end + fs = [getfield(mod2, Symbol(:f, i)) for i in 1:8] + typed = GPUCompiler.@device_code_typed begin + @sync for f in fs, _ in 1:2 + Threads.@spawn Native.code_execution(f, (Int,)) + end + end + @test length(typed) == 8 + @test Set(job.source.def.name for job in keys(typed)) == Set(Symbol(:f, i) for i in 1:8) + + # simultaneous scopes in unrelated tasks do not observe each other's jobs + ready = Channel(2) + proceed = Channel(2) + tasks = map((mod.f, mod.g)) do f + Threads.@spawn GPUCompiler.@device_code_typed begin + put!(ready, nothing) + take!(proceed) + Native.code_execution(f, (Int,)) + end + end + take!(ready) + take!(ready) + put!(proceed, nothing) + put!(proceed, nothing) + outputs = fetch.(tasks) + @test all(length(output) == 1 for output in outputs) + @test Set(only(keys(output)).source.def.name for output in outputs) == Set((:f, :g)) + + # 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} diff --git a/test/runtests.jl b/test/runtests.jl index 75838c96..5db1a0e2 100644 --- a/test/runtests.jl +++ b/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"))