fix(ir): carry the split through an IfStmt's merge variable - #2623
Conversation
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. |
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (6)
Included review availability: Your plan provides up to 2 included reviews per hour; 1 remains after this review. 📝 WalkthroughWalkthroughThe change adds ChangesIfStmt merge repair
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to This PR corrects lane-local handling for tiles merged through conditional branches and rejects incompatible merges at compile time. No actionable merge-blocking risk remains after normal checks and review. Sequence Diagram(s)sequenceDiagram
participant LowerAutoVectorSplit
participant IfStmtBranches
participant RepairIfReturnVars
participant TileStore
LowerAutoVectorSplit->>IfStmtBranches: Lower then and else bodies
IfStmtBranches-->>RepairIfReturnVars: Return lowered Yield tile information
RepairIfReturnVars->>RepairIfReturnVars: Validate compatible branch tile info
RepairIfReturnVars-->>LowerAutoVectorSplit: Return retyped tracked merge variable
LowerAutoVectorSplit->>TileStore: Use repaired variable with lane offset
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Linked Issues checkExplanation The PR satisfies issue Full details: Docstring CoverageExplanation Docstring coverage is 43.75% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 16 functions across 4 files. (2 skipped: 2 unsupported.) 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 |
![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: 4b0089ea9f
ℹ️ 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".
f909d0a to
1c707bc
Compare
LowerAutoVectorSplit rebuilds and tracks a ForStmt's loop-carried state so
a halved carry stays lane-local, but did the same for no IfStmt. Both
branch bodies were lowered and their Yield values substituted, while
IfStmt::return_vars_ kept its original full-width type, was never added
to tile_vars, and got no old->new replacement.
Two things went wrong at once. The merge variable contradicted both Yield
values -- the same declared-type-versus-operand defect as gh#2203, on the
merge path -- and because it stayed untracked, a following tile.store got
no lane offset:
v: pl.Tile[[64, 128], ...] = pl.tile.load(data, [0 + subblock_idx * 64, 0], ...)
if flag > 0:
merged: pl.Tile[[128, 128], ...] = pl.yield_(doubled) # branches are halved
else:
merged: pl.Tile[[128, 128], ...] = pl.yield_(v)
out_store = pl.tile.store(merged, [0, 0], out_0) # BOTH lanes write row 0
so the two AIV lanes overlapped their writes instead of taking a half each.
Repair the merge from the lowered branches, the way the ForStmt arm
repairs its carry from the iter_args: retype to the per-lane extent,
register in tile_vars, and record the replacement. A Yield value still
references the pre-halving var and the halving path registers both the old
and the new var in tile_vars, so looking the yielded value up there is what
tells a lane-local branch value from a shared one.
The branches must agree. One yielding a halved value while the other
yields a full-width one has no single merge type: halving is wrong for the
else path and not halving is wrong for the then path, and picking either
silently gives one lane the wrong extent. Reject that with a diagnostic
naming the merge variable instead.
Both branches always exist. SSAVerifier::VerifyIfStmt rejects an IfStmt
that defines return_vars without an else branch, and ConvertToSSA
synthesizes the else Yield for a source-level no-else phi -- so the repair
needs no else-less case, and asserts the invariant instead of supporting a
shape that cannot legally reach this pass.
Both the explicit path (split_axis::ProcessStmt) and the AUTO
affinity-gated arm (LowerStmts) need the call, for the same reason the loop
carries needed two call sites: the AUTO arm recurses through its own walk
and cannot reach ProcessStmt's IfStmt branch.
Closes hw-native-sys#2608
2 participants
Summary
LowerAutoVectorSplitrebuilds and tracks aForStmt's loop-carried state so a halved carry stays lane-local, but did the same for noIfStmt. Both branch bodies were lowered and theirYieldvalues substituted, whileIfStmt::return_vars_kept its original full-width type, was never added totile_vars, and got no old→new replacement.Two things went wrong at once:
Yieldvalues — the same declared-type-versus-operand defect as [Pass Bug] LowerAutoVectorSplit halves a vector op's result without sharding its full-width operand #2203, on the merge path;tile.storegets no lane offset, so the two AIV lanes overlap their writes instead of taking a half each.Behavior change: a branch-merged tile inside an automatically split region is now halved and lane-offset correctly; a merge whose branches disagree is rejected at compile time with a
ValueError. No in-contract program changes behavior. No interface or migration step.Changes
src/ir/transforms/utils/split_axis_utils.{h,cpp}: addRepairIfReturnVars, mirroring whatRepairIterArgs/RepairReturnVarsdo for loop carries — retype the merge to the per-lane extent, register it intile_vars, record the replacement. AYieldvalue still references the pre-halving var, and the halving path registers both the old and the new var intile_vars, so looking the yielded value up there is what tells a lane-local branch value from a shared one.src/ir/transforms/utils/split_axis_utils.cpp: call it fromProcessStmt'sIfStmtbranch.src/ir/transforms/lower_auto_vector_split_pass.cpp: call it from the AUTO affinity-gated arm too — that arm recurses through its own walk and cannot reachProcessStmt'sIfStmtbranch, the same reason the loop carries needed two call sites.docs/{en,zh}/dev/passes/21-lower_auto_vector_split.md: the "Carries and dropped axes" section becomes "Carries, merges and dropped axes" and documents the merge rule.tests/ut/ir/transforms/test_lower_auto_vector_split.py: three cases (below).The two rules the merge needs
Branches must agree. One yielding a halved value while the other yields a full-width one has no single merge type: halving is wrong for the else path, not halving is wrong for the then path, and picking either silently gives one AIV lane the wrong extent. Rejected, with the merge variable named.
An
ifwith noelseis not a disagreement.return_vars_is aDefField, so an else-lessifdefines the merge name on the taken path only — its one branch is the sole constraint on the type, and following it is what makes that path well-typed. (Verified that shape parses and reaches the pass; it does.)Verification
Run at the pushed commit, in a worktree-local build:
pytest tests/ut/ir/transforms/ tests/ut/ir/operators/ tests/ut/codegen/ -q -n 16 -p no:randomly: 6050 passed, 2 skipped.tests/lint/checks: all passed.clang-format --dry-run --Werror,ruff check,ruff format --checkon the changed sources: passed.Two pre-existing failures on this machine, both reproduced with every change in this PR reverted and therefore unrelated:
tests/ut/language/test_unified_ops.py::...::test_symlinked_import_path_still_names_the_caller— the editable install's meta-path finder hijacks the subprocess's symlinkedimport pypto.tests/ut/codegen/test_orchestration_codegen_graph.py::test_generated_orchestration_compiles_against_the_pinned_runtime— the localruntimesubmodule is checked out atdbdd041ewhile the branch pins15f5cbd9, so the generated file is type-checked against the wrong headers (static assertion failed: all arguments must be simpler::hbg::Tensor). The submodule is deliberately left unstaged.Tests added
test_if_merge_variable_is_retyped_and_tracked— both branches yield a halved value; asserts the merge is[64, 128]and the store carries+ subblock_idx * 64. Fails onmainwith the merge left at[128, 128]and the store at[0, 0].test_if_branches_disagreeing_on_the_merge_are_rejected— one halved branch, one full-width parameter; asserts the diagnostic namesmerged.test_if_without_else_takes_the_merge_type_from_its_only_branch— pins the else-less rule.No system tests were run locally; CI covers them.
Closes #2608