Redesign forward mode around Lifted{P,N,V} and chunked NDual duals

[yebai]

CI Summary — GitHub Actions

Documentation Preview

Mooncake.jl documentation for PR #1215 is available at:
https://chalk-lab.github.io/Mooncake.jl/previews/PR1215/

Performance

Performance Ratio:
Ratio of time to compute gradient and time to compute function.
Warning: results are very approximate! See here for more context.

┌────────────────────────────┬──────────┬──────────┬─────────────┬─────────┬─────────────┬────────┐
│                      Label │   Primal │ Mooncake │ MooncakeFwd │  Zygote │ ReverseDiff │ Enzyme │
│                     String │   String │   String │      String │  String │      String │ String │
├────────────────────────────┼──────────┼──────────┼─────────────┼─────────┼─────────────┼────────┤
│                   sum_1000 │ 210.0 ns │     1.38 │        5.49 │   0.571 │        2.81 │   5.34 │
│                  _sum_1000 │ 981.0 ns │     7.04 │        1.86 │  4360.0 │        41.8 │   1.05 │
│               sum_sin_1000 │  7.04 μs │     2.86 │        1.51 │    1.52 │        11.3 │   1.82 │
│              _sum_sin_1000 │  6.18 μs │     2.86 │        1.67 │   260.0 │        12.8 │   2.13 │
│                   kron_sum │ 221.0 μs │     12.4 │        15.3 │    10.8 │       329.0 │   18.8 │
│              kron_view_sum │ 302.0 μs │     10.9 │        19.4 │    24.1 │       299.0 │   12.4 │
│      naive_map_sin_cos_exp │  2.36 μs │     2.54 │        1.72 │ missing │        7.46 │   2.11 │
│            map_sin_cos_exp │  2.29 μs │     2.87 │        2.14 │     1.5 │        6.71 │    2.7 │
│      broadcast_sin_cos_exp │  2.36 μs │     2.74 │        1.64 │    3.81 │        1.44 │   2.07 │
│                 simple_mlp │ 260.0 μs │     5.96 │        76.6 │    2.45 │        10.9 │   2.73 │
│                     gp_lml │ 231.0 μs │      8.5 │        11.8 │    4.03 │     missing │   4.98 │
│ turing_broadcast_benchmark │  1.67 ms │     5.51 │        2.92 │ missing │        35.8 │    2.4 │
│         large_single_block │ 421.0 ns │     5.81 │        2.83 │  4750.0 │        36.0 │   2.05 │
└────────────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

Related: TuringLang/Bijectors.jl#480

[sunxd3]

This is obviously great work. Here are some issues worked out between Claude, GPT, and me. Very happy to discuss the points deeper.

The main text is written by Claude, sorry about the readability!

---

Issues

traced objref round-trips segfault (src/rules/foreigncall.jl:125-175)

using Mooncake
f5(x) = (r = Ref(x); p = pointer_from_objref(r);
         r2 = unsafe_pointer_to_objref(p)::Base.RefValue{Float64}; r2[])
cache = Mooncake.prepare_derivative_cache(f5, 1.0)
Mooncake.value_and_derivative!!(cache, (f5, Mooncake.NoTangent()), (1.0, 1.0))
# signal 11, deterministic — also mutation-free, with GC disabled, and in debug mode

Root cause: the pointer_from_objref(::Ref) frule returns Lifted{Ptr{Nothing},1,Tuple{Ptr{Float64}}} (lanes retyped to point into the partials buffer), while the transform annotates the slot with the canonical lifted_type(Val(1), Ptr{Nothing}) == Lifted{Ptr{Nothing},1,Tuple{Ptr{Nothing}}}. Both are concrete and Lifted is invariant, so the runtime value violates its own slot annotation — undefined behaviour in the compiled OpaqueClosure. The eltype mismatch is unique to the objref family (a pointer(::Array) chain has matching eltypes), which is why CI never trips it, and debug mode is blind to it too: verify_lifted_type has a blanket Ptr exemption.

Two adjacent problems in the same code path:

• the direct (non-traced) unsafe_pointer_to_objref frule rebuilds the NDualRef from a fresh snapshot buffer, severing aliasing with the original Ref's tangent storage — writes through the recovered alias are lost. This is a regression vs main (the width-1 Dual era recovered the aliased tangent).
• once the lanes are made canonical (fully isbits), nothing roots the primal or the tangent object across the round-trip window, so a fix needs a rooting story as well.

nested mutable structs cannot be lane-extracted or seeded (src/tangents/lifted.jl:287-316)

mutable struct Inner; a::Float64 end
struct Outer; i::Inner; x::Float64 end
Mooncake.tangent(Mooncake.zero_lifted(Val(2), Outer(Inner(1.0), 2.0)), 1)
# MethodError: Cannot `convert` … MutableDualTangentView … to … MutableTangent
Mooncake.tangent(Mooncake.zero_lifted(Val(2), [Inner(1.0)]), 1)   # same
# and the FunctionWrappers ext hits the same root cause at zero_lifted time:
# zero_lifted(Val(1), FunctionWrapper(MutableScale(2)))  →  MethodError

tangent(slot, lane) returns the MutableDualTangentView write proxy for any nested MutableDual, and there is no convert to the declared reverse MutableTangent backing. These accessors feed the chunked gradient/Jacobian extraction paths in interface.jl, so this is probably the highest user-impact item: structs with mutable struct fields (typical model objects) are exactly the inputs people will try first. Tuple-nested mutables work via _unlift_seed's MutableDual path — which also points at the fix: lane-projected materialisation for the struct-field/array-element cases, keeping the view proxy for direct rule-body access.

repeated mutable arguments: silent wrong JVPs, and forward ≠ reverse (src/interface.jl:712-716, :1642-1643)

mutable struct MSA; x::Float64 end
fs2(a, b) = a.x + 2b.x
a = MSA(1.0)

fcache = Mooncake.prepare_derivative_cache(fs2, a, a)
Mooncake.value_and_gradient!!(fcache, fs2, a, a)   # gradients: (x = 3.0), (x = 3.0)

rcache = Mooncake.prepare_gradient_cache(fs2, a, a)
Mooncake.value_and_gradient!!(rcache, fs2, a, a)   # gradients: (x = 1.0), (x = 2.0)

Two related causes: the shared lift cache is first-lift-wins, so the second slot reuses the first slot's Lifted and its seed is discarded; and the chunked-gradient sweep's whole-tuple zero_lifted dedups repeated mutable primals while total_dof/the scatter count each argument independently. Array inputs are not affected (I verified 1.0/2.0 correct) — this is the struct-lift cache path only.

Whatever the intended aliasing contract is, the two modes have to agree on it, and the current forward answer matches neither reverse nor the documented forward contract (slot-local independent JVP directions). One constraint for the fix: the shared cache is deliberate for the rule-capture case (a reverse rule captured in grad_f shares fwds_oc/pb_oc state whose forward tangents must be shared — the comment at interface.jl:712 says why), so the fix is per-input-slot caches that keep within-slot sharing, not removing the cache.

silently wrong numbers in three rule families

CUDA mul! drops active alpha/beta tangents (ext/MooncakeCUDAExt/MooncakeCUDAExt.jl:2226-2263, :2343-2379). The frule signatures accept alpha::Lifted{<:Number}/beta::Lifted{<:Number} but the bodies read only primal(alpha)/primal(beta) — the coefficient tangents never enter the JVP. GPU MWEs: zero JVPs where the true values are 44/34 (GEMM dα/dβ) and 14/15 (GEMV). This is a regression from main, which required NoTangent here and failed loudly. Fix is either the missing dα·op(A)op(B) + dβ·C_old terms (hoisted once, per-lane) or restricting the coefficient slots to NoDual V so active coefficients keep failing loudly.

trsm! loses BLAS's α = 0 strong-zero contract (src/rules/blas.jl:1587-1606).

using Mooncake, LinearAlgebra.BLAS
const A = [1.0 1.0; 0.0 0.0]   # singular triangular — fine at α = 0: BLAS never references A
f(B) = (BLAS.trsm!('L', 'U', 'N', 'N', 0.0, A, B); sum(B))
f([1.0 2.0; 3.0 4.0])          # 0.0 (exact zeros)
cache = Mooncake.prepare_derivative_cache(f, [1.0 2.0; 3.0 4.0])
Mooncake.value_and_derivative!!(cache, (f, Mooncake.NoTangent()), ([1.0 2.0; 3.0 4.0], ones(2, 2)))
# (NaN, NaN) — the returned *value* is wrong, not just the derivative

The (otherwise correct) hoist-unscaled-solve optimisation performs the solve and then scales by α, so α = 0 with a singular/garbage A NaN-poisons both the primal result and the partials. Needs an explicit iszero(α) branch; trmm! has a related hazard under BLAS's "B need not be set at α = 0" contract (0·NaN in the unit-diagonal correction).

clamp diverges from Base, silently (src/nfwd/Nfwd.jl:1045-1054). Three ways, all verified:

• mixed precision narrows the bounds to T before comparing: clamp(x, 0.50000001, 1.0) on a Float32 dual at 0.5f0 returns 0.5f0 (derivative 0) where Base returns 0.50000001::Float64 — wrong value and type, an inner-value invariant violation;
• <=/>= means every exact boundary hit returns the bound (derivative 0) where Base's x > hi ? hi : x < lo ? lo : x returns x (derivative 1);
• signed zero: clamp(-0.0, 0.0, 1.0) returns +0.0 (derivative 0); Base returns -0.0 (derivative 1). Crossed bounds (lo > hi) also resolve differently from Base, so forward and reverse disagree on the primal itself.