[Fixup] Elim-pushes: detect prologue init push by position, not by zero value, to preserve real x = 0.0 body stores

Author: duburcqa

quadrants/transforms/auto_diff/auto_diff_common.h

@@ -61,14 +61,11 @@ class NonLinearOps {
 };
 
 // ----------------------------------------------------------------------------
-// Recomputable chain analyzer + cloner: decide whether a forward SSA value can
-// be reconstructed in the reverse-pass scope from already-stack-backed allocas,
-// kernel args, constants, and loop indices, and clone the DAG at a target
-// insertion point. Cross-stage shared infrastructure: used by
-// `EliminateRecomputableAdStackPushes` (forward_state_spill stage) to drop
-// pushes whose values are recomputable, and by `BackupSSA::generic_visit`
-// (post_adjoint_cleanup stage) to clone such chains in place of cross-block
-// SSA reads.
+// Recomputable chain analyzer + cloner: decide whether a forward SSA value can be reconstructed in the reverse-pass
+// scope from already-stack-backed allocas, kernel args, constants, and loop indices, and clone the DAG at a target
+// insertion point. Cross-stage shared infrastructure: used by `EliminateRecomputableAdStackPushes`
+// (forward_state_spill stage) to drop pushes whose values are recomputable, and by `BackupSSA::generic_visit`
+// (post_adjoint_cleanup stage) to clone such chains in place of cross-block SSA reads.
 // ----------------------------------------------------------------------------
 
 // Returns true iff `stmt`'s transitive operand DAG terminates at recomputable leaves via side-effect-free
@@ -107,9 +104,7 @@ class RecomputableChainAnalyzer {
     // Recomputable leaves: ConstStmt, ArgLoadStmt, AdStackLoadTopStmt, AdStackAllocaStmt. LoopIndexStmt is
     // intentionally excluded - cloning a LoopIndexStmt copies the reference to the forward RangeForStmt,
     // but the cloned consumer lives inside the reverse RangeForStmt (a separate stmt with its own loop
-    // index), so the cloned read points at undefined state and silently double-accumulates gradients
-    // (`test_adstack_sum_linear`). A future MakeAdjoint coordination pass that tracks
-    // forward-RangeFor-to-reverse-RangeFor mapping could lift this restriction.
+    // index), so the cloned read points at undefined state and silently double-accumulates gradients.
     //
     // AdStackLoadTopStmt as a leaf is correct under the dominance + control-flow-consumer + self-load
     // guards in `EliminateRecomputableAdStackPushes::run_one_pass`. The reasoning:
@@ -146,11 +141,9 @@ class RecomputableChainAnalyzer {
     // before reverse runs, so the reverse re-read sees the kernel's final post-write state. If a global is
     // mutated by the forward and then re-read by the reverse clone, the values can differ.
     //
-    // Most rigid-step kernel chain leaves are reads of input parameters (mass, inertia, joint params,
-    // morphology) that the kernel does NOT write. For those, the re-read returns the same value.
-    // Conservative analysis: a future safety check could prove the SNode is read-only in the kernel; until
-    // then this path relies on the test suite's gradient-correctness asserts to surface any mutated-global
-    // miscompilation.
+    // Chain leaves that pass this branch are typically reads of kernel-input parameters not written by the
+    // kernel; for those, the re-read returns the same value as the forward read. The pass does not statically
+    // verify SNode read-only-ness, and mutated-global cases would silently produce wrong gradients.
     bool is_interior = stmt->is<UnaryOpStmt>() || stmt->is<BinaryOpStmt>() || stmt->is<TernaryOpStmt>() ||
                        stmt->is<MatrixPtrStmt>() || stmt->is<GlobalPtrStmt>() || stmt->is<ExternalPtrStmt>() ||
                        stmt->is<GlobalLoadStmt>();
@@ -214,9 +207,8 @@ class RecomputableChainCloner {
 };
 
 // ----------------------------------------------------------------------------
-// ADTransform: shared base for reverse-mode (MakeAdjoint) and forward-mode
-// (MakeDual) IR builders. All methods are inline so derived classes in
-// separate translation units can use them without ODR concerns.
+// ADTransform: shared base for reverse-mode (MakeAdjoint) and forward-mode (MakeDual) IR builders. Methods are inline
+// so derived classes in separate translation units can use them without ODR concerns.
 // ----------------------------------------------------------------------------
 class ADTransform : public IRVisitor {
  protected:

quadrants/transforms/auto_diff/eliminate_recomputable_pushes.cpp

@@ -4,27 +4,23 @@ namespace quadrants::lang {
 
 class Block;
 
-// Stage: Forward-state spill. Pre-MakeAdjoint passes that prepare the forward
-// IR so the reverse pass can re-read the per-iteration values it needs.
+// Stage: Forward-state spill. Pre-MakeAdjoint passes that prepare the forward IR so the reverse pass can re-read the
+// per-iteration values it needs.
 
-// Hoist forward-pass SSA defs that the reverse pass will re-read (operands of
-// non-linear ops, indices of GlobalPtr / ExternalPtr, range-for bounds, if
-// conds) into AllocaStmt + LocalStore + LocalLoad. Demand-driven: defs no
-// reverse formula reads stay in pure SSA form.
+// Hoist forward-pass SSA defs that the reverse pass will re-read (operands of non-linear ops, indices of GlobalPtr /
+// ExternalPtr, range-for bounds, if conds) into AllocaStmt + LocalStore + LocalLoad. Demand-driven: defs no reverse
+// formula reads stay in pure SSA form.
 void promote_ssa_to_local_var(Block *block);
 
-// Replace each AllocaStmt the reverse pass needs to read across iterations
-// with an AdStackAllocaStmt + AdStackPushStmt / AdStackLoadTopStmt. The
-// AllocaJudger (file-private) decides which allocas need the promotion.
+// Replace each AllocaStmt the reverse pass needs to read across iterations with an AdStackAllocaStmt +
+// AdStackPushStmt / AdStackLoadTopStmt. The AllocaJudger (file-private) decides which allocas need the promotion.
 void replace_local_var_with_stacks(Block *block, int ad_stack_size);
 
-// Drop AdStackAllocaStmts whose pushed value is recomputable from already-
-// stack-backed allocas + kernel args + constants + loop indices. Trades
-// forward-pass adstack memory traffic for cloned arithmetic in the reverse
-// scope, which `BackupSSA::generic_visit` synthesizes on demand via the
-// shared RecomputableChainCloner. Must run after `replace_local_var_with_stacks`
-// (so the analyzer sees the AdStack shape) and before `make_adjoint` (so the
-// reverse pass is generated against the cleaned forward IR).
+// Drop AdStackAllocaStmts whose pushed value is recomputable from already-stack-backed allocas + kernel args +
+// constants + loop indices. Trades forward-pass adstack memory traffic for cloned arithmetic in the reverse scope,
+// which `BackupSSA::generic_visit` synthesizes on demand via the shared RecomputableChainCloner. Must run after
+// `replace_local_var_with_stacks` (so the analyzer sees the AdStack shape) and before `make_adjoint` (so the reverse
+// pass is generated against the cleaned forward IR).
 void eliminate_recomputable_ad_stack_pushes(Block *ib);
 
 }  // namespace quadrants::lang