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