Internalize non-kernel deferred codegen entrypoints

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

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
 

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

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

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

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