Restore current_job across nested compilation. (#855)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: maleadt

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "1.22.3"
+version = "1.22.4"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [workspace]

src/driver.jl

@@ -82,44 +82,53 @@ function compile_unhooked(output::Symbol, @nospecialize(job::CompilerJob); kwarg
         error("No active LLVM context. Use `JuliaContext()` do-block syntax to create one.")
     end
 
-    @tracepoint "Validation" begin
-        check_method(job)   # not optional
-        job.config.validate && check_invocation(job)
-    end
+    # save/restore the `current_job` global that passes read: compilation nests (the relink
+    # pass rebuilds the runtime via `compile_unhooked` on a cold cache), and a leaked inner
+    # `kernel=false` job would make the outer kernel-state passes misfire.
+    global current_job
+    prev_job = @isdefined(current_job) ? current_job : nothing
+    try
+        @tracepoint "Validation" begin
+            check_method(job)   # not optional
+            job.config.validate && check_invocation(job)
+        end
 
-    prepare_job!(job)
+        prepare_job!(job)
 
 
-    ## LLVM IR
+        ## LLVM IR
 
-    ir, ir_meta = emit_llvm(job)
+        ir, ir_meta = emit_llvm(job)
 
-    if output == :llvm
-        if job.config.strip
-            @tracepoint "strip debug info" strip_debuginfo!(ir)
-        end
+        if output == :llvm
+            if job.config.strip
+                @tracepoint "strip debug info" strip_debuginfo!(ir)
+            end
 
-        return ir, ir_meta
-    end
+            return ir, ir_meta
+        end
 
 
-    ## machine code
+        ## machine code
 
-    format = if output == :asm
-        LLVM.API.LLVMAssemblyFile
-    elseif output == :obj
-        LLVM.API.LLVMObjectFile
-    else
-        error("Unknown assembly format $output")
-    end
-    asm, asm_meta = emit_asm(job, ir, format)
+        format = if output == :asm
+            LLVM.API.LLVMAssemblyFile
+        elseif output == :obj
+            LLVM.API.LLVMObjectFile
+        else
+            error("Unknown assembly format $output")
+        end
+        asm, asm_meta = emit_asm(job, ir, format)
 
-    if output == :asm || output == :obj
-        return asm, (; asm_meta..., ir_meta..., ir)
-    end
+        if output == :asm || output == :obj
+            return asm, (; asm_meta..., ir_meta..., ir)
+        end
 
 
-    error("Unknown compilation output $output")
+        error("Unknown compilation output $output")
+    finally
+        current_job = prev_job
+    end
 end
 
 # primitive mechanism for deferred compilation, for implementing CUDA dynamic parallelism.

test/ptx.jl

@@ -10,6 +10,28 @@
     end
 end
 
+@testset "kernel state survives a runtime rebuild during optimization" begin
+    # Clearing the cache forces the relink pass inside `optimize!` to rebuild the runtime
+    # (nested compilation); the kernel must still get its state argument afterwards.
+    mod = @eval module $(gensym())
+        function kernel(x)
+            x < 1 && throw(DivideError())
+            return
+        end
+    end
+    old_cache = GPUCompiler.compile_cache
+    ir = try
+        GPUCompiler.compile_cache = nothing
+        sprint() do io
+            PTX.code_llvm(io, mod.kernel, Tuple{Int}; kernel=true, dump_module=true)
+        end
+    finally
+        GPUCompiler.compile_cache = old_cache
+    end
+    @test occursin("gpu_report_exception", ir)
+    @test occursin("[1 x i64] %state", ir)
+end
+
 @testset "kernel functions" begin
 @testset "kernel argument attributes" begin
     mod = @eval module $(gensym())