Skip to content

Call primals of functions with custom rules natively #3783

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
vchuravy wants to merge 2 commits into from
Draft

Call primals of functions with custom rules natively #3783

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f0b81b1
Call primals of functions with custom rules natively
claude Oct 8, 2026
f4dc7b9
Keep a ruled call on jl_invoke when compileable_specialization widene…
vchuravy Oct 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions src/Enzyme.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,9 +148,11 @@ global without the prefix; see `insert_julia_value!`). On GPUCompiler 2.x the
values of the slots are read straight from the relocation records of the
modules the compilation emitted (`relocations`), and `bakes` says whether the
back-end lets their host addresses into the code. `native_callees` holds the
derivative-free callees `ci_cache_populate` left out of the module, for
`bind_native_callees!` to call natively, keyed by the CodeInstance or
MethodInstance codegen invokes.
callees `ci_cache_populate` left out of the module, for `bind_native_callees!` to call
natively, keyed by the CodeInstance or MethodInstance codegen invokes: their native
CodeInstance, its specialized entry point, and how Enzyme handles their calls
(`:inactive` for a derivative-free callee, `:frule` or `:rrule` for one with a custom
rule).

A context belongs to a single `compile_unhooked` invocation, and holds the
world age that compilation runs in. It is not passed as an argument:
Expand All @@ -172,7 +174,7 @@ mutable struct EnzymeContext
inserted_values::Dict{String, Any}
relocations::Vector{Any}
bakes::Bool
native_callees::IdDict{Any, Tuple{Core.CodeInstance, Ptr{Cvoid}}}
native_callees::IdDict{Any, Tuple{Core.CodeInstance, Ptr{Cvoid}, Symbol}}
EnzymeContext(world::Integer) = new(
world,
LLVM.Module[],
Expand All @@ -184,7 +186,7 @@ mutable struct EnzymeContext
Dict{String, Any}(),
Any[],
true,
IdDict{Any, Tuple{Core.CodeInstance, Ptr{Cvoid}}}()
IdDict{Any, Tuple{Core.CodeInstance, Ptr{Cvoid}, Symbol}}()
)
end

Expand Down
120 changes: 65 additions & 55 deletions src/compiler.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -628,69 +628,79 @@ function nested_codegen!(mode::API.CDerivativeMode, mod::LLVM.Module, @nospecial
end


function prepare_llvm(interp, mod::LLVM.Module, job, meta, enzyme_ctx::EnzymeContext)
for (mi, k) in meta.compiled
k_name = GPUCompiler.safe_name(k.specfunc)
if !haskey(functions(mod), k_name)
continue
end
llvmfn = functions(mod)[k_name]
"""
mark_compiled!(mod::LLVM.Module, llvmfn::LLVM.Function, mi::MethodInstance, RT::Type, world::UInt)

RT = return_type(interp, mi)
Give `llvmfn`, the function Julia's codegen emitted into `mod` for `mi` with return type
`RT`, the attributes Enzyme reads from a compiled Julia function: `enzymejl_mi` and
`enzymejl_rt`, the read-only-or-throw of a function with an easy rule, and the markers of
an sret union and of the GC roots of the return, on the function and its call sites.
"""
function mark_compiled!(mod::LLVM.Module, llvmfn::LLVM.Function, mi::Core.MethodInstance, @nospecialize(RT), world::UInt)
_, _, returnRoots0 = get_return_info(RT)
returnRoots = returnRoots0 !== nothing

_, _, returnRoots0 = get_return_info(RT)
returnRoots = returnRoots0 !== nothing
attributes = function_attributes(llvmfn)
# A function that already carries `enzymejl_mi` is not this emission's
# output: a derivative embedded for nested differentiation may reuse
# the name of a fresh function, since Julia's function name counters
# restart per compilation.
fresh = !has_fn_attr(llvmfn, StringAttribute("enzymejl_mi"))
push!(
attributes,
StringAttribute("enzymejl_mi", string(convert(UInt, pointer_from_objref(mi)))),
)
push!(
attributes,
StringAttribute("enzymejl_rt", string(convert(UInt, unsafe_to_pointer(RT)))),
)
if cached_has_easy_rule(Interpreter.simplify_kw(mi.specTypes), world)
push!(attributes, LLVM.StringAttribute("enzyme_LocalReadOnlyOrThrow"))
# The rule accesses no more than the body does, so Enzyme may keep
# inferring the function's full parameter attributes from the body.
push!(attributes, LLVM.StringAttribute("enzyme_custom_full_attributes"))
end

attributes = function_attributes(llvmfn)
# A function that already carries `enzymejl_mi` is not this emission's
# output: a derivative embedded for nested differentiation may reuse
# the name of a fresh function, since Julia's function name counters
# restart per compilation.
fresh = !has_fn_attr(llvmfn, StringAttribute("enzymejl_mi"))
push!(
attributes,
StringAttribute("enzymejl_mi", string(convert(UInt, pointer_from_objref(mi)))),
)
push!(
attributes,
StringAttribute("enzymejl_rt", string(convert(UInt, unsafe_to_pointer(RT)))),
)
if cached_has_easy_rule(Interpreter.simplify_kw(mi.specTypes), job.world)
push!(attributes, LLVM.StringAttribute("enzyme_LocalReadOnlyOrThrow"))
# The rule accesses no more than the body does, so Enzyme may keep
# inferring the function's full parameter attributes from the body.
push!(attributes, LLVM.StringAttribute("enzyme_custom_full_attributes"))
end
if startswith(LLVM.name(llvmfn), "japi3") || startswith(LLVM.name(llvmfn), "japi1") || startswith(LLVM.name(llvmfn), "jlcapi")
return nothing
end

if startswith(LLVM.name(llvmfn), "japi3") || startswith(LLVM.name(llvmfn), "japi1") || startswith(LLVM.name(llvmfn), "jlcapi")
continue
end
if fresh
check_emitted_specsig(mod, llvmfn, mi, RT)
end

if fresh
check_emitted_specsig(mod, llvmfn, mi, RT)
if is_sret_union(RT)
attr = StringAttribute("enzymejl_sret_union_bytes", string(union_alloca_type(RT)))
push!(parameter_attributes(llvmfn, 1), attr)
for u in LLVM.uses(llvmfn)
u = LLVM.user(u)
@assert isa(u, LLVM.CallInst)
LLVM.API.LLVMAddCallSiteAttribute(u, LLVM.API.LLVMAttributeIndex(1), attr)
end
end

if is_sret_union(RT)
attr = StringAttribute("enzymejl_sret_union_bytes", string(union_alloca_type(RT)))
push!(parameter_attributes(llvmfn, 1), attr)
for u in LLVM.uses(llvmfn)
u = LLVM.user(u)
@assert isa(u, LLVM.CallInst)
LLVM.API.LLVMAddCallSiteAttribute(u, LLVM.API.LLVMAttributeIndex(1), attr)
end
if returnRoots
attr = StringAttribute("enzymejl_returnRoots", string(length(eltype(returnRoots0).parameters[1])))
push!(parameter_attributes(llvmfn, 2), attr)
for u in LLVM.uses(llvmfn)
u = LLVM.user(u)
@assert isa(u, LLVM.CallInst)
LLVM.API.LLVMAddCallSiteAttribute(u, LLVM.API.LLVMAttributeIndex(2), attr)
end
end

if returnRoots
attr = StringAttribute("enzymejl_returnRoots", string(length(eltype(returnRoots0).parameters[1])))
push!(parameter_attributes(llvmfn, 2), attr)
for u in LLVM.uses(llvmfn)
u = LLVM.user(u)
@assert isa(u, LLVM.CallInst)
LLVM.API.LLVMAddCallSiteAttribute(u, LLVM.API.LLVMAttributeIndex(2), attr)
end
end
fixup_1p12_sret!(llvmfn)
return nothing
end

fixup_1p12_sret!(llvmfn)
function prepare_llvm(interp, mod::LLVM.Module, job, meta, enzyme_ctx::EnzymeContext)
for (mi, k) in meta.compiled
k_name = GPUCompiler.safe_name(k.specfunc)
if !haskey(functions(mod), k_name)
continue
end
llvmfn = functions(mod)[k_name]
mark_compiled!(mod, llvmfn, mi, return_type(interp, mi), job.world)
end

# We explicitly save the type of alloca's before they get lowered
Expand Down Expand Up @@ -6522,7 +6532,7 @@ function compile_unhooked_impl(output::Symbol, job::CompilerJob{<:EnzymeTarget})

# A derivative linked into this module as a deferred job calls its rules
# natively. Differentiating it again needs their bodies.
materialize_native_invokes!(mode, mod)
materialize_native_invokes!(mode, mod, Core.Compiler.method_table(primal_interp))

LLVM.@dispose pb=LLVM.NewPMPassBuilder() begin
registerEnzymeAndPassPipeline!(pb)
Expand All @@ -6543,7 +6553,7 @@ function compile_unhooked_impl(output::Symbol, job::CompilerJob{<:EnzymeTarget})
end
# `check_ir` links in the derivatives that `enzyme_call` embeds. Their
# natively called rules need bodies too.
materialize_native_invokes!(mode, mod)
materialize_native_invokes!(mode, mod, Core.Compiler.method_table(primal_interp))

disableFallback = String[]

Expand Down
61 changes: 59 additions & 2 deletions src/compiler/callconv.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -625,7 +625,7 @@ function invoke_codegen!(
end

"""
materialize_native_invokes!(mode, mod)
materialize_native_invokes!(mode, mod, method_table)

Give every natively called function that `mod` declares a body, so that
the differentiation of `mod` can differentiate through it.
Expand All @@ -639,17 +639,28 @@ is an opaque call to the outer differentiation. This emits the rule with
`nested_codegen!`, as an `:inline` rule, and defines the declaration as an
always-inline wrapper that forwards to it. The wrapper drops the
`pgcstack` parameter when the emitted rule takes none.

Neither a derivative-free callee nor one the differentiation of `mod` has a
custom rule for (see `bind_native_callees!`) needs a body: `mod` calls the
rule, or the declaration as is. A function with a custom rule only for
another mode, such as the primal of a reverse rule in a derivative that is
differentiated in forward mode, gets one like a rule, and loses the marks
the inner differentiation gave it for the rule handlers (see
[`drop_stale_custom_marks!`](@ref)).
"""
function materialize_native_invokes!(mode::API.CDerivativeMode, mod::LLVM.Module)
function materialize_native_invokes!(mode::API.CDerivativeMode, mod::LLVM.Module, method_table)
enzyme_ctx = enzyme_context()
world = enzyme_ctx.world
drop_stale_custom_marks!(mode, mod, world, method_table)
marker = StringAttribute("enzymejl_native_invoke")
for fn in collect(functions(mod))
isdeclaration(fn) || continue
has_fn_attr(fn, marker) || continue
# A callee Enzyme never differentiates needs no body (see `bind_native_callees!`).
has_fn_attr(fn, StringAttribute("enzymejl_native_inactive")) && continue
mi, RT = enzyme_custom_extract_mi(fn)
# Nor does one the rule handlers handle (see `handle_compiled`).
has_mode_rule(mode, mi, world, method_table) && continue
llvmf = nested_codegen!(mode, mod, mi, true)
check_specsig(llvmf, mi, enzyme_custom_extract_mi(llvmf)[2])
# `nested_codegen!` defers linking the emitted module until after
Expand Down Expand Up @@ -697,3 +708,49 @@ function materialize_native_invokes!(mode::API.CDerivativeMode, mod::LLVM.Module
end
return nothing
end

"""
has_mode_rule(mode, mi, world, method_table) -> Bool

Say if differentiating in `mode` applies a custom rule to calls of `mi`: a forward
rule in forward mode and a reverse rule otherwise, as `handle_compiled` decides.
"""
function has_mode_rule(mode::API.CDerivativeMode, mi::Core.MethodInstance, world::UInt, method_table)::Bool
specTypes = Interpreter.simplify_kw(mi.specTypes)
if mode == API.DEM_ForwardMode
return cached_has_frule(specTypes, world, method_table)
else
return cached_has_rrule(specTypes, world, method_table)
end
end

"""
drop_stale_custom_marks!(mode, mod, world, method_table)

Remove the marks `handle_compiled` gives a function with a custom rule (`enzyme_math`
set to `enzyme_custom`, and the preserved primal) from every function of `mod` that the
differentiation of `mod` in `mode` has no rule for. Only a derivative linked into `mod`
for nested differentiation has such marks yet: its own differentiation, in another mode,
had a rule for the function. The differentiation of `mod` must differentiate through
the function instead, and `handle_compiled` marks again the functions it does have a
rule for.
"""
function drop_stale_custom_marks!(mode::API.CDerivativeMode, mod::LLVM.Module, world::UInt, method_table)
for fn in functions(mod)
is_custom = false
for attr in collect(function_attributes(fn))
if attr isa StringAttribute && kind(attr) == "enzyme_math" && LLVM.value(attr) == "enzyme_custom"
is_custom = true
break
end
end
is_custom || continue
mi, _ = enzyme_custom_extract_mi(fn, false)
mi isa Core.MethodInstance || continue
has_mode_rule(mode, mi, world, method_table) && continue
fattrs = function_attributes(fn)
delete!(fattrs, StringAttribute("enzyme_math"))
delete!(fattrs, StringAttribute(PRESERVEPRIMAL_ATTR_KIND))
end
return nothing
end
Loading
Loading