[AutoDiff] Adstack chain-clone safety: mutated-SNode leaf reject + load_top consumer-aware guard

[duburcqa]

Adstack chain-clone safety: mutated-SNode leaf reject + load_top consumer-aware guard

Closes two distinct miscompiles in reverse-mode adjoint where BackupSSA re-clones a recomputable forward chain at the reverse cursor and the cloned chain reads a value that has since been overwritten. Re-enables EliminateRecomputableAdStackPushes, which #633 disabled as a stop-gap.

TL;DR

BackupSSA::generic_visit re-clones a recomputable forward chain in place of an alloca spill on the reverse side. Two shapes break that contract today:

1. The chain bottoms out in a GlobalLoadStmt of a SNode the same kernel writes elsewhere. The cloned load reads post-write state instead of the iter-k value the original AdStackLoadTopStmt would have returned.
2. The chain bottoms out in an AdStackLoadTopStmt whose stack receives a subsequent forward push in a strict ancestor scope of the consumer. The matching reverse AdStackPopStmt lands after the consumer's reverse code, so the cloned load_top reads the post-push top.

Both are silent NaN / inf gradient corruption on otherwise-valid forward kernels. This PR adds the two missing eligibility checks and re-enables EliminateRecomputableAdStackPushes (ERAP) which was disabled in #633 pending exactly this fix.

Why

EliminateRecomputableAdStackPushes is the load-bearing ERAP pass for the dynamic-loop adjoint perf path: dropping every spill whose value is recomputable from already-stack-backed allocas + args + const + loop-index turned the rigid step from 70 to 143 FPS. #633 disabled it after a Genesis test_differentiable_push[cpu] failure surfaced the chain-clone correctness gap. Re-enabling unblocks downstream perf work without paying the correctness tax.

RecomputableChainAnalyzer and BackupSSA::generic_visit share a single notion of recomputability today; both must agree on which leaves are clone-safe at the reverse cursor. The two new guards close the only known cases where they disagreed.

Surface API

None. All changes are internal to quadrants/transforms/auto_diff/. No new flags, no new env vars, no API churn for callers of qd.kernel.grad or qd.ad.Tape.

Mechanism

Two orthogonal eligibility tightenings, both layered on top of the existing recomputable-chain analysis without changing the cost model.

1. Mutated-SNode chain-leaf reject (auto_diff_common.h, eliminate_recomputable_pushes.cpp, post_adjoint_cleanup.cpp)

RecomputableChainAnalyzer::collect_mutated_snodes(Block *ib) walks the IB once and returns the set of SNodes targeted by any GlobalStoreStmt or AtomicOpStmt. is_recomputable(stmt, cache, mutated_snodes) rejects a GlobalLoadStmt chain leaf if its GlobalPtrStmt's SNode is in the mutated set (or null). ERAP and BackupSSA both pass the mutated set into the analyzer.

This rules out the post-write read shape. The leaf stays adstack-backed, ERAP keeps the original push/pop, and the reverse pops the iter-k value verbatim.

2. Consumer-aware load_top clone guard (post_adjoint_cleanup.cpp)

BackupSSA::is_load_top_safe_for_consumer(load_top, consumer) returns false iff some forward push P to the same stack strictly contains consumer's parent block AND precedes consumer's enclosing-scope stmt in P's parent block (so pop_P lands after consumer's reverse code). Document order is captured in a forward_preorder_ index built once at construction.

is_chain_clone_safe(op, consumer) walks the recomputable DAG and applies the same predicate to every AdStackLoadTopStmt leaf, gating the chain-clone path. The bare AdStackLoadTopStmt operand path uses the same predicate directly. On rejection, both paths fall back to the per-IB alloca spill, which captures the forward-time SSA value verbatim.

Tests

Two new regression tests in tests/python/test_adstack.py:

• test_backup_ssa_load_top_with_subsequent_pushes pins guard #2 (consumer-aware load_top): without the guard, BackupSSA clones a bare AdStackLoadTopStmt whose stack receives a later push in a strict ancestor scope, the reverse reads the post-push top, and the gradient blows up.
• test_eliminate_recomputable_pushes_rejects_mutated_snode_chain_leaf pins guard #1 (mutated-SNode leaf): three top-level for-loops over the same field SNode (atomic-add, gated divide-by-self, consumer); without the guard, ERAP drops the adstack carrying the leaf's value and the reverse re-clone reads field after the divide-by-self has overwritten the gated cells.

Both tests assert the analytic gradient (not just isfinite) with pytest.approx(rel=1e-6, abs=1e-6). Both fail on origin/main with ERAP re-enabled and pass with the two guards in place.

Side-effect audit

• RecomputableChainAnalyzer::is_recomputable signature gained a mutated_snodes parameter. Every call site in this repo now passes the IB-level set; no external callers exist.
• EliminateRecomputableAdStackPushes now computes mutated_snodes once per outer pass (single linear walk over the IB), so the per-pass cost is a strict additive on top of the pre-[AutoDiff] Disable EliminateRecomputableAdStackPushes pending mutated-SNode chain-leaf fix #633 path. No call-site cost in the steady state.
• BackupSSA builds forward_preorder_ once per IB at construction (single preorder walk over the IB). All is_load_top_safe_for_consumer queries are O(pushes-to-stack) with O(1) document-order lookups. Memoized on (op, consumer_parent) for the chain path.
• The new comments reference only IR concepts (independent block, AdStackPushStmt, MakeAdjoint reverse cursor) - no external project, no test, no history, no version.

[duburcqa]

@claude review

[github-actions]

Total: 5 file(s) changed, +285 -58 code lines.

[github-actions]

Diff coverage: 72% · 47 lines, 13 missing

[duburcqa]

@claude review

[hughperkins]

checklist:

• only affects adstack/autodiff (so no need for genesis unit tests or benchmarks)
• no user-facing chagnes, so no change in user-docs needed
• non-test files modified all seem reasonably-sized

=> ok to merge