fix(ir): Gather the operand's split axis at the V->C boundary - #2129
Merged
lyfne123 merged 2 commits intoJul 27, 2026
Merged
Conversation
|
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: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthrough
ChangesAUTO vector-to-cube boundary
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
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: 4272b02a1d
ℹ️ 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".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/ir/transforms/lower_auto_vector_split_pass.cpp`:
- Around line 339-349: Update the C→V shard-result construction that initializes
TileInfo to also store the computed split_dim, rather than leaving split_dim at
its default value. Preserve the existing HalfDimExtent(tt->shape_[split_dim])
initialization while ensuring LEFT_RIGHT shards retain split_dim == 1 through
the subsequent V→C gather path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 85f4d9ae-da9c-473d-888c-ac55291ba8d9
📒 Files selected for processing (4)
docs/en/dev/passes/18-lower_auto_vector_split.mddocs/zh-cn/dev/passes/18-lower_auto_vector_split.mdsrc/ir/transforms/lower_auto_vector_split_pass.cpptests/ut/ir/transforms/test_lower_auto_vector_split.py
Hzfengsy
added a commit
to Hzfengsy/pypto
that referenced
this pull request
Jul 24, 2026
- Carry split_dim into the C->V shard result's TileInfo. The V->C arm now gathers along the operand's tracked split_dim, so a shard result fed straight into a V->C boundary under LEFT_RIGHT was gathering dim 0 (split_dim defaulted to 0) instead of dim 1 — doubling rows [128,64] -> [256,64] and tripping the shape invariant. Flagged by both review bots. - Add a regression test for the LEFT_RIGHT direct-shard->gather flow (the per-op halving path the cross-core ST already covers seeds split_dim correctly; this exercises the shard arm's own seeding).
LowerAutoVectorSplit rewrites a vector->cube tile.move into tile.aic_gather plus the original cube-placement move, but kept that move's original result type without checking the gather actually reassembled back to it. Two ways that broke: - An un-halved operand (a Vec parameter moved straight to cube, or a tile whose split dim is a singleton the affinity gate preserves) has no half to gather. Doubling it produced a [256, 128] operand under a move still typed [128, 128]; tile.move is shape-preserving, so that IR did not survive print->parse. Now rejected with an actionable ValueError. - The gather doubled the *function* split axis, but a tile.reshape can migrate the split axis (the rms_norm [N,1]<->[1,N] column reshape) and TileInfo::split_dim tracks where it landed. A [1, 8] lane-local operand under UP_DOWN gathered to [2, 8] where the move expected [1, 16]. The gather now follows the operand's tracked split dim (dim 0 -> split=1, dim 1 -> split=2), so the migrated case lowers correctly instead of emitting a contradiction. An internal postcondition asserts the gather result matches the kept move type, so any residual mismatch fails loudly rather than silently. Convert the pass's transform-output tests to the mandated @pl.program Before/Expected style. The tile-level IR this pass consumes and produces is expressible in the DSL, including the outlined pl.tile.aiv_shard(x, split=N) boundary form; memory spaces are ordinary pl.Mem.* annotations and no memref exists this early. The explicit SplitAivScopeStmt region tests stay hand-built: the parser always wraps a pl.split_aiv region in a scope that this pass rejects by design, so no DSL spelling delivers a bare region. _lower now keeps the print->parse roundtrip instrument on. That check is what surfaced the boundary defect, and the file previously suppressed it.
- Carry split_dim into the C->V shard result's TileInfo. The V->C arm now gathers along the operand's tracked split_dim, so a shard result fed straight into a V->C boundary under LEFT_RIGHT was gathering dim 0 (split_dim defaulted to 0) instead of dim 1 — doubling rows [128,64] -> [256,64] and tripping the shape invariant. Flagged by both review bots. - Add a regression test for the LEFT_RIGHT direct-shard->gather flow (the per-op halving path the cross-core ST already covers seeds split_dim correctly; this exercises the shard arm's own seeding).
Hzfengsy
force-pushed
the
fix/lower-auto-vector-split-vc-boundary
branch
from
July 25, 2026 04:54
caa4d8a to
ee2ed27
Compare
lyfne123
approved these changes
Jul 27, 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
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.
Summary
LowerAutoVectorSplitrewrites a vector→cubetile.moveintotile.aic_gatherplus the original cube-placement move, but kept that move's original result type without ever checking the gather actually reassembled back to it. Two ways that broke:Vecparameter moved straight to cube (or a tile whose split dim is a singleton the affinity gate deliberately preserves) has no half to gather. Doubling it produced a[256, 128]operand under a move still typed[128, 128].tile.moveis shape-preserving, so the result was IR that contradicts its own op contract and does not survive print→parse. Now rejected with an actionableValueError.tile.reshapecan migrate it — the rms_norm[N,1]↔[1,N]column reshape — andTileInfo::split_dimtracks where it landed. A[1, 8]lane-local operand underUP_DOWNgathered to[2, 8]where the move expected[1, 16]. The gather now follows the operand's tracked split dim (dim 0 → split=1,dim 1 → split=2), so this case lowers correctly instead of emitting a contradiction.An internal postcondition now asserts the gather result matches the kept move type, so any residual mismatch fails loudly instead of silently.
The parser was right to reject the old output; the pass was wrong. Per first-principles, the fix is in the pass.
Test authoring
The pass's transform-output tests move to the mandated
@pl.programBefore/Expected style. The tile-level IR this pass consumes and produces is expressible in the DSL: memory spaces are ordinarypl.Mem.*annotations, the lowered boundary has a dedicated outlined form (pl.tile.aiv_shard(x, split=N)) the printer emits for exactly this shape, and no memref exists this early (init_mem_refruns ten passes later).The explicit
SplitAivScopeStmtregion tests stay hand-built, for a verified reason: the parser always wraps apl.split_aivregion in a scope when the enclosing function ispl.FunctionType.InCore, and this pass rejects a scope-nested region by design — so no DSL spelling delivers a bare region. Verified for both top-level and loop-nested regions._lowernow keeps the print→parse roundtrip instrument on. That check is what surfaced the defect, and the file previously suppressed it via a blanketPassContext([]). One test opts out, documented against its upstream cause (OutlineIncoreScopesemits an InCore function that reads/rebinds a capturedpl.Outtensor without declaring it a parameter — pre-existing, filed separately).Test count 52 → 54: the 6 parametrized generator cases map 1:1 to named tests (a DSL
Beforecan't be parametrized over shape — annotations are read from the AST),test_store_offset_*was strengthened from a duplicate of the C→V program to a non-zero base that actually distinguishes additive from replaced offsets, plus two new regression tests for the fixes above.Testing
tests/ut— 7595 passed, 2 skippedtests/ut/ir/transforms— 2391 passed (post-rebase)tests/st/codegen/torch/test_torch_codegen_cross_core.py— 27 passed (numeric golden vs torch, both V→C modes)Documentation
docs/en/dev/passes/18-lower_auto_vector_split.mdand the zh-CN mirror document the half-operand precondition, the tracked-split-axis rule, and quote the diagnostic verbatim.