fix(scatter): tscatter mask-form PTOAS syntax + cmp blayout roundtrip (#1498)#1513
Open
Little-oil wants to merge 8 commits into
Open
fix(scatter): tscatter mask-form PTOAS syntax + cmp blayout roundtrip (#1498)#1513Little-oil wants to merge 8 commits into
Little-oil wants to merge 8 commits into
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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughFixes packed-mask TileView default inference to use physical tile_shape, adjusts pto.tscatter MLIR emission (maskPattern placement and operand typing), adds an INT16 overflow check during tensor.scatter lowering, and adds/updates runtime and unit tests for mask-form scatter and conversion roundtrip. ChangesScatter operations: parser, codegen, and testing stack
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request fixes a print-to-parse roundtrip failure for packed-mask tiles by resolving implicit tile view defaults using the physical tile shape instead of the valid shape. It also implements the mask-form scatter (pto.tscatter) for A2/A3 backends, updates the PTO codegen to place the mask pattern inside the input arguments list, and adds comprehensive tests. The feedback suggests using the INTERNAL_CHECK_SPAN macro to automatically include source location information in the error message and dynamically printing the destination data type in the INT16 index range guard check to keep the error message accurate.
Little-oil
pushed a commit
to Little-oil/pypto
that referenced
this pull request
May 26, 2026
- tscatter mask form: emit maskPattern inside ins() after src, before the type
annotation (ins(%src, {maskPattern...} : src_ty) outs(%dst)) — device-verified.
- INTERNAL_CHECK -> INTERNAL_CHECK_SPAN(op->span_) for the scatter type-annotation
symmetry check (gemini).
- INT16 flat-index guard: print the actual element dtype instead of hardcoding
"2-byte" (INT16 covers 1- and 2-byte; gemini).
![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/op_conversion_registry.cpp`:
- Around line 1415-1420: Replace the unsafe multiplication CHECK(n * cols <=
32768) with a division-based bound to avoid signed int64_t overflow: define a
constant like kMaxFlat = 32767 (or 32768 per desired semantics), compute int64_t
max_rows = (cols == 0 ? kMaxFlat : kMaxFlat / cols), and then CHECK(n <=
max_rows). In the CHECK error message (the same CHECK site in
op_conversion_registry.cpp) do not recompute n*cols; instead report n, cols and
the computed max_rows (or kMaxFlat) to explain the limit and suggest splitting
the scatter.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 57ac265f-2584-4af3-bf24-4f84ecbbba35
📒 Files selected for processing (2)
src/backend/common/pto_ops_common.cppsrc/ir/transforms/op_conversion_registry.cpp
Little-oil
pushed a commit
to Little-oil/pypto
that referenced
this pull request
May 26, 2026
- INT16 scatter flat-index guard: bound rows via division (kMaxFlat/cols) instead of n*cols to avoid signed int64 overflow; handle cols==0
Little-oil
pushed a commit
to Little-oil/pypto
that referenced
this pull request
May 26, 2026
- tscatter mask form: emit maskPattern inside ins() after src, before the type
annotation (ins(%src, {maskPattern...} : src_ty) outs(%dst)) — device-verified.
- INTERNAL_CHECK -> INTERNAL_CHECK_SPAN(op->span_) for the scatter type-annotation
symmetry check (gemini).
- INT16 flat-index guard: print the actual element dtype instead of hardcoding
"2-byte" (INT16 covers 1- and 2-byte; gemini).
Little-oil
pushed a commit
to Little-oil/pypto
that referenced
this pull request
May 26, 2026
- INT16 scatter flat-index guard: bound rows via division (kMaxFlat/cols) instead of n*cols to avoid signed int64 overflow; handle cols==0
Little-oil
force-pushed
the
feat/scatter-review-fixes
branch
from
668210d to
96ca32c
Compare
Little-oil
pushed a commit
to Little-oil/pypto
that referenced
this pull request
May 27, 2026
- tscatter mask form: emit maskPattern inside ins() after src, before the type
annotation (ins(%src, {maskPattern...} : src_ty) outs(%dst)) — device-verified.
- INTERNAL_CHECK -> INTERNAL_CHECK_SPAN(op->span_) for the scatter type-annotation
symmetry check (gemini).
- INT16 flat-index guard: print the actual element dtype instead of hardcoding
"2-byte" (INT16 covers 1- and 2-byte; gemini).
Little-oil
pushed a commit
to Little-oil/pypto
that referenced
this pull request
May 27, 2026
- INT16 scatter flat-index guard: bound rows via division (kMaxFlat/cols) instead of n*cols to avoid signed int64 overflow; handle cols==0
Little-oil
force-pushed
the
feat/scatter-review-fixes
branch
from
96ca32c to
8de941a
Compare
Little-oil
pushed a commit
to Little-oil/pypto
that referenced
this pull request
May 28, 2026
- tscatter mask form: emit maskPattern inside ins() after src, before the type
annotation (ins(%src, {maskPattern...} : src_ty) outs(%dst)) — device-verified.
- INTERNAL_CHECK -> INTERNAL_CHECK_SPAN(op->span_) for the scatter type-annotation
symmetry check (gemini).
- INT16 flat-index guard: print the actual element dtype instead of hardcoding
"2-byte" (INT16 covers 1- and 2-byte; gemini).
Little-oil
pushed a commit
to Little-oil/pypto
that referenced
this pull request
May 28, 2026
- INT16 scatter flat-index guard: bound rows via division (kMaxFlat/cols) instead of n*cols to avoid signed int64 overflow; handle cols==0
Little-oil
force-pushed
the
feat/scatter-review-fixes
branch
from
8de941a to
3198967
Compare
Contributor
Author
|
blocked by PTO-ISA |
added 7 commits
May 30, 2026 16:02
…hw-native-sys#1498) The python printer elides a TileView field when it matches GetImplicitTileView(tile_type.shape_, ...) — i.e. the implicit view derived from the *physical* tile shape. The text parser, however, recomputed the implicit blayout/slayout/fractal from `valid_shape` when it was given, desynchronising the two for packed-mask tiles. A cmp/cmps result has physical shape e.g. [16, 8] but valid_shape [16, 1]: the printer omits its (row_major) blayout, while the parser saw valid_shape's cols==1 and filled col_major, so print->parse failed structural equality with "TileView blayout mismatch". Infer the implicit defaults from the physical tile shape (falling back to valid_shape only when the shape is unavailable), matching the printer. Un-skips TestConvertScatterOp::test_scatter_conversion, which exercises this path via the scatter DPS-preserve blend and now round-trips cleanly.
…view fixes
Mask-form codegen: pto.tscatter requires the maskPattern attribute *inside*
ins() after the src operand (same shape as pto.tgather's mask form), e.g.
`ins(%src, {maskPattern = #pto.mask_pattern<P0101>} : src_ty) outs(%dst)`.
The previous trailing-attribute form made PTOAS fail with "expected ',' after
src operand in ins(...)". The codegen UT now asserts maskPattern appears inside
ins() (before outs()) so the layout can't regress.
Also addresses review feedback on the index form:
- Guard the INT16 flat-index range in the tensor.scatter lowering (n*cols must
stay <= 32768) so an oversized 2-byte tile fails loudly instead of silently
overflowing to wrong destination addresses.
- Add an INTERNAL_CHECK that the src/indexes type annotations are both present
or both absent (a one-sided annotation is a codegen bug, not valid input).
- Document the duplicate-index ascending last-wins ordering as a pto.tscatter
ABI guarantee that the lowering and ST reference both rely on.
Tests: add mask-form ST (TestScatterMaskForm, P0101/P1010) on the A2/A3
backend (Ascend910B); A5/Ascend950 rejects the mask form. dst is zero-init so
the expected unselected columns are correct regardless of whether tscatter
zeros or preserves them.
- tscatter mask form: emit maskPattern inside ins() after src, before the type
annotation (ins(%src, {maskPattern...} : src_ty) outs(%dst)) — device-verified.
- INTERNAL_CHECK -> INTERNAL_CHECK_SPAN(op->span_) for the scatter type-annotation
symmetry check (gemini).
- INT16 flat-index guard: print the actual element dtype instead of hardcoding
"2-byte" (INT16 covers 1- and 2-byte; gemini).
- INT16 scatter flat-index guard: bound rows via division (kMaxFlat/cols) instead of n*cols to avoid signed int64 overflow; handle cols==0
Fold the RoPE even/odd reassembly repro into TestScatterMaskForm: write two compact inputs into one dst by chaining P0101 then P1010 mask scatters, pinning that the second scatter preserves the first's writes (dst[:, 0::2] = even, dst[:, 1::2] = odd). even/odd use disjoint positive ranges so a swapped pattern or clobbered column is caught.
Little-oil
force-pushed
the
feat/scatter-review-fixes
branch
from
fc39366 to
24a053a
Compare
… / WIP) Skip test_scatter_fp16/bf16/int16: the index-form 2-byte path currently fails on device due to a pto-isa bug (not this PR's codegen); re-enable once pto-isa lands the fix. The fp32/int32 4-byte cases still run. Skip test_scatter_mask_chain: the chained P0101->P1010 reassembly into a single dst is still being root-caused; the single-pattern P0101/P1010 mask cases still run.
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
1 participant
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.
Post-merge fixes for the scatter operators added in #1426, plus a fix for the
cmp/cmps TileView round-trip gap (#1498) that had forced a related test to be skipped.
1.
pto.tscattermask-form PTOAS syntax (functional bug)The mask-form codegen emitted the
maskPatternattribute as a trailing dictafter
outs(...). PTOAS rejects that withexpected ',' after src operand in ins(...). The attribute must ride insideins()after the src operand,exactly like
pto.tgather's mask form:The codegen UT now asserts
maskPatternappears insideins()(beforeouts()so the layout cannot regress.
2. cmp/cmps packed-mask TileView blayout roundtrip — Closes #1498
The python printer elides a TileView field when it matches
GetImplicitTileView(tile_type.shape_, ...)(implicit view from the physicaltile shape). The text parser instead recomputed the implicit blayout from
valid_shapewhen present, desynchronising the two for packed-mask tiles: acmp/cmps result with physical shape
[16, 8]butvalid_shape [16, 1]had itsrow_majorblayout omitted by the printer, while the parser inferredcol_majorfromvalid_shape'scols==1→print->parsefailed withTileView blayout mismatch.Fix: the parser now infers the implicit defaults from the physical tile shape
(falling back to
valid_shapeonly when the shape is unavailable), matching theprinter. This un-skips
TestConvertScatterOp::test_scatter_conversion, whichexercises the path through the scatter DPS-preserve blend.
3. Index-form review fixes
tensor.scatterlowering (n*cols <= 32768)so an oversized 2-byte tile fails loudly instead of overflowing to wrong addresses.
INTERNAL_CHECKthat the tscatter src/indexes type annotations are both presentor both absent (a one-sided annotation is a codegen bug).
pto.tscatterABI guarantee the lowering and ST reference rely on.
Tests
TestScatterIndexForm) across the dst/src + indexes dtype matrix(fp32 / int32 / fp16 / bf16 / int16), plus the repeated-index last-wins case and
the single-row regression for [Bug] pl.tensor.scatter on 1-row src tile triggers tile.ci cols!=1 ISA check #1586.
TestScatterMaskForm, P0101/P1010, plus a chained P0101→P1010reassembly case) on the A2/A3 backend (
Ascend910B); A5/Ascend950rejects themask form.
dstis zero-init so the expected unselected columns hold regardlessof zero-vs-preserve semantics.
test_scatter_conversionun-skipped, now passing under the autouse roundtrip instrument.tests/ut/ir/transformssuite (all under the roundtrip instrument): 1330 passed,25 skipped; parser/printer/type_resolver suites green.
test_scatter_fp16/test_scatter_bf16/test_scatter_int16: currently failing on device due to a pto-isa bug inthe 2-byte (fp16/bf16/int16) lowering path, not in this PR's codegen. Skipped via
@pytest.mark.skippending a pto-isa fix; the fp32/int32 (4-byte) index-formcases pass. Will be re-enabled once pto-isa lands the fix.
test_scatter_mask_chain: the chained P0101→P1010 scatterinto a single
dst(RoPE even/odd reassembly) has a failure still being root-caused;skipped for now. The single-pattern P0101/P1010 mask cases pass.