From 759a52abb8092b4b04f2538e512e82e986adf871 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Fri, 26 Jun 2026 17:30:04 +0200 Subject: [PATCH 1/2] Turn Julia passes into closures --- src/optim.jl | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/optim.jl b/src/optim.jl index 1245ba1c..154877c6 100644 --- a/src/optim.jl +++ b/src/optim.jl @@ -59,7 +59,7 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module; opt_level= @dispose pb=NewPMPassBuilder() begin tti === nothing || LLVM.target_transform_info!(pb, tti) - register!(pb, GPULowerCPUFeaturesPass()) + register!(pb, GPULowerCPUFeaturesPass(job)) register!(pb, GPULowerPTLSPass()) register!(pb, GPULowerGCFramePass()) register!(pb, GPULinkRuntimePass()) @@ -168,7 +168,7 @@ function buildEarlyOptimizerPipeline(mpm, @nospecialize(job::CompilerJob), opt_l end end end - add!(mpm, GPULowerCPUFeaturesPass()) + add!(mpm, GPULowerCPUFeaturesPass(job)) if opt_level >= 1 add!(mpm, NewPMFunctionPassManager()) do fpm if opt_level >= 2 @@ -425,8 +425,10 @@ end ## custom passes # lowering intrinsics -function cpu_features!(mod::LLVM.Module) - job = current_job::CompilerJob +struct CPUFeatures + current_job::CompilerJob +end +function (self::CPUFeatures)(mod::LLVM.Module) changed = false argtyps = Dict( @@ -470,7 +472,7 @@ function cpu_features!(mod::LLVM.Module) return changed end -GPULowerCPUFeaturesPass() = NewPMModulePass("GPULowerCPUFeatures", cpu_features!) +GPULowerCPUFeaturesPass(current_job) = NewPMModulePass("GPULowerCPUFeatures", CPUFeatures(current_job)) function link_runtime!(mod::LLVM.Module) job = current_job::CompilerJob From ecf6feed34d2efe5372535611828183c7b1016d2 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 28 Jun 2026 10:18:28 +0200 Subject: [PATCH 2/2] Remove current_job. --- src/driver.jl | 61 +++++++++++++----------------- src/gcn.jl | 7 ++-- src/irgen.jl | 45 +++++++++++----------- src/mcgen.jl | 17 ++++----- src/optim.jl | 103 +++++++++++++++++++++++++++----------------------- src/spirv.jl | 6 +-- 6 files changed, 117 insertions(+), 122 deletions(-) diff --git a/src/driver.jl b/src/driver.jl index ec4ea9f8..30529c70 100644 --- a/src/driver.jl +++ b/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. diff --git a/src/gcn.jl b/src/gcn.jl index 6edd8b79..e32b3c21 100644 --- a/src/gcn.jl +++ b/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 diff --git a/src/irgen.jl b/src/irgen.jl index 65d3f729..80dac84d 100644 --- a/src/irgen.jl +++ b/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`). diff --git a/src/mcgen.jl b/src/mcgen.jl index 77a40d85..2c5e85c3 100644 --- a/src/mcgen.jl +++ b/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) diff --git a/src/optim.jl b/src/optim.jl index 154877c6..7a07d200 100644 --- a/src/optim.jl +++ b/src/optim.jl @@ -53,21 +53,18 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module; opt_level= tm = llvm_machine(job.config.target) tti = llvm_targetinfo(job.config.target) - global current_job - current_job = job - @dispose pb=NewPMPassBuilder() begin tti === nothing || LLVM.target_transform_info!(pb, tti) register!(pb, GPULowerCPUFeaturesPass(job)) - register!(pb, GPULowerPTLSPass()) - register!(pb, GPULowerGCFramePass()) - register!(pb, GPULinkRuntimePass()) - register!(pb, GPULinkLibrariesPass()) - register!(pb, GPUFinishRuntimeIntrinsicsPass()) - register!(pb, AddKernelStatePass()) - register!(pb, LowerKernelStatePass()) - register!(pb, CleanupKernelStatePass()) + register!(pb, GPULowerPTLSPass(job)) + register!(pb, GPULowerGCFramePass(job)) + register!(pb, GPULinkRuntimePass(job)) + register!(pb, GPULinkLibrariesPass(job)) + register!(pb, GPUFinishRuntimeIntrinsicsPass(job)) + register!(pb, AddKernelStatePass(job)) + register!(pb, LowerKernelStatePass(job)) + register!(pb, CleanupKernelStatePass(job)) add!(pb, NewPMModulePassManager()) do mpm buildNewPMPipeline!(mpm, job, opt_level) @@ -325,12 +322,12 @@ function buildIntrinsicLoweringPipeline(mpm, @nospecialize(job::CompilerJob), op # lower GC intrinsics if !uses_julia_runtime(job) add!(mpm, NewPMFunctionPassManager()) do fpm - add!(fpm, GPULowerGCFramePass()) + add!(fpm, GPULowerGCFramePass(job)) end if job.config.libraries - add!(mpm, GPULinkRuntimePass()) - add!(mpm, GPULinkLibrariesPass()) - add!(mpm, GPUFinishRuntimeIntrinsicsPass()) + add!(mpm, GPULinkRuntimePass(job)) + add!(mpm, GPULinkLibrariesPass(job)) + add!(mpm, GPUFinishRuntimeIntrinsicsPass(job)) end end @@ -339,11 +336,11 @@ function buildIntrinsicLoweringPipeline(mpm, @nospecialize(job::CompilerJob), op # and thus additional uses of the kernel state intrinsics. if job.config.kernel # TODO: now that all kernel state-related passes are being run here, merge some? - add!(mpm, AddKernelStatePass()) + add!(mpm, AddKernelStatePass(job)) add!(mpm, NewPMFunctionPassManager()) do fpm - add!(fpm, LowerKernelStatePass()) + add!(fpm, LowerKernelStatePass(job)) end - add!(mpm, CleanupKernelStatePass()) + add!(mpm, CleanupKernelStatePass(job)) end if !uses_julia_runtime(job) @@ -351,7 +348,7 @@ function buildIntrinsicLoweringPipeline(mpm, @nospecialize(job::CompilerJob), op add!(mpm, NewPMFunctionPassManager()) do fpm add!(fpm, ADCEPass()) end - add!(mpm, GPULowerPTLSPass()) + add!(mpm, GPULowerPTLSPass(job)) end add!(mpm, NewPMFunctionPassManager()) do fpm @@ -426,7 +423,7 @@ end # lowering intrinsics struct CPUFeatures - current_job::CompilerJob + job::CompilerJob end function (self::CPUFeatures)(mod::LLVM.Module) changed = false @@ -446,7 +443,7 @@ function (self::CPUFeatures)(mod::LLVM.Module) # determine whether this back-end supports FMA on this type has_fma = if haskey(argtyps, typnam) typ = argtyps[typnam] - have_fma(job.config.target, typ) + have_fma(self.job.config.target, typ) else # warn? false @@ -472,17 +469,19 @@ function (self::CPUFeatures)(mod::LLVM.Module) return changed end -GPULowerCPUFeaturesPass(current_job) = NewPMModulePass("GPULowerCPUFeatures", CPUFeatures(current_job)) +GPULowerCPUFeaturesPass(job) = NewPMModulePass("GPULowerCPUFeatures", CPUFeatures(job)) -function link_runtime!(mod::LLVM.Module) - job = current_job::CompilerJob - job.config.libraries || return false - uses_julia_runtime(job) && return false +struct LinkRuntime + job::CompilerJob +end +function (self::LinkRuntime)(mod::LLVM.Module) + self.job.config.libraries || return false + uses_julia_runtime(self.job) && return false # GC lowering can introduce new calls to GPU runtime functions after the runtime # was linked initially. Link again now so those calls resolve to definitions before # later intrinsic-lowering passes inspect or rewrite the runtime call graph. - runtime = load_runtime(job) + runtime = load_runtime(self.job) # `RemoveNIPass` stripped non-integral address spaces from `mod`'s datalayout, but the # cached runtime kept them; align it (as with target libraries) to avoid a warning. triple!(runtime, triple(mod)) @@ -490,33 +489,37 @@ function link_runtime!(mod::LLVM.Module) link!(mod, runtime; only_needed=true) return true end -GPULinkRuntimePass() = NewPMModulePass("GPULinkRuntime", link_runtime!) +GPULinkRuntimePass(job) = NewPMModulePass("GPULinkRuntime", LinkRuntime(job)) -function link_libraries!(mod::LLVM.Module) - job = current_job::CompilerJob - job.config.libraries || return false +struct LinkLibraries + job::CompilerJob +end +function (self::LinkLibraries)(mod::LLVM.Module) + self.job.config.libraries || return false # Runtime functions materialized by GC lowering can introduce target-library # references after the normal library-linking phase has run. Give back-ends a # second chance to resolve those references before they lower their own runtime # intrinsics. This is a no-op for back-ends without a link_libraries! override. - if has_legacy_link_libraries(job) + if has_legacy_link_libraries(self.job) undefined_fns = [LLVM.name(f) for f in functions(mod) if isdeclaration(f) && !LLVM.isintrinsic(f)] - link_libraries!(job, mod, undefined_fns) + link_libraries!(self.job, mod, undefined_fns) else - link_libraries!(job, mod) + link_libraries!(self.job, mod) end return true end -GPULinkLibrariesPass() = NewPMModulePass("GPULinkLibraries", link_libraries!) +GPULinkLibrariesPass(job) = NewPMModulePass("GPULinkLibraries", LinkLibraries(job)) -function finish_runtime_intrinsics!(mod::LLVM.Module) - job = current_job::CompilerJob - return finish_runtime_intrinsics!(job, mod) +struct FinishRuntimeIntrinsics + job::CompilerJob end -GPUFinishRuntimeIntrinsicsPass() = - NewPMModulePass("GPUFinishRuntimeIntrinsics", finish_runtime_intrinsics!) +function (self::FinishRuntimeIntrinsics)(mod::LLVM.Module) + return finish_runtime_intrinsics!(self.job, mod) +end +GPUFinishRuntimeIntrinsicsPass(job) = + NewPMModulePass("GPUFinishRuntimeIntrinsics", FinishRuntimeIntrinsics(job)) # lower object allocations to to PTX malloc # @@ -526,8 +529,10 @@ GPUFinishRuntimeIntrinsicsPass() = # is currently very architecture/CPU specific: hard-coded pool sizes, TLS references, etc. # such IR is hard to clean-up, so we probably will need to have the GC lowering pass emit # lower-level intrinsics which then can be lowered to architecture-specific code. -function lower_gc_frame!(fun::LLVM.Function) - job = current_job::CompilerJob +struct LowerGCFrame + job::CompilerJob +end +function (self::LowerGCFrame)(fun::LLVM.Function) mod = LLVM.parent(fun) changed = false @@ -557,7 +562,7 @@ function lower_gc_frame!(fun::LLVM.Function) changed = true end - @compiler_assert isempty(uses(alloc_obj)) job + @compiler_assert isempty(uses(alloc_obj)) self.job end # we don't care about write barriers @@ -570,12 +575,12 @@ function lower_gc_frame!(fun::LLVM.Function) changed = true end - @compiler_assert isempty(uses(barrier)) job + @compiler_assert isempty(uses(barrier)) self.job end return changed end -GPULowerGCFramePass() = NewPMFunctionPass("GPULowerGCFrame", lower_gc_frame!) +GPULowerGCFramePass(job) = NewPMFunctionPass("GPULowerGCFrame", LowerGCFrame(job)) # lower the `julia.ptls_states` intrinsic by removing it, since it is GPU incompatible. # @@ -584,8 +589,10 @@ GPULowerGCFramePass() = NewPMFunctionPass("GPULowerGCFrame", lower_gc_frame!) # # TODO: maybe don't have Julia emit actual uses of the TLS, but use intrinsics instead, # making it easier to remove or reimplement that functionality here. -function lower_ptls!(mod::LLVM.Module) - job = current_job::CompilerJob +struct LowerPTLS + job::CompilerJob +end +function (self::LowerPTLS)(mod::LLVM.Module) changed = false intrinsic = "julia.get_pgcstack" @@ -606,4 +613,4 @@ function lower_ptls!(mod::LLVM.Module) return changed end -GPULowerPTLSPass() = NewPMModulePass("GPULowerPTLS", lower_ptls!) +GPULowerPTLSPass(job) = NewPMModulePass("GPULowerPTLS", LowerPTLS(job)) diff --git a/src/spirv.jl b/src/spirv.jl index da9306a6..3400fef6 100644 --- a/src/spirv.jl +++ b/src/spirv.jl @@ -144,7 +144,7 @@ end # the LLVM to SPIR-V translator does not support the freeze instruction # (SPIRV-LLVM-Translator#1140) - rm_freeze!(mod) + rm_freeze!(job, mod) # translate to SPIR-V input = tempname(cleanup=false) * ".bc" @@ -249,8 +249,7 @@ end # remove freeze and replace uses by the original value # (KhronosGroup/SPIRV-LLVM-Translator#1140) -function rm_freeze!(mod::LLVM.Module) - job = current_job::CompilerJob +function rm_freeze!(@nospecialize(job::CompilerJob), mod::LLVM.Module) changed = false @tracepoint "remove freeze" begin @@ -271,7 +270,6 @@ end # convert alloca [N x i128] to alloca [N x <2 x i64>] # SPIR-V doesn't support i128 types, but we can represent them as vectors function convert_i128_allocas!(mod::LLVM.Module) - job = current_job::CompilerJob changed = false @tracepoint "convert i128 allocas" begin