[Bug] Fix the reverse-mode gradient of a loop-carried vector rebuilt across a dynamic loop.#799
Conversation
`determine_ad_stack_size` runs before the scalarize pass. For an adaptive tensor-typed adstack it records a symbolic `size_expr` and leaves `max_size` at the seed value of 1 for the runtime to overwrite from that tree. When scalarize splits the tensor stack into one scalar `AdStackAllocaStmt` per component it copied `max_size` but dropped `size_expr`, so the launcher found no tree and sized the component heap to the seed. A scalar component stack pushed inside a runtime-bounded loop then overflowed as soon as the loop ran more than once - a loud assertion on SPIR-V, silently corrupted gradients on a `__debug__`- disabled build. Clone `size_expr` onto every scalar component. The components are pushed and popped in lockstep with the tensor stack, so they share its depth exactly. Fixes #793.
|
…across a dynamic loop. A multi-component `qd.Vector` carried across a field-bounded loop and rebuilt from its own previous components produced a wrong reverse-mode gradient (the second-order term scaled by a spurious extra factor), silently on `__debug__`-disabled builds. The scalar-valued analogue was already correct. Two defects, both specific to the tensor component read: 1. The adstack judge (`AdStackAllocaJudger`) decides a loop-carried local needs an adstack via the "loaded AND stored inside a dynamic loop" rule, keyed on `local_loaded_`. That flag was set only for a direct `LocalLoad` of the alloca. A tensor local is read component-wise through a `MatrixPtr`, so the flag never fired and the loop-carried tensor stayed a single overwrite-each-iteration slot with no per-iteration history. Match the `MatrixPtr`-rooted load so tensor locals promote to an adstack like scalars do. 2. `ReplaceLocalVarWithStacks` subscripts a stack-backed tensor top by materializing it into a fresh plain alloca (so the forward store-to-load walker can trace it). That alloca is not recomputable, so `BackupSSA` spilled the component read to a single overwrite-each-iteration slot and every reverse iteration read the last forward iteration's component. Recognize this materialization shape and reconstruct the component in reverse by cloning the recomputable full-tensor top per iteration and re-subscripting. Fixes #797.
…ebuilt across a dynamic loop.
8e17bda to
d42a9ab
Compare
|
MacBook Pro: |
|
Production CI: |
|
Genesis PR #2842 (94f257277) x quadrants PR #799 (d42a9ab) MacBook Pro: |
|
ok to merge |
2 participants
Issue: #797
Stacked on #796 (the scalarize
size_exprpropagation): promoting a loop-carried tensor to an adstack relies on that fix to size the scalar component stacks. Review/merge #796 first; this PR's base will retarget tomainonce #796 lands.Brief Summary
copilot:summary
Walkthrough
copilot:walkthrough