Forwards-Mode Design

The purpose of this document is to explain how forwards-mode AD in Mooncake.jl is implemented. It should do so to a sufficient level of depth to enable the interested reader to read the forwards-mode AD code in Mooncake.jl and understand what is going on.

This document

specifies the semantics of a "rule" for forwards-mode AD,
specifies how to implement rules by-hand for primitives,
specifies how to derive rules from IRCode algorithmically in general,
discusses batched ("chunked") forwards-mode,
discusses some notable technical differences between our forwards-mode AD implementation details and reverse-mode AD implementation details, and
concludes with a brief comparison with ForwardDiff.jl.

Forwards-Rule Interface

Loosely, a rule for a function simultaneously

performs the same computation as the original function, and
computes the Fréchet derivative.

This is best explained through a worked example. Consider a function call

z = f(x, y)

where f itself may contain data / state which is modified by executing f. rule_for_f is some callable which claims to be a forwards-rule for f. For rule_for_f to be a valid forwards-rule for f, it must be applicable to Lifteds as follows:

z_dz = rule_for_f(lift(f, df), lift(x, dx), lift(y, dy))::Lifted

where:

rule_for_f is a callable. It might be written by-hand, or derived algorithmically.
df, dx, and dy are tangents for f, x, and y respectively. They are the inputs to the derivative of (f, x, y); the tangent carried by z_dz is the output.
z_dz is a Lifted containing the primal z and the component of the derivative of (f, x, y) in the direction (df, dx, dy) associated with z.
running rule_for_f leaves f, x, and y in the same state that running f does.

We refer readers to Algorithmic Differentiation to explain what we mean when we talk about the "derivative" above.

Note that rule_for_f is an as-yet-unspecified callable which we introduced purely to specify the interface that a forwards-rule must satisfy. In Hand-Written Rules and Derived Rules below, we introduce two concrete ways to produce rules for f.

The `Lifted` slot and the forward value `V`

A forwards-mode argument is a Mooncake.Lifted slot:

struct Lifted{P, N, V}
    primal::P
    value::V
end

P is the primal type, N is the chunk width (the number of derivative directions propagated at once — N == 1 for a standard single-direction rule, N > 1 for batched mode), and V is the forward value: the representation of N tangents of P packed together.

For a concrete P it must always hold that V === dual_type(Val(N), P). dual_type plays the same role for forwards-mode that tangent_type plays for reverse-mode, and it is recursively coherent with it: where the reverse representation of a component is tangent_type(component), the forward representation is dual_type(Val(N), component), mirroring each other shape-for-shape. Concretely:

a real scalar P → NDual{P, N} (a value plus an NTuple{N} of per-lane partials);
an array Array{T, D} → NDualArray{T, N, D, ...}, a wrapper holding the primal array (a genuine Array{T, D} usable directly in a ccall, aliasing the user's storage) and the N lane partials in one slot-local element-major block of shape (N, size...), so each element's partials form a contiguous column (per-lane access is a strided view, a.partials[k]);
a MemoryRef{T} (Julia 1.11+) → NDualMemoryRef{T, N, ...}, holding the primal ref plus the backing MemoryRef of the enclosing Memory/Array V's element-major block (not a reshaped copy of it), the number of block columns, and the column that pairs with the referenced element. Storing the backing ref keeps the array.ref projection a pure field read — no reshape, no header allocation — even for a rank-D > 1 array (its block is Array{T, D+1}, whose backing ref is still a plain MemoryRef{T}, so the ref V's type stays coherent regardless of the source's rank). The backing is shared with the enclosing V, so mutations through any V land in one storage, mirroring the primal aliasing;
a struct → ImmutableDual / MutableDual wrapping the per-field forward values;
tuples / named-tuples → element-wise recursion;
a non-differentiable P (integers, Symbol, Module, types, …) → Mooncake.NoDual, the forwards-mode analogue of reverse-mode's NoTangent.

Rules read and write slots through the accessors rather than touching the fields directly: primal(slot), tangent(slot) (the whole V), and tangent(slot, lane) (the lane-th direction's partial). This keeps rule bodies independent of the inner-V shape — see Hand-Written Rules.

Block layout across backends, and performance

Keeping an array's N per-lane partials in a separate block (rather than interleaving NDuals into the primal container) is what lets the primal stay a genuine Array/CuArray usable directly in a ccall, and lets each lane be handed to BLAS/LAPACK as a strided operand. The orientation of that block is chosen per backend to keep the hot path allocation-free, behind the _block_type/_block_dims/_block_reshape helpers so rule bodies never see the difference:

Julia 1.11+ (CPU): an Array{T, D+1} of shape (N, size...) — element-major, so scalar getindex/setindex! (the dominant cost in element-wise forward AD) touches one contiguous per-element column, while BLAS still gets the whole (N, len) lane matrix.
Julia 1.10 (CPU): a flat Vector{T} of length N·length. A 1.10 Array grows its buffer in place, and a shaped Matrix cannot; the flat vector can, so length-changing ops (push!/resize!) stay amortised-O(1). Shape-consuming helpers wrap it to (N, size...) on demand.
CUDA: lane-major (size..., N), so each lane is a contiguous last-dim slice — a GPU never scalar-indexes, but the low-level CUDA primitives (unsafe_copyto!, batched cuBLAS) require contiguous per-lane buffers.

Performance is judged per lane, not per pass: a width-N frule!! computes N directional derivatives, so it necessarily uses N partial buffers and, for direction-scaling ops (matmul, solves, reductions), up to N× a VJP's derivative work — this is fundamental to propagating N directions, not overhead. The target is that a width-1 forward pass matches one reverse VJP, and that nothing lane-invariant (the primal, factorizations, shared products) repeats across lanes. On the posteriordb log-density suite, chunked forward is within a small factor of ForwardDiff at a matched chunk size across every model, with the element-major block beating a parallel-arrays layout on the scalar-heavy targets.

Testing

Suppose that we have (somehow) produced a supposed forwards-rule. To check that it is correctly implemented, we must ensure that

all primal state after running the rule is approximately the same as all primal state after running the primal, and
the inner product between all tangents (both output and input) and a random tangent vector after running the rule is approximately the same as the estimate of the same quantity produced by finite differencing or reverse-mode AD.

We already have the functionality to do this in a very general way (see Mooncake.TestUtils.test_rule), which additionally checks the chunk widths N = 1, 2, 3 and the canonical-V coherence invariant.

Hand-Written Rules

Hand-written rules are implemented by writing methods of two functions: is_primitive and frule!!.

`is_primitive`

is_primitive(context, ::Type{ForwardMode}, signature::Type{<:Tuple}, world) must return true if AD must attempt to differentiate a call by passing the arguments to frule!!, and false otherwise. context is a context type such as Mooncake.MinimalCtx or Mooncake.DefaultCtx — the same contexts used for reverse-mode; the mode is selected by the ForwardMode argument. The Mooncake.@is_primitive macro is used to declare new primitives, e.g.

@is_primitive MinimalCtx ForwardMode Tuple{typeof(sin), Float64}

The declaration must stay in lockstep with the frule!! method coverage: a broader @is_primitive than the rule's methods fails only at call time with a MethodError.

`frule!!`

Methods of frule!! do the actual differentiation, and must satisfy the Forwards-Rule Interface discussed above. A frule!! must return the canonical dual_type(Val(N), typeof(result)) shape (use zero_dual / zero_lifted for a zero derivative); it must never pair a NoDual with a differentiable result, nor double-wrap a Lifted.

In what follows, we will refer to frule!!s for signatures. For example, the frule!! for signature Tuple{typeof(sin), Float64} is the rule which would differentiate calls like sin(5.0).

Simple Scalar Function

Recall that for $y = \sin(x)$ we have that $\dot{y} = \cos(x) \dot{x}$. The scalar NDual overloads compute the value and its partials in a single call (here sin of an NDual returns an NDual whose .value is sin(x) and whose partials carry $\cos(x)\dot{x}$). So the frule!! for signature Tuple{typeof(sin), Float64} is, for any chunk width N:

function frule!!(::Lifted{typeof(sin), N}, x::Lifted{Float64, N, NDual{Float64, N}}) where {N}
    dy = sin(tangent(x))                 # one call: value AND all N lane partials
    return Lifted{Float64, N}(dy.value, dy)   # read the primal back from dy (inner-value invariant)
end

Note that the primal sin(x) is read out of dy.value rather than recomputed — f(tangent(x)) has already evaluated it, so a separate sin(primal(x)) would run a transcendental twice. This "do not recompute the primal" discipline is the dominant cost in element-wise forward AD.

Pre-allocated Matrix-Matrix Multiply

Recall that for $Z = X Y$ we have that $\dot{Z} = X \dot{Y} + \dot{X} Y$. Because the forward value of an array is an NDualArray (the primal is a genuine Array, and the lane partials live in one element-major block accessed per lane as tangent(·).partials[k], a strided view), we can apply the primal mul! once and then a per-lane mul! over the lane partials — writes through the views land in the slot's block:

function frule!!(
    ::Lifted{typeof(mul!), N},
    Z::Lifted{P, N, <:NDualArray}, X::Lifted{P, N, <:NDualArray}, Y::Lifted{P, N, <:NDualArray},
) where {N, P<:Matrix{Float64}}
    mul!(primal(Z), primal(X), primal(Y))            # primal update, once
    for k in 1:N
        dZ = tangent(Z).partials[k]                   # write-through lane view
        dX, dY = tangent(X).partials[k], tangent(Y).partials[k]
        mul!(dZ, primal(X), dY)                       # X * Ẏ
        mul!(dZ, dX, primal(Y), 1.0, 1.0)            # + Ẋ * Y
    end
    return Z
end

The in-place primal update is hoisted out of the per-lane loop: repeating it would corrupt the shared primal seen by later lanes. (In practice we would implement a rule for a lower-level function like LinearAlgebra.BLAS.gemm!, rather than mul!. The actual BLAS/LAPACK rules go one step further: a single lane of the element-major block is a stride-N view, which the pointer-based BLAS wrappers cannot consume, so those rules operate on the dense block itself — a lane-invariant linear map applies to all N lanes in one wide BLAS call by right-multiplying the (N, len) lane matrix by the map's transpose. See _partials_block and the per-rule comments in src/rules/blas.jl.)

Derived Rules

This is the "automatic" / "algorithmic" bit of AD! This is the second way of producing concrete callable objects which satisfy the Forwards-Rule Interface discussed above. The object which we will ultimately construct is an instance of Mooncake.DerivedFRule.

Worked Example: Julia Function

Before explaining how derived rules are produced algorithmically, we explain by way of example what a derived rule should look like if we work things through by hand.

A derived rule for a function such as

function f(x)
    y = g(x)
    z = h(x, y)
    return z
end

should be something of the form

function rule_for_f(::Lifted{typeof(f)}, x::Lifted)
    y = rule_for_g(zero_lifted(Val(1), g), x)
    z = rule_for_h(zero_lifted(Val(1), h), x, y)
    return z
end

Observe that the transformation is simply

replace all variables with Lifted variables,
replace all constants (e.g. g and h) with constant Lifteds (with zero tangent),
replace all calls with calls to rules.

In general, all control flow should be identical between primal and rule.

Worked Example: IRCode

The above example is expressed in terms of Julia code, but we will be operating on Julia Compiler.IRCode, so it is helpful to consider how the above example translates into this form. If we call f on a Float64, and suppose that g and h both return Float64s, the primal Compiler.IRCode will look something like the following:

julia> Base.code_ircode_by_type(Tuple{typeof(f), Float64})
1-element Vector{Any}:
2 1 ─ %1 = invoke Main.g(_2::Float64)::Float64
3 │   %2 = invoke Main.h(_2::Float64, %1::Float64)::Float64
4 └──      return %2
   => Float64

Recall that _2 is the second argument, in this case a Float64, and %1 and %2 are SSAValues. Roughly speaking, the forwards-mode IR for the (ficticious) function rule_for_f should look something like:

julia> Base.code_ircode_by_type(Tuple{typeof(rule_for_f), Lifted{typeof(f), 1, NoDual}, Lifted{Float64, 1, NDual{Float64, 1}}})
1-element Vector{Any}:
2 1 ─ %1 = invoke rule_for_g($(zero_lifted(Val(1), Main.g)), _3::Lifted{Float64, 1, NDual{Float64, 1}})::Lifted{Float64, 1, NDual{Float64, 1}}
3 │   %2 = invoke rule_for_h($(zero_lifted(Val(1), Main.h)), _3::Lifted{Float64, 1, NDual{Float64, 1}}, %1::Lifted{Float64, 1, NDual{Float64, 1}})::Lifted{Float64, 1, NDual{Float64, 1}}
4 └──      return %2
   => Lifted{Float64, 1, NDual{Float64, 1}}

Observe that:

All Arguments have been incremented by 1. i.e. _2 has been replaced with _3. This corresponds to the fact that the arguments to the rule have all been shuffled along by one, and the rule itself is now the first argument.
Everything has been turned into a Lifted.
Constants such as g appear as constant Lifteds with a zero tangent (zero_lifted(Val(1), Main.g)).

Here, as before, we have not specified exactly what rule_for_f, rule_for_g, and rule_for_h are. This is intentional – they are just callables satisfying the Forwards-Rule Interface. In the following we show how to derive rule_for_f, and show how rule_for_g and rule_for_h might be methods of Mooncake.frule!!, or themselves derived rules.

Rule Derivation Outline

Equipped with some intuition about what a derived rule ought to look like, we examine how we go about producing it algorithmically.

Rule derivation is implemented via the function Mooncake.build_frule. This function accepts a context and a signature / Base.MethodInstance / MistyClosure (plus a chunk_size keyword, see Batch Mode) and, roughly speaking, does the following:

Look up the optimised Compiler.IRCode.
Apply a series of standardising transformations to the IRCode.
Transform each statement according to a set of rules to produce a new IRCode.
Apply standard Julia optimisations to this new IRCode.
Put this code inside a MistyClosure in order to produce an executable object.
Wrap this MistyClosure in a DerivedFRule to handle various bits of book-keeping around varargs.

In order:

Looking up the `Compiler.IRCode`.

This is done using Mooncake.lookup_ir. This function has methods which return the IRCode associated to:

signatures (e.g. Tuple{typeof(f), Float64})
Base.MethodInstances (relevant for :invoke expressions – see Statement Transformation below)
MistyClosures.MistyClosure objects, which is essential when computing higher order derivatives and Hessians by applying Mooncake.jl to itself.

Standardisation

We apply the following transformations to the Julia IR. They can all be found in ir_normalisation.jl:

Mooncake.foreigncall_to_call: convert Expr(:foreigncall, ...) expressions into Expr(:call, Mooncake._foreigncall_, ...) expressions.
Mooncake.new_to_call: convert Expr(:new, ...) expressions to Expr(:call, Mooncake._new_, ...) expressions.
Mooncake.splatnew_to_call: convert Expr(:splatnew, ...) expressions to Expr(:call, Mooncake._splat_new_...) expressions.
Mooncake.intrinsic_to_function: convert Expr(:call, ::IntrinsicFunction, ...) to calls to the corresponding function in Mooncake.IntrinsicsWrappers.

The purpose of converting Expr(:foreigncall...), Expr(:new, ...) and Expr(:splatnew, ...) into Expr(:call, ...)s is to enable us to differentiate such expressions by adding methods to frule!!(::Lifted{typeof(Mooncake._foreigncall_)}), frule!!(::Lifted{typeof(Mooncake._new_)}), and frule!!(::Lifted{typeof(Mooncake._splat_new_)}), in exactly the same way that we would for any other regular Julia function.

The purpose of translating Expr(:call, ::IntrinsicFunction, ...) is to do with type stability – see the docstring for the Mooncake.IntrinsicsWrappers module for more info.

Statement Transformation

Each statement which can appear in the Julia IR is transformed by a method of Mooncake.modify_fwd_ad_stmts!. Consequently, this transformation phase simply corresponds to iterating through all of the expressions in the IRCode, applying Mooncake.modify_fwd_ad_stmts! to each, modifying the IRCode in place. To understand how to modify IRCode and insert new instructions, see Oxinabox's Gist.

We provide here a high-level summary of the transformations for the most important Julia IR statements, and refer readers to the methods of Mooncake.modify_fwd_ad_stmts! for the definitive explanation of what transformation is applied, and the rationale for applying it. In particular there are quite a number more statements which can appear in Julia IR than those listed here and, for those we do list here, there are typically a few edge cases left out.

Expr(:invoke, method_instance, f, x...) and Expr(:call, f, x...)

:call expressions correspond to dynamic dispatch, while :invoke expressions correspond to static dispatch. That is, if you see an :invoke expression, you know for sure that the compiler knows enough information about the types of f and x to prove exactly which specialisation of which method to call. This specialisation is method_instance. Conversely, a :call expression typically occurs when the compiler has not been able to deduce the exact types of f and x, and therefore has to wait until runtime to figure out what to call, resulting in dynamic dispatch.

As we saw earlier, the idea is to translate these kinds of expressions into something vaguely along the lines of

Expr(:call, rule_for_f, f, x...)

There are three cases to consider, in order of preference:

Primitives:

If is_primitive returns true when applied to the signature constructed from the static types of f and x, then we simply replace the expression with Expr(:call, frule!!, f, x...), regardless whether we have an :invoke or :call expression. (Due to the Standardisation steps, it regularly happens that we see :call expressions in which we actually do know enough type information to do this, e.g. for Mooncake._new_ :call expressions).

Static Dispatch:

In the case of :invoke nodes we know for sure at rule compilation time what rule_for_f must be. We derive a rule for the call by passing method_instance to Mooncake.build_frule. (In practice, we might do this lazily, but while retaining enough information to maintain type stability, mirroring the Mooncake.LazyDerivedRule used in reverse-mode).

Dynamic Dispatch:

If we have a :call expression and are not able to prove that is_primitive will return true, we must defer dispatch until runtime. We do this by replacing the :call expression with a call to a Mooncake.DynamicFRule, which simply constructs (or retrieves from a cache) the rule at runtime. Reverse-mode utilises a similar strategy via Mooncake.DynamicDerivedRule.

The above was written in terms of f and x. In practice, of course, we encounter various kinds of constants (e.g. Base.sin), Arguments (e.g. _3), and Core.SSAValues (e.g. %5). The translation rules for these are:

constants are turned into constant Lifteds in which the tangent is zero (via zero_lifted(Val(N), v)),
Arguments are incremented by 1,
SSAValues are left as-is.

Core.GotoNodes

These remain entirely unchanged.

Core.GotoIfNot

These require minor modification. Suppose that a Core.GotoIfNot of the form Core.GotoIfNot(%5, 4) is encountered in the primal. Since %5 will be a Lifted in the derived rule, we must pull out its primal field and pass that to the conditional instead. Therefore, these statements get lowered to two lines in the derived rule. For example, Core.GotoIfNot(%5, 4) becomes a call to Mooncake._primal on the (incremented) condition, immediately followed by the GotoIfNot on the result:

%n = Mooncake._primal(%5)
Core.GotoIfNot(%n, 4)

Core.PhiNode

Core.PhiNode looks something like the following in the general case:

φ (#1 => %3, #2 => _2, #3 => 4, #4 => #undef)

They map from a collection of basic block numbers (#1, #2, etc) to values. The values can be Core.Arguments, Core.SSAValues, constants (literals and QuoteNodes), or undefined.

Core.PhiNodes in the primal are mapped to Core.PhiNodes in the rule. They contain exactly the same basic block numbers, and apply the following translation rules to the values:

Core.SSAValues are unchanged.
Core.Arguments are incremented by 1 (as always).
constants are translated into constant Lifteds (with a zero tangent, via zero_lifted).
undefined values remain undefined.

The node's declared type is also widened to lifted_type(Val(N), ...). So the above example would be translated into something like

φ (#1 => %3, #2 => _3, #3 => $(zero_lifted(Val(1), 4)), #4 => #undef)

The related PiNode, UpsilonNode, and PhiCNode statements are handled analogously (increment Arguments, zero_lifted constants, widen the declared type with lifted_type).

Optimisation

The IR generated in the previous step will typically be uninferred, and suboptimal in a variety of ways. We fix this up by running inference and optimisation on the generated IRCode. This is implemented by Mooncake.optimise_ir!.

Put IRCode in MistyClosure

Now that we have an optimised IRCode object, we need to turn it into something that can actually be run. This can, in general, be straightforwardly achieved by putting it inside a Core.OpaqueClosure. This works, but Core.OpaqueClosures have the disadvantage that once you've constructed a Core.OpaqueClosure using an IRCode, it is not possible to get it back out. Consequently, we use MistyClosures, in order to keep the IRCode readily accessible if we want to access it later.

Put the MistyClosure in a DerivedFRule

See the implementation of DerivedRule (used in reverse-mode) for more context on this. This is the "rule" that users get.

Batch Mode

So far, we have mostly assumed that we would only apply forwards-mode to a single tangent vector at a time (chunk width N == 1). However, in practice, it is typically best to pass a collection of tangents through at a time — for example, computing a gradient or a Jacobian via forwards-mode requires one derivative direction per input degree of freedom, and propagating them in chunks amortises the cost of the primal computation.

No separate "batched" transformation is needed: the forward value V = dual_type(Val(N), P) is already parameterised by the chunk width N. Every statement transformation above threads the width (Val(info.width)) into the constant zero_lifteds and the lifted_type annotations, and the hand-written frule!!s loop over the N lanes (tangent(slot, k)) as shown in the matrix-multiply example. A width-N rule is obtained simply by passing chunk_size = N to Mooncake.build_frule; the NDual / NDualArray forward values then carry N partials per primal instead of one.

Forwards vs Reverse Implementation

The implementation of forwards-mode AD is quite dramatically simpler than that of reverse-mode AD. Some notable technical differences include:

forwards-mode AD only makes use of the (forward) tangent system — the Lifted slot and its dual_type forward value — whereas reverse-mode also makes use of the fdata / rdata system.
forwards-mode AD comprises only line-by-line transformations of the IRCode. In particular, it does not require the insertion of additional basic blocks, nor the modification of the successors / predecessors of any given basic block (the GotoIfNot rewrite inserts a node within the same block). Consequently, there is no need to make use of the CFGBlock infrastructure built up for reverse-mode AD – everything can be straightforwardly done at the Compiler.IRCode level.

Comparison with ForwardDiff.jl

With reference to the limitations of ForwardDiff.jl, there are a few noteworthy differences between ForwardDiff.jl and this implementation:

:foreigncalls pose much less of a problem for Mooncake's forward-mode than for ForwardDiff.jl, because we can write a rule for any method of any function. In essence, you can only (reliably) write rules for ForwardDiff.jl via dispatch on ForwardDiff.Dual.
the target function can be of any arity in Mooncake.jl, but must be unary in ForwardDiff.jl.
there are no limitations on the argument type constraints that Mooncake.jl can handle, while ForwardDiff.jl requires that argument type constraints be <:Real or arrays of <:Real.
No special storage types are required with Mooncake.jl, while ForwardDiff.jl requires that any container you write to is able to contain ForwardDiff.Duals. (Mooncake's array forward value, NDualArray, keeps the partials in a separate element-major block rather than interleaving them into the primal container.)