Fix deferred_codegen registration

Co-authored-by: Gabriel Baraldi <[email protected]>

Author: vchuravy

src/GPUCompiler.jl

@@ -66,6 +66,28 @@ function __init__()
     global compile_cache = dir
 
     Tracy.@register_tracepoints()
+
+    # Register deferred_codegen as a global function so that it can be called with `ccall("extern deferred_codegen"`
+    @dispose jljit=JuliaOJIT() begin
+        jd = JITDylib(jljit)
+
+        address = LLVM.API.LLVMOrcJITTargetAddress(
+            reinterpret(UInt, @cfunction(deferred_codegen, Ptr{Cvoid}, (Ptr{Cvoid},))))
+        flags = LLVM.API.LLVMJITSymbolFlags(
+            LLVM.API.LLVMJITSymbolGenericFlagsExported, 0)
+        name = mangle(jljit, "deferred_codegen")
+        symbol = LLVM.API.LLVMJITEvaluatedSymbol(address, flags)
+        map = if LLVM.version() >= v"15"
+            LLVM.API.LLVMOrcCSymbolMapPair(name, symbol)
+        else
+            LLVM.API.LLVMJITCSymbolMapPair(name, symbol)
+        end
+
+        mu = LLVM.absolute_symbols(Ref(map))
+        LLVM.define(jd, mu)
+        addr = lookup(jljit, "deferred_codegen")
+        @assert addr != C_NULL "Failed to register deferred_codegen"
+    end
 end
 
 end # module

src/driver.jl

@@ -129,7 +129,8 @@ const deferred_codegen_jobs = Dict{Int, Any}()
 
 # We make this function explicitly callable so that we can drive OrcJIT's
 # lazy compilation from, while also enabling recursive compilation.
-Base.@ccallable Ptr{Cvoid} function deferred_codegen(ptr::Ptr{Cvoid})
+# Julia 1.11 and co broke @ccallable so we have to do this manually in __init__
+function deferred_codegen(ptr::Ptr{Cvoid})::Ptr{Cvoid}
     ptr
 end
 

test/utils.jl

@@ -190,4 +190,6 @@ end
 
 @testset "Mock Enzyme" begin
     Enzyme.deferred_codegen_id(typeof(identity), Tuple{Vector{Float64}})
+    # 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