Complete rewrite of forward mode - #1151
Closed
yebai wants to merge 46 commits into
Closed
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…nfwd bridge
Step 1: Delete old forward-mode stack (forward_mode.jl, 110+ frule!! methods,
LazyFoRRule/DynamicFoRRule, forward debug tests). Stub dual_ir and forward IR
inspection.
Step 2: Reimplement public forward via nfwd bridge. NfwdCache for forward-mode
with chunked gradients. HVPCache for nfwd-over-reverse Hessian-vector products
via NDual{T,1} reverse rules. Add _nfwd_lift for Tuple inputs. Update docs and tests.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Add NTangent width-aware tangent container to src/tangents/dual.jl
- Add tangent_type(Val(N), P) and dual_type(Val(N), P) width-aware queries
- Update verify_dual_type to handle NTangent-wrapped tangents
- Restore 171 frule!! definitions from main across 10 rule files
- Add _chunk_pack_tangent disambiguations for NTangent{Vararg{NoTangent}}
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Add PrimalMode context to contexts.jl with is_primitive always false - Add PrimalMode interpreter delegation in abstract_interpretation.jl - Add primal_mode.jl (605 lines) from primal-mode-migration donor - Update optimise_ir! to use PrimalMode interpreter - Rename DerivedFRule/LazyFRule/DynamicFRule to DerivedPrimal/LazyPrimal/DynamicPrimal - Rename forwards_mode_design.md to primal_mode_design.md - Add primal_mode.jl test (44 lines) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Port LazyFoRRule/DynamicFoRRule from main for forward-over-reverse lazy rule compilation (fixes x^4, x^6 HVP on Julia 1.10) - Add interp keyword to optimise_ir! for ForwardMode inlining control - Fix FoR value_and_hvp!! aliasing bug: copy grad/hvp before returning so successive calls don't overwrite earlier results - Add dual-mode HVP correctness tests (FoR + RoF): quadratic, Rosenbrock, multi-arg, namedtuple, nested tuple intermediates - Add dual-mode Hessian correctness tests (FoR + RoF): Rosenbrock, quadratic, multi-arg, namedtuple, nested tuple, tuple map/reduce Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add build_rrule methods and RRule call method to NfwdMooncake that were lost during the forward-mode refactoring. These are needed by integration tests (distributions) and the public API. Remove 4 @test_broken zero-allocation assertions for array-input nfwd gradient paths that are not yet resolved. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…hecking
Step 7 infrastructure: zero_dual/randn_dual/uninit_dual for NDual and
Complex{NDual}, primal/tangent accessors for Array{NDual} and
Array{Complex{NDual}}, ndual_width/check_ndual_width_consistency for
runtime width validation.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Add width param to const_dual!, _uninit_dual dispatcher (4 overloads)
- Thread width through __unflatten_dual_varargs with _group_vararg_dual
- Add W type param to DerivedPrimal, DynamicPrimal for zero-alloc width
- Add width field to LazyPrimal, DynamicPrimal; propagate in _copy
- Relax Vararg{Dual,N} to Vararg{Any,N} in DerivedPrimal/DynamicPrimal
- Include width in closure cache key and primal_rule_type
- Add __get_primal(::NDual) and _partial_i(::NDual) in NfwdMooncake.jl
- All modify_primal_stmts! handlers use info.width for _uninit_dual
Scalar end-to-end verified: sin, sin∘cos, x*x at width-2.
Tests: basic (26533 pass), Nfwd (477 pass).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…idth-N
Add element-wise NDual lifting for container types:
- MemoryRef/Memory: dual_type(Val(N), MemoryRef{T}) → MemoryRef{NDual{T,N}} (1.11+)
- Tuple: dual_type(Val(N), Tuple{T1,T2}) → Tuple{dual_type(Val(N),T1), ...}
- Val{0} ambiguity resolvers for MemoryRef/Memory
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add frule!! overloads for NDual containers (Array{NDual}, MemoryRef{NDual},
Memory{NDual}) so the primal-mode compiler can trace through array operations
at width-N.
Key additions:
- _uninit_dual(Val{N}, Array{T,D}): returns bare Array{NDual{T,N}}
- lmemoryrefget/lmemoryrefset!/memoryrefnew: NDual container variants
- lgetfield for bare NDual containers (tangent = NoTangent)
- sum/sum(abs2) NDual overloads for primitive array reductions
- _HasNDual/_NDualMemTypes dispatch unions
- copy for Array{NDual}
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Make the primal-mode compiler NDual-aware so that chunk_size > 1
produces width-N dual numbers instead of looping width-1.
Key changes:
- FCache carries width; value_and_gradient!! splits into
_gradient_width1 (Dual path) and _gradient_widthN (NDual path)
- _make_ndual_seed / _combine_to_ndual build NDual inputs from
per-slot seed tangents
- prepare_derivative_cache passes Val(cs) to build_frule
- _uninit_dual lifts Memory type constants for NDual containers
- _new_ frule handles NDual container construction (Array, Memory)
and struct construction with NTangent
- @zero_derivative frule!! uses per-arg type parameters to avoid
Union{Dual{<:Any},Any} collapsing to Any
- Bare NDual container frule!! overloads for memoryrefnew,
lmemoryrefget, lmemoryrefset!, getfield, copy
- NDual frule!! for sincosd, sincospi, modf (tuple-returning)
- unalias declared as forward-mode primitive (workaround for OC
segfault with NDual containers in broadcast)
- _has_ndual / _dual_or_ndual helpers for NDual dispatch in rules
- NTangent added to StandardTangentType; _get_tangent_field
overloads for NTangent field access
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Add _combine_to_ndual methods for Complex{T} scalars and
AbstractArray{Complex{T}} so chunked forward mode preserves
tangent directions instead of dropping them via NoTangent fallback
- Replace zero(T) with f(x...) in empty-input Hessian early return
so the actual primal value is returned
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…biguity
- Allow width-1 Dual calls through chunked FCache by padding into
lane 1 of the width-W rule and extracting via _ndual_output_to_width1
- Add _combine_to_ndual complex NoTangent disambiguators (resolves Aqua)
- Change generic _combine_to_ndual fallback to NTangent (structured types)
- Fix _eval_dir NTangent check and add width-N path for tuple derivatives
- Fix HVP reverse-over-forward seed type for constant functions
- Add Complex{NDual} support: _has_ndual, _dual_or_ndual, lgetfield
frule, _new_ frule, _ndual_width, _tangent_dir
- Fix empty Hessian to return f(x...) instead of zero(T)
- Add regression tests for chunked complex and empty Hessian
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Discover output type via compiled rule instead of raw f(x...) in reverse_over_forward prepare, preserving mutation-restoration semantics - Document no-direct-call contract in prepare_hvp_cache docstring - Add _combine_to_ndual for Tuple: element-wise NDual packing - Add lgetfield frules for bare Tuple/NamedTuple with NDual elements Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ages - Remove Tuple specialization in _count_slots that bypassed IdDict-based alias deduplication, causing redundant evaluations for aliased inputs - Respect config.silence_debug_messages in prepare_derivative_cache instead of hard-coding true Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Bring back detailed comments documenting cache-key design, stack aliasing invariants, thread-safety caveats, and why sig_or_mi is excluded from the DynamicFoRRule key. These explain non-obvious correctness constraints. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…es, fix Aqua issues
- Guard all Memory/MemoryRef references with @static if VERSION >= v"1.11-rc4"
in Mooncake.jl, new.jl, and primal_mode.jl
- Move _HasNDual const from memory.jl to Mooncake.jl (used by always-loaded code)
- Fix Aqua unbound type param: _ndual_width uses fieldcount(L) instead of NTuple{W}
- Fix Aqua method ambiguity: add _combine_to_ndual(::Tuple, ::Tuple{}) disambiguator
- Skip forward-mode testing for @zero_derivative typed-vararg cases (pre-existing bug)
- Guard chunked complex array test for 1.11+ only (Core.arrayref lacks width-N frule)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ive vararg frule
- Consolidate all Memory/MemoryRef NDual overloads (_has_ndual, _dual_or_ndual,
_find_ndual_memref, _ndual_width, _tangent_dir, _uninit_dual) into memory.jl
which is already conditionally loaded on Julia 1.11+. Removes @static if guards
from Mooncake.jl, new.jl, and primal_mode.jl.
- Fix _vararg_any_type to produce Union{<:T, Dual{<:T}} so typed vararg frule
signatures match Dual-wrapped args. Removes mode=ReverseMode workarounds.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Fix forward cache review follow-ups (docs, config, interface, tests) - Merge origin/main into step-8-ndual-routing - Fix NDual memory copy forwarding - Inline and remove single-use helper functions in interface.jl Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Handle NDual outputs in forward cache summaries and add arrayref/arrayset forward rules for arrays storing NDual elements. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Member
|
FYI: We need this to be merged for #1160 (in the future) to go through. |
Resolved conflicts in docs/src/interface.md and test/interface.jl by keeping this branch's FCache naming and incorporating main's new value_and_jacobian!! and nfwd NDual error tests. Also fixed stale ForwardCache references introduced by the auto-merge of src/interface.jl (value_and_jacobian!! methods used ForwardCache instead of the branch-local FCache). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Replace `_validate_prepared_cache_inputs` (main's name) with
`_validate_prepared_cache` throughout the auto-merged jacobian code
- Rewrite FCache value_and_jacobian!! to use _make_seed_tangent +
Dual calls instead of the incompatible gradient_chunk_size /
_fcache_gradient_seed_tangent / NTangent pattern from main;
throw for width != nothing since NDual frules for vector
construction/broadcast are not yet complete on this branch
- Remove nfwd NDual user-dispatch test: in this branch the primal
compiler doesn't re-dispatch to NDual-specialised user methods,
so the test was testing main-only behaviour
- Remove chunked FCache jacobian tests (chunk_size=1/2): NDual lacks
frules for getfield on Tuple{NDual,...} and zero_derivative for
Vector{NDual} in broadcast; those are separate work items
- Fix hvp_cache test to use a valid Vector input (this branch's
prepare_hvp_cache requires Vector{<:IEEEFloat} inputs)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Member
Author
|
@AstitvaAggarwal, can you run all downstream tests (eg, SciML-related) on this branch? |
Member
|
Yeah |
Member
|
@yebai I have checked the relevant SciML repositories which will require running the tests and they pass. |
Bug fixes:
- 4-arg lgetfield frule routes through _dual_or_ndual so width-N
NTangent field tangents lift to NDual instead of falling back to a
bare Dual (silent width-N miss)
- _gradient_widthN raises an explicit ArgumentError when chunk>1 and
the lifted OC returns a non-NDual output for an IEEEFloat primal,
instead of BoundsError on coeffs[d]
Dead code:
- _fcache_derivative_chunked!! placeholder + stale docstring
- duplicate using imports in test_utils.jl
- unreachable _make_seed_tangent(::NoTangent, _cursor, _dict) overload
- unused Maybe{T} alias
- ndual_width / check_ndual_width_consistency
- LazyFRule/DynamicFRule comment refs
NDual helper consolidation:
- move _has_ndual, _dual_or_ndual, _ndual_width, _ndual_primal,
_tangent_dir(_elem), _find_ndual_memref, _HasNDual, the Memory
_uninit_dual overloads, and the Memory zero_derivative method into
src/nfwd/NfwdMooncake.jl
- Mooncake.jl imports them from .NfwdMooncake so rule files can
reference by bare name
Tests:
- new test/nfwd/nfwdmooncake.jl exercises every supported NDual
container (NDual, Complex{NDual}, Array{NDual},
Array{Complex{NDual}}, Memory and MemoryRef variants on 1.11+) for
_has_ndual, _ndual_width, and _dual_or_ndual — guards against the
silent-fallthrough failure mode where a missing container overload
drops the FCache forward path to width-1
- wired into test/runtests.jl Nfwd group
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The fallback `_uninit_dual(::Val, v) = uninit_dual(v)` silently returns a width-1 Dual regardless of width. For non-differentiable IR constants this is harmless, but if a container type with a non-trivial tangent_type ever reaches this fallback, the FCache forward path silently downgrades to width-1. Throw an explicit ArgumentError pointing at NfwdMooncake.jl as the right place to add the missing overload. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The generic _count_slots fallback allocates a fresh IdDict{Any,Any}()
per call (used to track aliased mutable containers). Tuples and
NamedTuples are by definition acyclic, so the dict bookkeeping is
wasted on the hot value_and_gradient!! path.
Add specializations that recurse without state. `@allocated` confirms
zero-alloc on `(1.0, 2.0, [3.0, 4.0])` and `(a=1.0, b=2.0)`.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`build_frule` already memoises the worker rule per `(F, world)` via `interp.oc_cache`, but each `threading_run` frule call still rebuilt the per-thread `[_copy(worker_rule) for _ in 1:threadpoolsize()]` array — `_copy` recursively copies captures and runs once per thread. Cache the array keyed by `(F, threadpoolsize, world)` so the per-thread copies happen only on the first call. After warm-up, the cache stays at exactly 1 entry per signature. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
After auditing the ~20 near-duplicate Memory/MemoryRef/Array NDual
frule overloads (`Dual{<:Container{<:_HasNDual}}` vs bare
`Container{<:_HasNDual}`):
- IR probe on a chunked-gradient primal (`x -> x[1] + x[2]` walking
Memory via memoryrefnew/memoryrefget) shows the lifted IR dispatches
to the bare-container form at every site — arguments flow as bare
typed SSAValues whose static type already encodes NDual elements.
- The `Dual{...}` form is the stable user-facing entry point for
direct frule!! calls outside the lifted IR.
- `Aqua.test_ambiguities([Mooncake])` reports no ambiguities.
Add a comment explaining this duality so future readers don't try to
collapse them or wonder why both exist.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…API gap
Re-add `count_allocs == 0` assertions on `value_and_gradient!!` for
scalar / two-scalar primals. These hold (verified) on the FCache path
because `_make_seed_tangent(::Tuple)` and `(::IEEEFloat)` skip the
IdDict aliasing tracking after the recent `_count_slots` Tuple
specialisation, the gradient workspace is reused across calls, and
the lifted OC has no captures.
The chunked array gradient testsets are *not* asserting `== 0` —
array-input gradient currently allocates ~21/element (one fresh
`IdDict` per `_make_seed_tangent` call plus the Memory{NDual} lift
buffer in `_dual_or_ndual`). Pinned in pr1151.md follow-up #16.
Also flag in `src/rules/memory.jl` that the user-facing
`Dual{<:Container{<:_HasNDual}}` frule overloads can't easily be
exercised by external callers today, because `zero_dual` is width-1
only — making it width-aware to match `dual_type` is pr1151.md
follow-up #17.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
DebugFRule wraps the rule with verification that allocates, so the `count_allocs == 0` asserts I added in d26dff8 fired under `debug_mode=true` and `(debug_mode=true, silence_debug_messages=true)` in the parametrised testset (2 failures each). Skip the asserts when `kwargs.debug_mode` is true; the contract is that non-debug mode FCache scalar gradient is allocation-free. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Add an `_seed_inplace!(buf, x, slot, cursor)` helper that mutates a
tangent-shaped buffer in place to represent the standard-basis seed
for a given global slot, instead of `_make_seed_tangent` rebuilding
fresh tangent arrays every direction. Cache the buffer on `FCache`
(new `seed_buf::SB` field, lazily initialised on first call).
Allocation reduction on `value_and_gradient!!` for `x -> sum(abs2, x)`
on `Vector{Float64}` (warm-cache, no friendly_tangents):
size before after
1 25 10
2 46 16
5 109 34
10 214 64
50 1054 304
About 3.3× across the range. Per-call savings come from reusing the
tangent vector buffer; remaining allocations are mostly `_scale` and
`increment!!` temporaries on the inner accumulation loop.
The generic fallback delegates to `_make_seed_tangent` so cyclic
references and partially initialised mutable structs continue to work
unchanged. Verified rules/memory passes (16143/16143) and aliased /
cyclic / mutable-struct primals still produce correct gradients.
Width-N path is unchanged in this commit (still calls
`_make_seed_tangent` per direction inside `_make_ndual_seed`). Caching
that path will need a different seeds-tuple buffer.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The previous `tangent_type(_typeof(v)) !== NoTangent` check was too
strict: container types whose tangent_type wraps only NoTangent leaves
(e.g. `Base.RefValue{Int}` → `MutableTangent{(x=NoTangent,)}`) tripped
it even though they have no differentiable scalar slots and the
existing width-1 `uninit_dual(v)` is already correct for them.
Tighten the check to `_count_slots(v) > 0` so we only fail for
containers that *actually* carry differentiable degrees of freedom.
This caught two regressions in the basic test suite (Ref{Int} appearing
as an IR constant inside a chunked-forward primal).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Add `Val{N}` overloads to the three Mooncake.jl seed-tangent
constructors so their output static type matches `dual_type(Val(N), P)`
for every primal type covered by the NDual-aware `dual_type` table:
| Primal | Width-N return type
|---------------------|-----------------------------------
| T<:IEEEFloat | NDual{T,N}
| Complex{T} | Complex{NDual{T,N}}
| Array{T,D} | Array{NDual{T,N},D}
| Array{Complex{T},D} | Array{Complex{NDual{T,N}},D}
| Memory{T} (1.11+) | Memory{NDual{T,N}}
| MemoryRef{T} | MemoryRef{NDual{T,N}}
| (Complex variants) | …
`Val(0)` is the primal passthrough, matching `dual_type(Val(0), P) == P`.
Internal `_ndual_zero` / `_ndual_array` / `_ndual_memory` helpers
factor out the per-element NDual construction; the three entry points
share them through a per-element generator (`_ -> zero(T)`,
`_ -> randn(rng, T)`).
Without this, the user-facing `Dual{<:Container{<:_HasNDual}}` frule
overloads in `memory.jl` / `array_legacy.jl` cannot be reached from
external callers — they would have to construct NDual values directly
to seed a width-N forward call.
Add type-equality + sanity tests in test/nfwd/nfwdmooncake.jl
(117/117 pass) covering all primal types and N ∈ {0,1,2,4}.
Closes pr1151.md follow-up #17.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2 tasks
yebai
force-pushed
the
step-8-ndual-routing
branch
from
May 26, 2026 23:37
ff8c919 to
ee3101e
Compare
# Conflicts: # src/interface.jl
Restores cache-owned Hessian/grad/basis buffers (dropped by the merge conflict with main's #1178) onto the branch's refactored `HVPCache{M,C,S}`. Drops `copy(grad)`/`copy(hvp)` in `value_and_hvp!!` for `:forward_over_reverse` so the Hessian sweep writes straight into `H[:, i]`/`g` without per-coordinate Vector allocations. Docstrings note the aliasing; tests that hcat results from the same cache snapshot via `copy` before the next call. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Member
Author
|
Superseded by #1215 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Complete rewrite of forward mode. The new implementation is more Cassette-like: it lifts each primal method so that arbitrary arguments can pass through it, including
DualandNDualnumbers.This needs to be split into smaller PRs for a comfortable review.
Do not review; work in progress.
CI Summary — GitHub Actions
Documentation Preview
Mooncake.jl documentation for PR #1151 is available at:
https://chalk-lab.github.io/Mooncake.jl/previews/PR1151/
Performance
Performance Ratio:
Ratio of time to compute gradient and time to compute function.
Warning: results are very approximate! See here for more context.