Make GPUInterpreter extensible

Currently Enzyme uses it's own AbstractInterpreter, in particular to
handle inlining blocking of functions with custom rules and to handle
nested autodiff operations.

- [ ] Create a version of Enzyme with this
- [ ] Support a version of `gpuc.deferred(meta)`

Author: vchuravy

src/driver.jl

@@ -42,7 +42,7 @@ end
 ## deferred compilation
 
 """
-    var"gpuc.deferred"(f, args...)::Ptr{Cvoid}
+    var"gpuc.deferred"(meta, f, args...)::Ptr{Cvoid}
 
 As if we were to call `f(args...)` but instead we are
 putting down a marker and return a function pointer to later
@@ -154,10 +154,11 @@ const __llvm_initialized = Ref(false)
 
     @timeit_debug to "IR generation" begin
         ir, compiled = irgen(job)
+        edge = Edge(inference_metadata(job), job.source)
         if job.config.entry_abi === :specfunc
-            entry_fn = compiled[job.source].specfunc
+            entry_fn = compiled[edge].specfunc
         else
-            entry_fn = compiled[job.source].func
+            entry_fn = compiled[edge].func
         end
         entry = functions(ir)[entry_fn]
     end
@@ -198,24 +199,28 @@ const __llvm_initialized = Ref(false)
             return val
         end
 
-        worklist = Dict{Any, Vector{LLVM.CallInst}}()
+        worklist = Dict{Edge, Vector{LLVM.CallInst}}()
         for use in uses(dyn_marker)
             # decode the call
             call = user(use)::LLVM.CallInst
-            dyn_mi_inst = find_base_object(operands(call)[1])
+            dyn_meta_inst = find_base_object(operands(call)[1])
+            @compiler_assert isa(dyn_meta_inst, LLVM.ConstantInt) job
+            dyn_mi_inst = find_base_object(operands(call)[2])
             @compiler_assert isa(dyn_mi_inst, LLVM.ConstantInt) job
+            dyn_meta = Base.unsafe_pointer_to_objref(
+                convert(Ptr{Cvoid}, convert(Int, dyn_meta_inst)))
             dyn_mi = Base.unsafe_pointer_to_objref(
-                convert(Ptr{Cvoid}, convert(Int, dyn_mi_inst)))
-            push!(get!(worklist, dyn_mi, LLVM.CallInst[]), call)
+                convert(Ptr{Cvoid}, convert(Int, dyn_mi_inst)))::MethodInstance
+            push!(get!(worklist, Edge(dyn_meta, dyn_mi), LLVM.CallInst[]), call)
         end
 
-        for dyn_mi in keys(worklist)
-            dyn_fn_name = compiled[dyn_mi].specfunc
+        for dyn_edge in keys(worklist)
+            dyn_fn_name = compiled[dyn_edge].specfunc
             dyn_fn = functions(ir)[dyn_fn_name]
 
             # insert a pointer to the function everywhere the entry is used
             T_ptr = convert(LLVMType, Ptr{Cvoid})
-            for call in worklist[dyn_mi]
+            for call in worklist[dyn_edge]
                 @dispose builder=IRBuilder() begin
                     position!(builder, call)
                     fptr = if LLVM.version() >= v"17"

src/interface.jl

@@ -89,6 +89,7 @@ Several keyword arguments can be used to customize the compilation process:
 struct CompilerConfig{T,P}
     target::T
     params::P
+    meta
 
     kernel::Bool
     name::Union{Nothing,String}
@@ -98,6 +99,7 @@ struct CompilerConfig{T,P}
 
     function CompilerConfig(target::AbstractCompilerTarget,
                             params::AbstractCompilerParams;
+                            meta = nothing,
                             kernel=true,
                             name=nothing,
                             entry_abi=:specfunc,
@@ -106,16 +108,16 @@ struct CompilerConfig{T,P}
         if entry_abi ∉ (:specfunc, :func)
             error("Unknown entry_abi=$entry_abi")
         end
-        new{typeof(target), typeof(params)}(target, params, kernel, name, entry_abi,
+        new{typeof(target), typeof(params)}(target, params, meta, kernel, name, entry_abi,
                                             always_inline, opt_level)
     end
 end
 
 # copy constructor
-CompilerConfig(cfg::CompilerConfig; target=cfg.target, params=cfg.params,
+CompilerConfig(cfg::CompilerConfig; target=cfg.target, params=cfg.params, meta=cfg.meta, 
                kernel=cfg.kernel, name=cfg.name, entry_abi=cfg.entry_abi,
                always_inline=cfg.always_inline, opt_level=cfg.opt_level) =
-    CompilerConfig(target, params; kernel, entry_abi, name, always_inline, opt_level)
+    CompilerConfig(target, params; meta, kernel, entry_abi, name, always_inline, opt_level)
 
 function Base.show(io::IO, @nospecialize(cfg::CompilerConfig{T})) where {T}
     print(io, "CompilerConfig for ", T)
@@ -124,6 +126,7 @@ end
 function Base.hash(cfg::CompilerConfig, h::UInt)
     h = hash(cfg.target, h)
     h = hash(cfg.params, h)
+    h = hash(cfg.meta, h)::UInt
 
     h = hash(cfg.kernel, h)
     h = hash(cfg.name, h)
@@ -178,15 +181,17 @@ runtime_module(@nospecialize(job::CompilerJob)) = error("Not implemented")
 # check if a function is an intrinsic that can assumed to be always available
 isintrinsic(@nospecialize(job::CompilerJob), fn::String) = false
 
+inference_metadata(@nospecialize(job::CompilerJob)) = job.config.meta
+
 # provide a specific interpreter to use.
 if VERSION >= v"1.11.0-DEV.1552"
 get_interpreter(@nospecialize(job::CompilerJob)) =
-    GPUInterpreter(job.world; method_table=method_table(job),
+    GPUInterpreter(job.world; meta=inference_metadata(job), method_table=method_table(job),
                    token=ci_cache_token(job), inf_params=inference_params(job),
                    opt_params=optimization_params(job))
 else
 get_interpreter(@nospecialize(job::CompilerJob)) =
-    GPUInterpreter(job.world; method_table=method_table(job),
+    GPUInterpreter(job.world; meta=inference_metadata(job), method_table=method_table(job),
                    code_cache=ci_cache(job), inf_params=inference_params(job),
                    opt_params=optimization_params(job))
 end
@@ -227,10 +232,11 @@ struct GPUCompilerCacheToken
     target_type::Type
     always_inline::Bool
     method_table::Core.MethodTable
+    metadata
 end
 
 ci_cache_token(@nospecialize(job::CompilerJob)) =
-    GPUCompilerCacheToken(typeof(job.config.target), job.config.always_inline, method_table(job))
+    GPUCompilerCacheToken(typeof(job.config.target), job.config.always_inline, method_table(job), inference_metadata(job))
 
 # the codeinstance cache to use -- should only be used for the constructor
 if VERSION >= v"1.11.0-DEV.1552"

src/irgen.jl

@@ -2,10 +2,11 @@
 
 function irgen(@nospecialize(job::CompilerJob))
     mod, compiled = @timeit_debug to "emission" compile_method_instance(job)
+    edge = Edge(inference_metadata(job), job.source)
     if job.config.entry_abi === :specfunc
-        entry_fn = compiled[job.source].specfunc
+        entry_fn = compiled[edge].specfunc
     else
-        entry_fn = compiled[job.source].func
+        entry_fn = compiled[edge].func
     end
     @assert entry_fn !== nothing
     entry = functions(mod)[entry_fn]
@@ -70,25 +71,25 @@ function irgen(@nospecialize(job::CompilerJob))
         entry = deprecation_marker
     end
     if job.config.entry_abi === :specfunc
-        func = compiled[job.source].func
+        func = compiled[edge].func
         specfunc = LLVM.name(entry)
     else
         func = LLVM.name(entry)
-        specfunc = compiled[job.source].specfunc
+        specfunc = compiled[edge].specfunc
     end
 
-    compiled[job.source] =
-        (; compiled[job.source].ci, func, specfunc)
+    compiled[edge] =
+        (; compiled[edge].ci, func, specfunc)
 
     # Earlier we sanitize global names, this invalidates the
     # func, specfunc names safed in compiled. Update the names now,
     # such that when when use the compiled mappings to lookup the
     # llvm function for a methodinstance (deferred codegen) we have 
     # valid targets.
-    for mi in keys(compiled)
-        mi == job.source && continue
-        ci, func, specfunc = compiled[mi]
-        compiled[mi] = (; ci, func=safe_name(func), specfunc=safe_name(specfunc))
+    for other in keys(compiled)
+        other == edge && continue
+        ci, func, specfunc = compiled[other]
+        compiled[other] = (; ci, func=safe_name(func), specfunc=safe_name(specfunc))
     end
 
     # TODO: Should we rewrite gpuc.lookup here?
@@ -111,11 +112,11 @@ function irgen(@nospecialize(job::CompilerJob))
         # internalize all functions, but keep exported global variables.
         linkage!(entry, LLVM.API.LLVMExternalLinkage)
         preserved_gvs = String[LLVM.name(entry)]
-        for mi in keys(compiled)
+        for other in keys(compiled)
             # delay internalizing of deferred calls since
             # gpuc.lookup is not yet rewriten.
-            mi == job.source && continue
-            _, _, specfunc = compiled[mi]
+            other == edge && continue
+            _, _, specfunc = compiled[other]
             push!(preserved_gvs, specfunc) # this could be deleted if we rewrite gpuc.lookup earlier
         end
         for gvar in globals(mod)