Internal API

EnzymeError

Common supertype for Enzyme-specific errors.

This type is made public so that downstream packages can add custom error hints for the most common exceptions thrown by Enzyme.

abs_ntuple_type(arg::LLVM.Value) -> Union{Nothing, Tuple{LLVM.Value, Any}}

If arg is a call julia.enzyme.ntuple_type(T, count) with T known statically, return (count, T); otherwise nothing.

Helper function for llvm-level rule generation. Will call callsamewithinvertedargifactive with corresponding extracted batches if width > 1, otherwise it will call it once.

box_inline_union!(B, alloctx, val, offset, UT) -> LLVM.Value

val is an aggregate that holds an isbits Union field of type UT inline at byte offset: the payload, then a selector byte with the 0-based index of the member in Base.uniontypes order. Return the selected member boxed, as a tracked pointer; a singleton member is its instance. The tape slot of a Union tape holds the union boxed, as jl_type_to_llvm declares it, and the reverse rule takes such a tape boxed too.

Helper function for llvm-level rule generation. Will call the same function (and optional postprocessing), if the argument at index cmpidx isn't active. This takes into account runtime activity as a reason the value may not be active.

If postprocess_const is set, the original function will always be called, but the postprocessing will be conditionally gated as follows.

If the relevant input is active (and verified by runtime activity), postprocess(B, result, args) will run as normal Otherwise postprocess_const(B, result, args) will run

clear_caches!()

Empty every cache Enzyme keeps in a global, dropping what this session compiled, looked up and rooted along with them.

Nearly all of it means something only to the session that filled it. A CompileResult holds the address the JIT gave a thunk, captured_constants roots objects because their addresses were written into that code, the rule and activity memos are keyed on world ages, and the jl_load_and_lookup handles are ones this process opened. Anything outliving the session must not carry them, which is why Enzyme's precompile workload ends with this call: what it left behind would otherwise be serialized into Enzyme's package image and inherited, dead, by every session that loads it.

This is meant for the end of precompilation and not for a live session. It hands back the thunks the JIT compiled and unroots the objects their code refers to by address, so a thunk still held anywhere is left pointing at objects that may now be collected.

Caches filled by __init__ rather than by compiling are left alone: they are rebuilt per session and so never reach an image.

copy_abi_attrs!(call::LLVM.CallInst, fn::LLVM.Function)

Copy the zeroext, signext and swiftself parameter attributes of fn to the call site call. restore_lookups replaces the callee with a constant address, after which only the call site's attributes describe the convention the callee expects: the extension of a small integer argument, and the register pgcstack goes in. A target without the swift calling convention takes pgcstack as a plain parameter, which carries no attribute to copy.

copy_metadata!(dst, src)

Attach every metadata node of the instruction src, except its debug location, to dst.

current_pgcstack() -> Ptr{Cvoid}

Give the pgcstack of the running task, for an llvmcall that takes it as an argument (see use_gcstack_arg!).

Julia computes current_task() from the pgcstack of the function that this is inlined into (on 1.13 that is the "gcstack" argument of the function). So the function gets no julia.get_pgcstack call from this. An llvmcall of julia.get_pgcstack would not do: Julia inlines it where it is called, which is the problem that use_gcstack_arg! avoids.

decay_args_readonly(st, fop, args)

Whether the call st merely reads through each of the argument positions in args. That is the case when the position itself is marked readonly / readnone, or when the callee as a whole only reads memory – which is how plain libcalls such as memcmp are annotated. Positions that carry an sret-like marker are never treated as read-only, since the callee writes the object back through them.

declare_ntuple_type!(mod::LLVM.Module)

Declare julia.enzyme.ntuple_type(T, count), which returns NTuple{count, T}. It stays a bare declaration while Enzyme differentiates, including in the module kept for nested differentiation, so that no optimization drops its constant T argument and abs_ntuple_type can always recover T. define_ntuple_type! gives it a body after differentiation.

define_ntuple_type!(mod::LLVM.Module)

Give julia.enzyme.ntuple_type, if mod declares it, an alwaysinline body. For count <= NTUPLE_TYPE_STACK_SIZE it fills a stack array with T and calls jl_apply_tuple_type_v; otherwise it calls jl_apply_tuple_type(jl_svec_fill(count, T)). Both find the interned type several times faster than jl_f_apply_type(NTuple, count, T), which first instantiates the NTuple UnionAll. Call this once differentiation is done, after the module for nested differentiation has been saved.

dereferenceable_root_ptr(v) -> Bool

Whether a pointer-sized load from v is safe even where the program would not have loaded from it: v points into an alloca, or at a constant offset into an argument whose dereferenceable attribute covers the load.

emit_ntuple_type!(B, count, T) -> LLVM.Value

Emit the type NTuple{count, T} for a runtime count::Int as a call to julia.enzyme.ntuple_type (see declare_ntuple_type! and define_ntuple_type!).

extract_roots_from_value!(builder, sret, roots)

Store the GC-tracked fields of the value sret into the returnRoots array roots, which must have room for CountTrackedPointers(value_type(sret)).count entries.

This is the split recombine_value! undoes: a caller that needs to hand a value on through the sret/returnRoots convention writes the tracked pointers here and the inline data into the sret buffer separately.

fixup_1p12_sret!(f::LLVM.Function)

Rewrite the untyped store of a return value that needs both an sret buffer and a returnRoots array into field-wise stores that skip the GC-tracked slots.

Julia 1.12 changed the convention for such a return: the callee writes the tracked pointers only into returnRoots and leaves the corresponding slots of the sret buffer undefined, so a caller has to recombine the two halves (which is what recombine_value! does). Codegen spells the write of the remaining, inline data as a single untyped llvm.memcpy out of an [N x i64] alloca, which also drags the undefined bytes into the tracked slots – and Enzyme would then take those for live jlvalues. Replacing the memcpy with stores of just the untracked fields keeps the tracked slots alone, matching what the convention promises the caller.

The rewrite is keyed on the ABI actually present in the IR rather than on a version bound: it only fires for a memcpy whose destination parameter carries an sret attribute for exactly RT.

gcstack_arg_index(fn::LLVM.Function) -> Int

Give the index of the pgcstack parameter of fn, or 0 when fn has none.

Julia's codegen marks that parameter swiftself where the target supports the swift calling convention, and turns the convention off where it does not (jl_codegen_output_t::use_swiftcc, false on RISC-V). Julia 1.13 also gives the parameter a gcstack string attribute, which specsig writes on every version. Hence recognize either mark.

has_gcstack_arg(fn::LLVM.Function) -> Bool

Say if fn takes pgcstack as a parameter (see gcstack_arg_index).

import_cached_autodiff!(mod, ptr, FT)

Make the derivative thunk compiled at ptr callable from mod, and return the function to call in its place.

cached_compilation keeps the bitcode of every thunk it compiles in autodiff_cache, so that a later compilation which sees a call to the thunk's address can use its IR instead of calling an opaque pointer. That blob is a fixed serialization: every module it is linked into gets the very same symbol names for the Julia functions it carries. compile_unhooked links all the modules nested_codegen! emitted into one module, so two of them importing the same blob define those symbols twice and LLVM.link! rejects the second one ("symbol multiply defined", EnzymeAD/Enzyme.jl#2788). The names collide even though nothing else does, because they were minted once, when the thunk was compiled, and are replayed verbatim on every import.

Import the blob once per compilation instead. The first module to need it gets the definitions; every later one gets a declaration of the entry, which the final link binds to that single definition. The entry stays externally visible for as long as those declarations do: compile_unhooked internalizes it again once every module is linked (see internalize_imported_thunks!), so it is inlined and discarded just as a lone import was.

instantiate_annotation(A, rt, width)

Fill in the free parameters of a (possibly partially applied) activity annotation A with element type rt and batch width width.

The batch annotations take a second parameter carrying the batch width. Applying only A{rt} to them leaves that parameter free, and a subsequent A{rt} binds the element type to it, yielding an annotation whose batch_size is a type rather than the width. Filling both explicitly keeps batch_size(A) == width, which the shadow-return ABI in create_abi_wrapper and enzyme_call asserts.

internalize_imported_thunks!(mod)

Undo the external linkage import_cached_autodiff! gave the entry of every cached thunk this compilation imported.

The entry is externally visible only so that the modules which did not import the blob can declare it and have the final link bind them to the one definition. Once mod holds every such module, nothing outside it refers to the entry any more, and leaving it visible would keep a copy of the thunk that post_optimize! may neither inline nor drop.

is_mutable_array(T::Type)::Bool

Return true if T refers to mutable memory of elements of type eltype(T), as an Array does. The activity of T then comes from eltype(T). A package can add methods for its array or pointer types, for example in a package extension. Enzyme calls this function in the world of the compilation.

is_readonly(attr::LLVM.Attribute)::Bool

Whether attr on its own establishes that the function or argument position it is attached to is only read from. That is readonly / readnone, and on LLVM 16+ a memory effect whose modref is read-only.

legalize_readonly_decay!(st, inst, args)

Rewrite the argument positions args of the read-only call st, all of which are the illegal addrspace(10) -> addrspace(0) cast inst, to use a legally derived pointer instead: decay to addrspace 11 and go through julia.pointer_from_objref, with the object gc-preserved across the call. Julia's late GC lowering understands that form; the raw cast it does not.

link_split_existing!(mod::LLVM.Module, newmod::LLVM.Module)

Link newmod into mod like LLVM.link!(mod, newmod), but set LLVMInternalLinkage on any function defined in both modules before linking. This allows LLVM's linker to natively internalize and resolve duplicate definitions without string comparisons or linker collisions.

mark_load_dereferenceable!(inst::LLVM.LoadInst, @nospecialize(source_typ), byref)

Julia only marks a loaded pointer to a heap object as dereferenceable when it comes from a mutable struct's field; a field of an immutable struct gets no such metadata. Enzyme's abs_typeof knows the exact Julia type of the loaded value, so if inst loads a tracked pointer to an object of the concrete type source_typ (byref == MUT_REF), mark it dereferenceable_or_null for that type's size, as Julia's codegen does in the mutable case. Whether it is also non-null is left to Julia's own nonnull, which it emits exactly when the field cannot be #undef; with both, LLVM may speculate loads through the pointer, which lets LICM hoist e.g. an array's Memory pointer out of loops it is not guaranteed to be loaded in, where Enzyme otherwise has to cache it per iteration.

mark_loads_dereferenceable!(fn::LLVM.Function)::Bool

Apply mark_load_dereferenceable! to every load of a tracked pointer in fn whose Julia type abs_typeof can determine. Runs as a pass right before the loop passes of the early pipeline, since earlier passes recreate such loads without their metadata.

memtransfer_truetype(world, ptr, sz)

The enzyme_truetype metadata for a memcpy/memmove/memset of sz bytes at ptr, or nothing if ptr cannot be traced back to a Julia object of known concrete type. The type is read off the Julia layout, so it is not limited to the offsets Enzyme's type analysis keeps.

nullify_rooted_values!(builder, sret)

Return the value sret with every GC-tracked field replaced by a null reference.

Used where only the inline data of a value is wanted, and the tracked fields are either held elsewhere (in a returnRoots array, see extract_roots_from_value!) or known not to be needed, so that leaving the original pointers in place would root objects that must not be kept alive.

precedes(a, b)

Whether instruction a comes before b in their common basic block.

recombine_value!(builder, sret, roots; must_cache=false)

Rebuild a whole return value from the two halves the sret/returnRoots calling convention splits it into.

A callee returning a type with both GC-tracked and inline fields writes the tracked pointers into the returnRoots array and the remaining data into the sret buffer, leaving the tracked slots of sret undefined. sret here is the value already loaded out of that buffer and roots the pointer to the root array; the tracked fields are loaded from roots and inserted back into their slots, and the completed value is returned. This is the inverse of extract_roots_from_value!; see recombine_value_ptr! for the variant that takes sret as a pointer.

must_cache marks the loads from roots as must-cache, for a caller that needs the recombined value to survive into the reverse pass.

recombine_value_ptr!(builder, jltype, sret, roots; must_cache=false)

Like recombine_value!, but loads the inline half out of the sret buffer rather than taking it as an already-loaded value.

Both sret and roots are pointers; a fresh jltype value is built by loading the untracked fields from sret and the tracked ones from roots.

restore_native_invokes!(mod::LLVM.Module)

Bind the declarations of natively called functions in mod to their entry addresses. restore_lookups(mod; native_invokes = false) skips them, so that the module can still be differentiated again while they are symbolic.

rewrite_abi_converter_calls!(mod::LLVM.Module)

Julia 1.12+ lowers @cfunction to a world-age-guarded dispatch site that calls jl_get_abi_converter to obtain a callable pointer for the target in the current world. That runtime resolver picks a code instance from the native JIT cache, which is the wrong one for GPUCompiler-emitted code: both its owner (ci->owner) and its ABI (compiled with gcstack_arg = true, i.e. pgcstack in the swiftself register) differ from this module (gcstack_arg = false), so the raw specsig pointer handed back is called with a mismatched ABI and reads a garbage GC stack out of the swiftself register (#3284).

Rewrite each such site to act like jl_apply_generic instead: Julia's codegen already emitted an in-module unspecialized dispatcher thunk for every site (stored in the cfuncdata global) that boxes the arguments and dispatches via jl_apply_generic. That thunk was compiled with this module's own ABI and resolves the callee from the correct cache in the current world, so replace every jl_get_abi_converter call with it and let the world-age guard fold away.

root_block_after(inst, bb)

Whether inst is in bb, or in a successor of bb that has no other predecessor, so that bb runs right before it.

root_load_in(inst, ptr, bb, T_prjlvalue)

Whether inst is a plain (non-volatile, non-atomic) load of a tracked pointer from ptr, placed in the block bb.

root_path_may_write(inst) -> Bool

Whether inst may write memory. Unlike mayWriteToMemory, which only reads the attributes of the call site, a call is also known not to write when its callee is read-only, as for intrinsics such as llvm.smax that LLVM leaves between a phi and the loads it moved into a successor.

roots_follow(args, ai, removedRoots) -> Bool

Whether the classified argument after args[ai] carries the inline roots of args[ai] and is folded back into it (removedRoots), so that args[ai] is handled through its data pointer rather than loaded whole.

shrinks_split_values() -> Bool

Whether this Julia trims trailing tracked pointers from the data half of a split value (see split_value_size).

split_value_into!(builder, val, sret, roots)

Store the value val through the sret/returnRoots convention: its GC-tracked fields into the roots array and every other field into the sret buffer, leaving the tracked slots of the buffer alone, as a caller reads them from roots only. The inverse of recombine_value_ptr!.

split_value_size(dl, T) -> Int

Size in bytes of the data half of a value of LLVM type T when Julia keeps it on the stack split into inline data and GC roots (split_value_size in cgutils.cpp).

Up to Julia 1.13.0 the data half kept the full layout, with the tracked slots left undefined. Since 1.13.1 (JuliaLang/julia#60388) codegen shrink-wraps it: pointer words at the very end of the layout are dropped, while interior pointer slots stay as padding so the offsets of the remaining fields are unchanged. The buffer behind the by-reference pointer of an argument with inline roots, and the stack copy a new makes of such a value, are only this many bytes; a callee may not describe or touch more. An sret buffer is not affected: its parameter is still declared dereferenceable for the full layout and callers allocate it whole, only the memcpy that fills it may stop short (see fixup_1p12_sret!).

tracked_pointer_offsets!(offs, dl, T, base=0)

Append to offs the byte offset, relative to base, of every GC-tracked pointer leaf of the LLVM type T in layout order.

unfold_root_phi_loads!(f::LLVM.Function) -> Bool

Turn a load through a phi of root-array pointers back into a phi of loads.

Julia 1.13.1+ keeps the inline GC roots of a split value lazily, as a pointer into whatever roots array already holds them (jl_gc_roots_t), so the roots a phi merges may be loaded right in its predecessors: from the roots argument of the function on one edge, from a local roots alloca on another. LLVM's instcombine then folds that phi of loads into one load of a phi ptr of the arrays. Enzyme promotes every alloca of the augmented primal to a heap allocation and cannot push that address-space change through a phi whose other operands are not allocas ("Illegal address space propagation"). Sinking the load back into the predecessors gives the alloca only plain loads again.

Only rewrite what is certainly equivalent: a phi ptr in address space 0 with an alloca among its incoming values, whose users are all non-atomic loads of tracked pointers, in its own block or in a successor whose only predecessor it is, that no memory write precedes. Each load is re-created before the terminator of every predecessor. That load is speculative on an edge whose predecessor has other successors, and on every edge when the original load is in the successor, so it is then only done for a pointer that is always dereferenceable: an alloca, or an argument whose dereferenceable bytes cover it.

use_gcstack_arg!(f::LLVM.Function, arg::LLVM.Argument)

Make the llvmcall f use its argument arg as the pgcstack of the code it inlines. The caller gives this argument with current_pgcstack, as a Ptr{Cvoid}.

Julia inlines an llvmcall into its caller, together with the alwaysinline functions that the llvmcall calls. A julia.get_pgcstack call in that code thus goes into the middle of the caller. On 1.13 the caller takes its pgcstack as its own "gcstack" argument. But the GC lowering prefers a julia.get_pgcstack call in the entry block to that argument, and pushes the GC frame of the caller only after the call. Then no safepoint before the call has the roots of the caller.

Thus inline the alwaysinline functions into f here, and replace every julia.get_pgcstack call in f with arg, as LowerPTLS does for a function that takes its pgcstack as an argument.

Enzyme.guess_activity(::Type{T}, mode::Enzyme.Mode)

Try to guess the most appropriate Annotation for arguments of type T passed to autodiff with a given mode.

enzyme_call_kind(interp, specTypes) -> Union{Nothing, Symbol}

Say how Enzyme handles a call of the signature specTypes, or nothing when it handles it like any other call:

• :primitive: Enzyme differentiates the call itself (is_primitive_func).
• :alwaysinline: the inliner must always inline it (is_alwaysinline_func).
• :inactive, :frule, :rrule: the signature has a rule of that kind, and interp has that kind of rule enabled.

Every place that decides whether a call needs Enzyme's handling asks this, so that a call reaches the pipeline the same way whichever path inference took to it: FutureCallinfoByType for an ordinary call, and invoked_ci_needs_rule for invoke(f, ci, args...).