fix(ir): thread a scope's store-target rename out of the loop that produced it - #2566
Merged
lyfne123 merged 2 commits intoAug 31, 2026
Conversation
…oduced it
`OutlineIncoreScopes` replaces a scope that writes a captured tensor with a call
whose result binds a *fresh* SSA name, and every later reference to that tensor
resolves to the fresh one through the function-wide `store_target_renames_` map.
When the scope sits inside a loop or an `if`, that name is bound inside the body,
so the map hands the statements *after* the control flow a Var that is out of
scope:
for i__idx_v0 in pl.range(4):
out__ssa_v1 = self.k_incore_0(a__ssa_v0, i__idx_v0, out__ssa_v0)
return out__ssa_v1 # defined inside the loop
The map is saved and restored around a *scope* body but not around a
control-flow body, and `ScopeOutliner` declared no `ForStmt` / `WhileStmt` /
`IfStmt` override at all, so the entry outlived the body that created it while
the loop header was never rewritten to carry the tensor.
Nothing miscompiled: the outlined callee returns its own `pl.Out` parameter, so
every SSA version denotes the same GM buffer. Only the def-use graph was wrong —
`SSAVerify` reports "used outside its defining scope" and `UseAfterDef` reports
"used before definition" on the result. It went unnoticed because the pass
declares `SSAForm` as produced but never invalidates it, so `PassPipeline`
memoizes the check away at pass 8.
The shape is far more general than a `manual_dep` submit: it needs neither
`manual_dep`, nor `deps=`, nor `pl.submit`, and it reaches the `if` form (where
the escaping definition exists only on the taken branch) and the ReturnStmt form
(the `serialized` example in the dependency guide).
Thread the rename out as a real carry instead. The value on entry seeds a new
`IterArg`, the body yields the fresh Var, and a new `return_var` becomes what the
following statements see. N sibling scopes writing one target share a slot;
nested loops re-thread the inner carry through the outer one; a target that
already *is* one of the loop's iter_args reuses that slot's `return_var`; an `if`
with no `else` gets one synthesised, yielding the value it came in with.
Two subtleties are handled at the code. The body rebind cannot use
`transform_utils::Substitute`, which re-resolves whatever it substitutes in: the
replacement is an `IterArg` whose `initValue_` is the Var being replaced, so
re-visiting turns the seed into a self-reference and mints a second, unbound
`IterArg`. And the `if` branches are alternatives, so the else branch restarts
from the pre-`if` map rather than inheriting the then branch's renames.
The new carry also exposed a pre-existing divergence in `ClassifyIterArgCarry`:
its alias rule filtered on the *call-site* `ArgDirection` and excluded `NoDep`,
while orchestration codegen aliases a call's result by the *callee's*
`ParamDirection`, which `pl.at(no_dep_args=[...])` never touches. A `no_dep`
carry therefore classified as `rebind`, and codegen materialised a fresh
`TaskTensor` for a slot it was simultaneously aliasing to the arg. `NoDep` is an
ordering claim, not an identity one, so it now counts as output-side at both
sites.
With that, generated code is byte-identical to main for the whole dependency
guide: the carries land in the iter_arg's alias class and classify trivial, so
they are SSA bookkeeping and cost nothing at runtime.
|
Warning Review limit reachedNext included review available in 50 minutes. View limit detailsLimit details: You’ve used all 2 included reviews currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Review configuration: ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (9)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2a15c7b967
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
… rename map
Review follow-ups on the store-target carry.
**The alias rule must read the callee, not the call site.** The first cut
admitted `ArgDirection::NoDep` as output-side, on the reasoning that `no_dep`
suppresses ordering rather than identity. That is true of the direction but not
of the slot: `pl.at(no_dep_args=[...])` accepts *any* tensor the scope captures,
read or written, and `DeriveCallDirections` overwrites the slot unconditionally.
A read-only capture therefore lands as `NoDep` too, and the positional
`TupleGetItemExpr` walk then counts it as the first output:
pl.submit(self.main_incore_0, x__ssa_v0, scratch__iter_v1,
attrs={"arg_directions": [pl.adir.no_dep, pl.adir.inout]})
`ret[0]` aliased `x` instead of `scratch`, and the carry misclassified as
`rebind`. Read the *callee's* `ParamDirection` instead — the only thing
orchestration codegen consults (`CollectOutIndices`) — so the two agree by
construction rather than by a call-site approximation. `CalleeReturnedArg` now
mirrors `GenerateTupleReturnAliases`: prefer the explicit returned-param map,
fall back to the position-th written param, and range-check against `args_` so a
Submit's runtime-allocated params (which have no caller arg) drop out.
**The retarget must not scan the whole map.** Each control-flow node swept every
entry of the function-wide `store_target_renames_`, which keeps growing, so a
function with N sequential loops writing N distinct tensors cost O(N^2) —
over the pass-complexity bound. Add `renamed_by_value_`, a reverse index from a
value to the keys currently resolving to it, and visit only the entries that
actually move. Buckets may hold stale keys, so every read re-checks the forward
map. All writes now go through `SetStoreTargetRename` to keep the two in step.
The `if` path dropped its three whole-map copies with it: rewinding the store
targets the frame already recorded is enough to make the branches alternatives
again, and costs O(renames). The branch-local post-store aliases are keyed on
Vars the other branch cannot name, so they stay and are retargeted with
everything else.
Generated code stays byte-identical to main for the dependency guide.
lyfne123
approved these changes
Aug 31, 2026
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.
What
OutlineIncoreScopesreplaces a scope that writes a captured tensor with a call whose result binds a fresh SSA name, and every later reference to that tensor resolves to the fresh one through the function-widestore_target_renames_map. When the scope sits inside a loop or anif, that name is bound inside the body, so the map hands the statements after the control flow a Var that is out of scope.The first example in the dependency guide is enough to show it — after pass 8:
→
SSAVerify: Variable 'out__ssa_v1' used outside its defining scopeThis PR threads the rename out as a real carry: the value on entry seeds a new
IterArg, the body yields the fresh Var, and a newreturn_varis what the following statements see.Why it went unnoticed
Nothing miscompiled. The outlined callee returns its own
pl.Outparameter, so every SSA version of the tensor denotes the same GM buffer and downstream passes resolve them all to one allocation. Only the def-use graph was wrong.It stayed silent because the pass declares
SSAFormas.producedbut never.invalidated, andPassPipelinecomputesto_verify = produced ∩ verified_props − verified.ConvertToSSA(pass 4) already verifiedSSAFormand nothing between 4 and 8 invalidates it, so the re-check at pass 8 is memoized away. CallingPropertyVerifierRegistry.verify({SSAForm}, prog)directly right after pass 8 fires immediately.This is not addressed here — closing that gap is a separate call, since re-verifying
SSAFormat pass 8 may surface other latent violations. Flagging it so it is not lost.Scope of the defect
Broader than the shape it was first reported on. It needs none of
manual_dep,deps=, orpl.submit— a plainpl.create_tensorwritten by an in-looppl.atand read afterwards reproduces it byte for byte. It also reaches:ifform, which is worse: the fresh Var exists only on the taken branch, yet the read after theifis unconditional;VisitExpr_(VarPtr)rather than the call-args lookup), which is theserializedexample indocs/en/user/performance/03-dependencies.md.Root cause
store_target_renames_is flat and function-wide, and is saved/restored only around a scope body (scope_outline_utils.cpp:1125/:1148) — never around a control-flow body.ScopeOutlinerdeclared noForStmt/WhileStmt/IfStmtoverride at all, so it inheritedIRMutator's plain descent: the entry outlived the body that created it, and the loop header was never rewritten to carry the tensor.How the carry behaves
return_varifwith noelseTwo subtleties, both documented at the code:
transform_utils::Substitute, which re-resolves whatever it substitutes in so chained maps settle. Here the replacement is anIterArgwhoseinitValue_is the Var being replaced, so re-visiting rewrites the seed into a self-reference and mints a second, unboundIterArg.CarryRebindMutatormakes a carry terminal in both directions.ifbranches are alternatives, not a sequence, so the else branch restarts from the pre-ifmap rather than inheriting the then branch's renames.The
ClassifyIterArgCarrychangeThe new carry exposed a pre-existing divergence. Pass 47's alias rule filtered on the call-site
ArgDirectionand excludedNoDep, while orchestration codegen aliases a call's result by the callee'sParamDirection(CollectOutIndices), whichpl.at(no_dep_args=[...])never touches. So ano_depcarry classified asrebindand codegen materialised a freshTaskTensorfor a slot it was simultaneously aliasing to the arg:NoDepis an ordering claim, not an identity one — the callee still writes through that same tensor and returns it.IsOutputSideDirectionnow coversOutputExisting | InOut | Output | NoDepat both call sites in that pass, which also removes an inconsistency between them (theTupleGetItemExprbranch already admittedOutput; the direct-call branch did not — no pass stampsOutputon an argument today, so that half is currently unreachable).Verification
main. Built a baseline worktree atorigin/mainand diffed the complete artifacts —.pto, ptoas.cpp, AIV kernels, orchestration.cpp— for all four dependency-guide kernels (serialized,tensor_claim,narrow_claim,region_claim). Zero diff. The carries land in the iter_arg's alias class and classify trivial, so iter_arg and return_var both emit as the init value's name.10605 passed, 3 skipped, 3 xfailedacrosstests/ut/;ctestclean;pre-commitclean.SSAForm+UseAfterDef.Reviewer notes
ClassifyIterArgCarrychange is what makes that true for theno_dep_argscase — without it,narrow_claimregresses to a materialised carry.TestOutlineScopeInControlFlow(loop→consumer, loop→ReturnStmt, appended-after-an-existing-carry, nested loops, sibling scopes,ifbranch, and the unchanged top-level case), asserting onSSAForm+UseAfterDefand on the resulting carry's shape; 1 for theNoDepalias rule.08-outline_incore_scopes.mdand47-classify_iter_arg_carry.md, EN and ZH.