Turn Julia passes into closures (#864)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: vchuravy

src/driver.jl

@@ -82,53 +82,44 @@ function compile_unhooked(output::Symbol, @nospecialize(job::CompilerJob); kwarg
         error("No active LLVM context. Use `JuliaContext()` do-block syntax to create one.")
     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
+    @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
-
-            return ir, ir_meta
+    if output == :llvm
+        if job.config.strip
+            @tracepoint "strip debug info" strip_debuginfo!(ir)
         end
 
+        return ir, ir_meta
+    end
 
-        ## 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)
 
-        if output == :asm || output == :obj
-            return asm, (; asm_meta..., ir_meta..., ir)
-        end
+    ## 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)
 
-        error("Unknown compilation output $output")
-    finally
-        current_job = prev_job
+    if output == :asm || output == :obj
+        return asm, (; asm_meta..., ir_meta..., ir)
     end
+
+
+    error("Unknown compilation output $output")
 end
 
 # primitive mechanism for deferred compilation, for implementing CUDA dynamic parallelism.

src/gcn.jl

@@ -44,7 +44,7 @@ pass_by_ref(@nospecialize(job::CompilerJob{GCNCompilerTarget})) = true
 
 function finish_module!(@nospecialize(job::CompilerJob{GCNCompilerTarget}),
                         mod::LLVM.Module, entry::LLVM.Function)
-    lower_throw_extra!(mod)
+    lower_throw_extra!(job, mod)
 
     if job.config.kernel
         # calling convention
@@ -156,8 +156,7 @@ end
 
 ## LLVM passes
 
-function lower_throw_extra!(mod::LLVM.Module)
-    job = current_job::CompilerJob
+function lower_throw_extra!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
     changed = false
     @tracepoint "lower throw (extra)" begin
 
@@ -179,7 +178,7 @@ function lower_throw_extra!(mod::LLVM.Module)
                     # replace the throw with a trap
                     @dispose builder=IRBuilder() begin
                         position!(builder, call)
-                        emit_exception!(builder, f_name, call)
+                        emit_exception!(job, builder, f_name, call)
                     end
 
                     # remove the call

src/irgen.jl

@@ -143,9 +143,7 @@ function irgen(@nospecialize(job::CompilerJob))
             end
         end
 
-        global current_job
-        current_job = job
-        can_throw(job) || lower_throw!(mod)
+        can_throw(job) || lower_throw!(job, mod)
 
         # resolve the `julia.gpu.debug_level` intrinsic (see `kernel_debug_level_value`) to
         # the job's configured level, so device code can branch on it as a compile-time
@@ -166,8 +164,7 @@ end
 #
 # once we have thorough inference (ie. discarding `@nospecialize` and thus supporting
 # exception arguments) and proper debug info to unwind the stack, this pass can go.
-function lower_throw!(mod::LLVM.Module)
-    job = current_job::CompilerJob
+function lower_throw!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
     changed = false
     @tracepoint "lower throw" begin
 
@@ -200,7 +197,7 @@ function lower_throw!(mod::LLVM.Module)
                 # replace the throw with a PTX-compatible exception
                 @dispose builder=IRBuilder() begin
                     position!(builder, call)
-                    emit_exception!(builder, name, call)
+                    emit_exception!(job, builder, name, call)
                 end
 
                 # remove the call
@@ -249,8 +246,7 @@ end
 # while `!mayThrow() && willReturn()`) from deleting the `signal_exception` call below and
 # folding away the guarding bounds-check branch. so the trap is the optimizer-correctness guard;
 # do not move its removal earlier than post-`optimize!`.
-function emit_exception!(builder, name, inst)
-    job = current_job::CompilerJob
+function emit_exception!(@nospecialize(job::CompilerJob), builder, name, inst)
     bb = position(builder)
     fun = LLVM.parent(bb)
     mod = LLVM.parent(fun)
@@ -822,12 +818,13 @@ end
 # so that the julia.gpu.state_getter` can be simplified to return an opaque pointer.
 
 # add a state argument to every function in the module, starting from the kernel entry point
-function add_kernel_state!(mod::LLVM.Module)
-    job = current_job::CompilerJob
-
+struct AddKernelState
+    job::CompilerJob
+end
+function (self::AddKernelState)(mod::LLVM.Module)
     # check if we even need a kernel state argument
-    job.config.kernel || return false
-    state = kernel_state_type(job)
+    self.job.config.kernel || return false
+    state = kernel_state_type(self.job)
     if state === Nothing
         return false
     end
@@ -1022,18 +1019,20 @@ function add_kernel_state!(mod::LLVM.Module)
 
     return true
 end
-AddKernelStatePass() = NewPMModulePass("AddKernelStatePass", add_kernel_state!)
+AddKernelStatePass(job) = NewPMModulePass("AddKernelStatePass", AddKernelState(job))
 
 # lower calls to the state getter intrinsic. this is a two-step process, so that the state
 # argument can be added before optimization, and that optimization can introduce new uses
 # before the intrinsic getting lowered late during optimization.
-function lower_kernel_state!(fun::LLVM.Function)
-    job = current_job::CompilerJob
+struct LowerKernelState
+    job::CompilerJob
+end
+function (self::LowerKernelState)(fun::LLVM.Function)
     mod = LLVM.parent(fun)
     changed = false
 
     # check if we even need a kernel state argument
-    state = kernel_state_type(job)
+    state = kernel_state_type(self.job)
     if state === Nothing
         return false
     end
@@ -1084,10 +1083,12 @@ function lower_kernel_state!(fun::LLVM.Function)
 
     return changed
 end
-LowerKernelStatePass() = NewPMFunctionPass("LowerKernelStatePass", lower_kernel_state!)
+LowerKernelStatePass(job) = NewPMFunctionPass("LowerKernelStatePass", LowerKernelState(job))
 
-function cleanup_kernel_state!(mod::LLVM.Module)
-    job = current_job::CompilerJob
+struct CleanupKernelState
+    job::CompilerJob
+end
+function (self::CleanupKernelState)(mod::LLVM.Module)
     changed = false
 
     # remove the getter intrinsic
@@ -1102,7 +1103,7 @@ function cleanup_kernel_state!(mod::LLVM.Module)
 
     return changed
 end
-CleanupKernelStatePass() = NewPMModulePass("CleanupKernelStatePass", cleanup_kernel_state!)
+CleanupKernelStatePass(job) = NewPMModulePass("CleanupKernelStatePass", CleanupKernelState(job))
 
 function kernel_state_intr(mod::LLVM.Module, T_state)
     state_intr = if haskey(functions(mod), "julia.gpu.state_getter")
@@ -1147,7 +1148,7 @@ end
 
 # device code can query the job's configured debug level as a compile-time constant via
 # `kernel_debug_level()`, which emits the `julia.gpu.debug_level` intrinsic; `lower_debug_level!`
-# (run from `irgen`, with `current_job` in scope) replaces it with the constant. this keeps
+# (run from `irgen`, with `job` in scope) replaces it with the constant. this keeps
 # the level part of the cache key (it lives in `CompilerConfig`), unlike reading the `-g`
 # global at parse time (which would bake the wrong level under pkgimage reuse across `-g`).
 

src/mcgen.jl

@@ -3,15 +3,12 @@
 # final preparations for the module to be compiled to machine code
 # these passes should not be run when e.g. compiling to write to disk.
 function prepare_execution!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
-    global current_job
-    current_job = job
-
     @dispose pb=NewPMPassBuilder() begin
-        register!(pb, ResolveCPUReferencesPass())
+        register!(pb, ResolveCPUReferencesPass(job))
 
         add!(pb, RecomputeGlobalsAAPass())
         add!(pb, GlobalOptPass())
-        add!(pb, ResolveCPUReferencesPass())
+        add!(pb, ResolveCPUReferencesPass(job))
         add!(pb, GlobalDCEPass())
         add!(pb, StripDeadPrototypesPass())
 
@@ -29,8 +26,10 @@ end
 # but at the same time the GPU can't resolve them at run-time.
 #
 # this pass performs that resolution at link time.
-function resolve_cpu_references!(mod::LLVM.Module)
-    job = current_job::CompilerJob
+struct ResolveCPUReferences
+    job::CompilerJob
+end
+function (self::ResolveCPUReferences)(mod::LLVM.Module)
     changed = false
 
     for f in functions(mod)
@@ -65,8 +64,8 @@ function resolve_cpu_references!(mod::LLVM.Module)
 
     return changed
 end
-ResolveCPUReferencesPass() =
-    NewPMModulePass("ResolveCPUReferences", resolve_cpu_references!)
+ResolveCPUReferencesPass(job) =
+    NewPMModulePass("ResolveCPUReferences", ResolveCPUReferences(job))
 
 
 function mcgen(@nospecialize(job::CompilerJob), mod::LLVM.Module, format=LLVM.API.LLVMAssemblyFile)