[Bug] Propagate the adstack size-expr onto scalarized tensor components.#796
Merged
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.
hughperkins
reviewed
Jul 20, 2026
| // `max_size`, the launcher finds no symbolic tree and sizes the heap to that seed, and a scalar stack pushed | ||
| // inside a runtime-bounded loop overflows once the loop runs more than once. | ||
| if (stmt->size_expr) { | ||
| scalar_ad_stack->size_expr = stmt->size_expr->clone(); |
Collaborator
There was a problem hiding this comment.
note to self: scoped to only adstack things because inside a visit(AdStackAllocaStmt *stmt) ✅
Collaborator
|
skimming => looks good to me
|
Contributor
Author
|
Production CI: |
Contributor
Author
|
MacBook Pro: |
Collaborator
|
ok to merge |
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.
Issue: #793
Brief Summary
copilot:summary
Walkthrough
copilot:walkthrough