Skip to content

Commit 60528bc

Browse files
committed
[AutoDiff] Disable EliminateRecomputableAdStackPushes pending mutated-SNode chain-leaf fix
1 parent 9977672 commit 60528bc

1 file changed

Lines changed: 7 additions & 8 deletions

File tree

quadrants/transforms/auto_diff/auto_diff.cpp

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -37,14 +37,13 @@ void auto_diff(IRNode *root, const CompileConfig &config, AutodiffMode autodiff_
3737
replace_local_var_with_stacks(ib, config.ad_stack_size);
3838
type_check(root, config);
3939

40-
// Drop AdStackAllocas whose pushed value is recomputable from already-stack-backed allocas + args + const +
41-
// loop-index. Trades forward-pass adstack memory traffic (one push + one load per iter per intermediate spill)
42-
// for cloned arithmetic in the reverse scope, which `BackupSSA::generic_visit` generates on demand via the same
43-
// RecomputableChainCloner path. Must run after `replace_local_var_with_stacks` (so the analyzer sees the
44-
// AdStackAlloca shape, not the alloca-with-store shape promote_ssa_to_local_var emits) and before
45-
// `make_adjoint` (so the reverse pass is generated against the cleaned forward IR with no spurious push/load
46-
// scaffolding).
47-
eliminate_recomputable_ad_stack_pushes(ib);
40+
// Disabled: `RecomputableChainAnalyzer` admits `GlobalLoadStmt` as a recomputable interior node without
41+
// verifying SNode read-only-ness. When the chain reaches a `GlobalLoadStmt` of an SNode written elsewhere in
42+
// the same kernel execution, `BackupSSA`'s reverse-side re-clone reads post-write state instead of the iter-k
43+
// value the original `AdStackLoadTopStmt` would have returned, silently corrupting gradients on MPM-style
44+
// kernels that mix per-particle reads with grid writes. Re-enable once the chain analyzer rejects
45+
// mutated-SNode loads (or the cost model gates on read-only-verified leaves).
46+
// eliminate_recomputable_ad_stack_pushes(ib);
4847
type_check(root, config);
4948

5049
make_adjoint(ib);

0 commit comments

Comments
 (0)