From d81f2c55890a835e2411dd0701afc56c4e0066b3 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Wed, 9 Sep 2026 14:56:36 +0200 Subject: [PATCH] Internalize non-kernel deferred codegen entrypoints Deferred codegen entrypoints were kept externally visible across the InternalizePass, so that linking could resolve them. Once linked (and, for Enzyme's wrappers, alwaysinlined) they are dead, but their external linkage kept GlobalDCEPass from dropping them. Back-ends like SPIR-V then still had to translate functions they cannot express. Only deferred jobs that are kernels themselves (e.g. child kernels for dynamic parallelism) are entrypoints in their own right. Classify the entrypoints once after deferred codegen, and use that set both for internalization and for the per-entrypoint `finish_ir!` loop. Other deferred functions are internalized like any other function, and dropped by the existing clean-up passes once inlined. The mock Enzyme test helper gains an opt-in `always_inline` flag that marks the generated function `alwaysinline`, like Enzyme's wrappers, so the removal can be tested on the native, PTX and SPIR-V targets. Assisted-by: Claude Code (Fable 5.1) --- src/driver.jl | 13 ++++++++++--- test/helpers/enzyme.jl | 36 ++++++++++++++++++++++++------------ test/native.jl | 20 ++++++++++++++++++-- test/ptx.jl | 6 ++++-- test/spirv.jl | 33 +++++++++++++++++++++++++++++++++ test/utils.jl | 2 +- 6 files changed, 90 insertions(+), 20 deletions(-) diff --git a/src/driver.jl b/src/driver.jl index 43168a75..d23b8d1d 100644 --- a/src/driver.jl +++ b/src/driver.jl @@ -292,6 +292,13 @@ const __llvm_initialized = Ref(false) erase!(dyn_marker) end + # only deferred jobs that are kernels themselves (e.g. child kernels for dynamic + # parallelism) are entrypoints in their own right. other deferred functions (e.g. the + # wrappers Enzyme generates) are only called from within this module, so they should + # be internalized like any other function and dropped once inlined, rather than kept + # around for back-ends that cannot express their signatures (e.g. SPIR-V). + entrypoints = filter(((job′, _),) -> job′ === job || job′.config.kernel, jobs) + if job.config.toplevel && job.config.libraries # load the runtime outside of a timing block (because it recurses into the compiler) if !uses_julia_runtime(job) @@ -320,7 +327,7 @@ const __llvm_initialized = Ref(false) # mark everything internal except for entrypoints and any exported # global variables. this makes sure that the optimizer can, e.g., # rewrite function signatures. - preserved_gvs = collect(values(jobs)) + preserved_gvs = collect(values(entrypoints)) for gvar in globals(ir) if linkage(gvar) == LLVM.API.LLVMExternalLinkage push!(preserved_gvs, LLVM.name(gvar)) @@ -400,8 +407,8 @@ const __llvm_initialized = Ref(false) # during deferred code generation. instead, process the deferred jobs # here. entry = finish_ir!(job, ir, entry) - for (job′, fn′) in jobs - job′ == job && continue + for (job′, fn′) in entrypoints + job′ === job && continue finish_ir!(job′, ir, functions(ir)[fn′]) end end diff --git a/test/helpers/enzyme.jl b/test/helpers/enzyme.jl index 302c4fac..e21c2a73 100644 --- a/test/helpers/enzyme.jl +++ b/test/helpers/enzyme.jl @@ -1,6 +1,7 @@ module Enzyme using ..GPUCompiler +using LLVM struct EnzymeTarget{Target<:AbstractCompilerTarget} <: AbstractCompilerTarget target::Target @@ -20,13 +21,17 @@ GPUCompiler.dwarf_version(target::EnzymeTarget) = GPUCompiler.dwarf_version(targ abstract type AbstractEnzymeCompilerParams <: AbstractCompilerParams end struct EnzymeCompilerParams{Params<:AbstractCompilerParams} <: AbstractEnzymeCompilerParams params::Params + # mark the generated function `alwaysinline`, like the wrappers Enzyme generates + always_inline::Bool end struct PrimalCompilerParams <: AbstractEnzymeCompilerParams end -EnzymeCompilerParams() = EnzymeCompilerParams(PrimalCompilerParams()) +EnzymeCompilerParams(params=PrimalCompilerParams(); always_inline=false) = + EnzymeCompilerParams(params, always_inline) -GPUCompiler.nest_params(::EnzymeCompilerParams, other::AbstractCompilerParams) = EnzymeCompilerParams(other) +GPUCompiler.nest_params(params::EnzymeCompilerParams, other::AbstractCompilerParams) = + EnzymeCompilerParams(other; params.always_inline) function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeTarget}) config = job.config @@ -47,21 +52,28 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT # ??? entry_abi ) primal_job = CompilerJob(job.source, primal_config, job.world) - return GPUCompiler.compile_unhooked(output, primal_job) + ir, meta = GPUCompiler.compile_unhooked(output, primal_job) # Normally, Enzyme would run here and transform the output of the primal job. + if output === :llvm && job.config.params.always_inline + push!(function_attributes(meta.entry), EnumAttribute("alwaysinline", 0)) + end + + return ir, meta end import GPUCompiler: deferred_codegen_jobs import Core.Compiler as CC -function deferred_codegen_id_generator(world::UInt, source, self, ft::Type, tt::Type) +function deferred_codegen_id_generator(world::UInt, source, self, ft::Type, tt::Type, + always_inline::Type) @nospecialize @assert CC.isType(ft) && CC.isType(tt) ft = ft.parameters[1] tt = tt.parameters[1] + always_inline = always_inline.parameters[1]::Bool - stub = Core.GeneratedFunctionStub(identity, Core.svec(:deferred_codegen_id, :ft, :tt), Core.svec()) + stub = Core.GeneratedFunctionStub(identity, Core.svec(:deferred_codegen_id, :ft, :tt, :always_inline), Core.svec()) # look up the method match method_error = :(throw(MethodError(ft, tt, $world))) @@ -98,15 +110,15 @@ function deferred_codegen_id_generator(world::UInt, source, self, ft::Type, tt:: new_ci.edges = Any[mi] # prepare the slots - new_ci.slotnames = Symbol[Symbol("#self#"), :ft, :tt] - new_ci.slotflags = UInt8[0x00 for i = 1:3] + new_ci.slotnames = Symbol[Symbol("#self#"), :ft, :tt, :always_inline] + new_ci.slotflags = UInt8[0x00 for i = 1:4] @static if isdefined(Core, :DebugInfo) - new_ci.nargs = 3 + new_ci.nargs = 4 end # We don't know the caller's target so EnzymeTarget uses the default NativeCompilerTarget. target = EnzymeTarget() - params = EnzymeCompilerParams() + params = EnzymeCompilerParams(; always_inline) config = CompilerConfig(target, params; kernel=false) job = CompilerJob(mi, config, world) @@ -130,13 +142,13 @@ function deferred_codegen_id_generator(world::UInt, source, self, ft::Type, tt:: return new_ci end -@eval function deferred_codegen_id(ft, tt) +@eval function deferred_codegen_id(ft, tt, always_inline) $(Expr(:meta, :generated_only)) $(Expr(:meta, :generated, deferred_codegen_id_generator)) end -@inline function deferred_codegen(f::Type, tt::Type) - id = deferred_codegen_id(f, tt) +@inline function deferred_codegen(f::Type, tt::Type; always_inline::Bool=false) + id = deferred_codegen_id(f, tt, Val(always_inline)) ccall("extern deferred_codegen", llvmcall, Ptr{Cvoid}, (Int,), id) end diff --git a/test/native.jl b/test/native.jl index ceb780eb..bb94c50a 100644 --- a/test/native.jl +++ b/test/native.jl @@ -2030,9 +2030,25 @@ end return end - ir = sprint(io->Native.code_llvm(io, dkernel, Tuple{Vector{Float64}}; debuginfo=:none)) + ir = sprint(io->Native.code_llvm(io, dkernel, Tuple{Vector{Float64}}; + debuginfo=:none, dump_module=true)) @test !occursin("deferred_codegen", ir) - @test occursin("call void @julia_kernel", ir) + # the deferred function is only called from within the module, so it gets internalized + @test occursin(r"define internal .*@julia_kernel", ir) + @test occursin(r"call .*@julia_kernel", ir) + + # Enzyme's wrappers are `alwaysinline`, and should be dropped entirely once inlined + function dkernel_inline(a) + ptr = Enzyme.deferred_codegen(typeof(kernel), Tuple{Vector{Float64}}; + always_inline=true) + ccall(ptr, Cvoid, (Vector{Float64},), a) + return + end + + ir = sprint(io->Native.code_llvm(io, dkernel_inline, Tuple{Vector{Float64}}; + debuginfo=:none, dump_module=true)) + @test !occursin("deferred_codegen", ir) + @test !occursin("@julia_kernel", ir) end @testset "Mock Enzyme deferred relocations" begin diff --git a/test/ptx.jl b/test/ptx.jl index d6b0810a..6f018708 100644 --- a/test/ptx.jl +++ b/test/ptx.jl @@ -210,9 +210,11 @@ end return end - ir = sprint(io->Native.code_llvm(io, dkernel, Tuple{Ptr{Float64}}; debuginfo=:none)) + ir = sprint(io->PTX.code_llvm(io, dkernel, Tuple{Ptr{Float64}}; + debuginfo=:none, dump_module=true, kernel=true)) @test !occursin("deferred_codegen", ir) - @test occursin("call void @julia_", ir) + @test occursin(r"define internal .*@julia_kernel", ir) + @test occursin(r"call .*@julia_kernel", ir) end end diff --git a/test/spirv.jl b/test/spirv.jl index 912f72f9..4a4e3279 100644 --- a/test/spirv.jl +++ b/test/spirv.jl @@ -275,6 +275,39 @@ end end +@testset "deferred codegen" begin + # Enzyme's deferred wrappers are `alwaysinline` and take arguments the SPIR-V back-end + # cannot always express. once inlined they have to be dropped, rather than left behind + # as external definitions that still need to be translated. + mod = @eval module $(gensym()) + import ..Enzyme + function child(a) + unsafe_store!(a, unsafe_load(a)^2) + return + end + function kernel(a) + ptr = Enzyme.deferred_codegen(typeof(child), Tuple{Ptr{Float64}}; + always_inline=true) + ccall(ptr, Cvoid, (Ptr{Float64},), a) + return + end + end + + @test @filecheck begin + @check_not "@{{(julia|j)_child}}" + @check "define spir_kernel void @{{.*}}" + @check_not "@{{(julia|j)_child}}" + SPIRV.code_llvm(mod.kernel, Tuple{Ptr{Float64}}; backend, kernel=true, + dump_module=true) + end + + # `dump_module=true` disables `only_entry`, which would skip deferred codegen + asm = sprint(io->SPIRV.code_native(io, mod.kernel, Tuple{Ptr{Float64}}; + backend, kernel=true, dump_module=true)) + @test occursin("OpEntryPoint Kernel", asm) + @test count("= OpFunction ", asm) == 1 +end + @testset "replace i128 allocas" begin mod = @eval module $(gensym()) # reimplement some of SIMD.jl diff --git a/test/utils.jl b/test/utils.jl index 6477f02a..7dcecaf0 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -283,7 +283,7 @@ end @testset "Mock Enzyme" begin - Enzyme.deferred_codegen_id(typeof(identity), Tuple{Vector{Float64}}) + Enzyme.deferred_codegen_id(typeof(identity), Tuple{Vector{Float64}}, Val(false)) # Check that we can call this function from the CPU, to support deferred codegen for Enzyme. @test ccall("extern deferred_codegen", llvmcall, UInt, (UInt,), 3) == 3 end