feat(ir): Require elementwise valid-region agreement - #2159
Conversation
|
Warning Review limit reached
Next review available in: 46 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (8)
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: e291e3ebcd
ℹ️ 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".
Implements brief 12 of the unified valid_shape series. One shared rule now decides the valid region of every elementwise result, for tensor and tile alike. Each deducer previously copied GetValidShape(arg0) onto the result under a standing TODO, which silently widened or narrowed the result when the operands disagreed and made commutative ops operand-order-dependent. InferElementwiseValidShape right-aligns operands to the result rank and classifies each per output dimension (missing leading axis, physical singleton broadcasting up, or non-broadcast contributor), then requires every non-broadcast contributor to hold a provably equal valid extent. An undecided symbolic relation rejects too: no runtime guard is emitted. This is the hardware contract, not a policy choice. TAdd_Impl asserts src0/src1 valid extents against dst's, so when the operands disagree no choice of result region satisfies the assert -- rejecting is the only sound answer. TPartInstr conversely tests each source per cell and zeroes cells in neither, which is why part_add/part_mul/part_max/part_min take the union and admit it only when it is an origin-anchored rectangle. MakeElementwiseTileView takes layout from an operand owning the full output shape, falling back to the layout implied by that shape, so a broadcast singleton can no longer dictate the result's layout. Also lands the two items deferred by the unary/reduction PR: - Fresh results no longer inherit pad. A computed result is written only across its valid region, so the source's promise about the cells outside it does not survive; exp of a zero-padded tile is not zero-padded. This changes all five InheritTileViewLayout families atomically (elementwise, unary, broadcast, gather, sort) -- doing it per-family makes the DSL fail to type-check across reassignments -- and it is codegen-visible, since TileView::pad flows into the emitted tile-buf type. The fillpad family still establishes its own pad immediately after the call. - Tensor scalar-binary ops (adds/subs/muls/divs/fmods) preserve validity, matching their tile equivalents. Scratch and predicate operands are excluded from value-validity agreement. DeduceTileOpTileScalarTileType gains third_is_addend because it serves both tile.rems (scratch tmp) and tile.addsc/subsc (real addend), and MakePackedPredicateTileType now bit-packs the agreed logical region rather than reading lhs directly. A fully valid input still yields a fully valid, view-less result, so no existing program changes shape. Adds 48 tests covering matching partials, static and symbolic disagreement, singleton full/empty/unknown validity, lower-rank and mutually broadcasting shapes, operand order, every tile multi-operand family, mask validity, the part_* union and its rejections, layout derivation, and fresh-result metadata.
Addresses two review findings on the elementwise valid-region rule. ClassifyAxis matched a literal ConstInt(1) result extent to decide that a unit operand axis was not widening. BroadcastShapes keeps the *first* operand's expression for provably equal dims, so a unit result axis can arrive in symbolic form (n - n + 1 against a literal 1). The literal operand was then classified as a broadcast singleton and its valid extent required to be 1 -- rejecting an empty axis that widens nothing, and only in one operand order. Prove the result extent instead of pattern-matching a literal. PrefersAsRepresentative preferred a ConstInt, which leaves the first-seen expression in place when the analyzer proves two *non-constant* extents equal (n against n + 0). TileView/TensorView equality and hashing compare structurally, so a commutative operation could still produce unequal types and trip a reassignment or join check downstream. Break that tie on a stable printed key. Both cases are pinned by tests that fail without the corresponding fix.
Rebasing onto the PTOAS op batch (hw-native-sys#2157) put two overlapping checks on the same programs. `ConvertTensorToTileOps` gained a divisor-cover check for `tensor.div` row broadcasting and a same-valid_shape check for plain tdiv; the elementwise valid-region rule already rejects those operand pairs while the program is still being built, so the pass never sees them and the `pytest.raises` around the pass no longer fires. Move the expectation to where the rejection now happens. Each test still asserts that exactly the same operand pair is refused, and the positive shared-symbol case still lowers through the pass unchanged. The pass-level checks stay in place for IR that reaches them without going through operator type deduction.
33814e9 to
24465df
Compare
Implements brief 12 of the unified
valid_shapeseries (Require elementwisevalid-region agreement). Follows #2031 (canonicalization), #2037 (proof/bounds),
#2043 (type boundaries), #2048 (window reads) and #2051 (unary/reduction).
Summary
One shared rule now decides the valid region of every elementwise result, for
tensor and tile alike. Previously each deducer copied
GetValidShape(arg0)ontothe result under a standing
TODO(YunjiQin): assumes both src tiles have the same valid_shape, which silently widened or narrowed the result whenever the operandsdisagreed, and made commutative operations operand-order-dependent.
InferElementwiseValidShaperight-aligns the operands to the result rankand classifies each per output dimension — missing leading axis, physical
singleton broadcasting up, or non-broadcast contributor — then requires every
non-broadcast contributor to hold a provably equal valid extent.
MakeElementwiseTileViewpairs that with layout taken from an operandwhose physical shape equals the output shape, falling back to the layout
implied by the output shape. A broadcast singleton can no longer dictate the
result's layout.
part_add/part_mul/part_max/part_mininstead take the union,admitted only when it is representable as an origin-anchored rectangle.
pad, and tensor scalar-binary ops nowpreserve validity — the two items feat(ir): preserve valid regions through unary and reduction ops #2051 deferred to this brief.
A fully valid input still yields a fully valid, view-less result, so no existing
program changes shape.
What the ISA says
Both the CPU model and the a2a3 NPU path assert the sources against the
destination (
cpu/TAdd.hpp:70-74,npu/a2a3/TAdd.hpp:65-70):That assert is what this PR's equality rule encodes. But see the open question
at the bottom —
PTO_ASSERTis_DEBUG-only, and the release kernel enforcesthe weaker
dst.valid <= every src.valid, which changes what the right rule is.TPartOp.hppconversely tests each source per cell and zeroes cells inneither:
That is exactly why the union must be a representable rectangle: claiming
[6, 8]for[4, 8] ∪ [6, 6]would type silently-zeroed cells as valid data.Observable behavior
Commutative operations now produce structurally equal result types in either
order: extents are chosen per dimension, preferring a constant representative.
New diagnostics
<op>: operands <a> and <b> do not provably agree on the valid extent of dimension <n> (<a>: X, <b>: Y). An elementwise operation reads every operand at each cell and writes one region, so every non-broadcast operand must have a provably equal valid extent. Narrow the wider operand, or fill its padding with a fillpad before combining<op>: operand <label> broadcasts dimension <n> from a physical extent of 1, but its valid extent there is X, which is not provably 1. A broadcast singleton replicates its one cell across the whole output dimension, so that cell must be validpart_*operands with no provable ordering on some dimension<op>: the operand valid regions (...) have no provably widest extent on dimension <n>, so their union cannot be expressed as a valid_shape...part_*operands not ordered by containment<op>: the operand valid regions (...) are not ordered by containment, so their union [...] is not an origin-anchored rectangle and would claim cells that are valid in no operand. Narrow one operand so that one region contains the otherUndecided symbolic equality rejects deliberately: no runtime guard is emitted, so
taking an unproved relation on trust would silently widen or narrow the region.
Operand audit
Scratch and predicate operands are excluded from value-validity agreement:
Two consequences worth calling out:
DeduceTileOpTileScalarTileTypeserved bothtile.rems(arg 2 = scratch)and
tile.addsc/subsc(arg 2 = real addend). It gained athird_is_addendflag; without it a scratch buffer would be held to agreement.MakePackedPredicateTileTypenow bit-packs the agreed logical region ratherthan reading lhs directly, so mask validity follows the same rule as value
validity.
Deviations from the brief
fractalis derived, not inherited. The brief says "inherit productionlayout/fractal".
InheritTileViewLayouthas never copiedfractal, and addingit would change unary/gather/sort too — outside this brief's inherited scope.
The derive-from-output-shape branch yields the
TileViewdefault (512), whichis correct for these Vec-space ops.
validity-semantics change and is what makes the brief's acceptance criterion
"tensor and tile operations follow the same validity rule" hold.
A pre-existing bug this surfaces
A partially-valid
[M, 1]column-broadcast add now fails to compile.Measured on both sides of the change:
tile.addresult carriespl.TileView(valid_shape=[1, 128])— 1 row x 128 cols, when the real data is 40 x 64. The store writes the wrong region.tile.add: operands lhs and rhs do not provably agree on the valid extent of dimension 1 (lhs: 128, rhs: 64)The old output was silently wrong, so the loud failure is strictly better.
Root cause is
ResolveBackendOpLayouts' column-vector repair(
src/ir/transforms/resolve_backend_op_layouts_pass.cpp:188-204), which reshapes[M,1] -> [1,M]and naively transposesvalid_shapeinstead of mapping it ontothe broadcast axis it actually occupies. That is layout-repair territory, not
elementwise deduction — the deducers only ever see the already-reshaped operand.
The row-broadcast
[1,N]form and all fully-valid broadcasts are unaffected.Test-scaffolding change
Two cases in
test_convert_tensor_to_tile_ops.pyadded a narrowed slice to afully-valid gm input — programs that would trip
PTO_ASSERTon device. Theaddthere is scaffolding; the subject istensor.slice -> tile.slicepreservingvalid_shape/pad_value. The gm operand is now narrowed to match, and its loadcarries that valid shape. No assertion was weakened.
Testing
tests/ut/ir/operators/{test_tile_ops,test_tensor_ops}.py:matching partials, static and symbolic disagreement, symbolic agreement,
singleton full/empty/unknown validity, lower-rank and mutually broadcasting
shapes, operand order, every tile multi-operand family (binary, shift, ternary,
tri-tile, tile-scalar-tile, sel, sels, cmp), mask validity,
part_*union andits rejections, and fresh-result metadata (no alias, stride, offset or pad).
tests/st/codegen+tests/st/examples: no new failures against thepre-change baseline.
pre-commitand the diff-basedclang-tidygate are clean.User-facing English and Chinese documentation is deferred to the series'
documentation PR, per the series rules.
dist-system-testsandpypto-lib-modelboth fail, and both reduce to the samepattern — a statically-full operand combined with a symbolically-narrowed one:
Both are real, currently-passing programs (Qwen3-14B decode still passes on
device with this branch; only prefill trips).
What the hardware actually does.
PTO_ASSERTis compiled out unless_DEBUG(
pto/common/debug.h:33-36), and the release kernel bounds its loops bydst'sextents while reading the sources by row-stride offset
(
npu/a2a3/TAdd.hpp:37-51,TADD_IMPLpassingdst.GetValidRow()/GetValidCol()).So the enforced release contract is not equality but
with the
_DEBUGequality assert acting as a stricter "you should have narrowedalready" check.
That makes all three candidate rules comparable:
dst = lhs.valid)valid_tokdst = min(...))min(64, valid_tok) <= bothIntersection accepts both programs and is sound, which suggests the brief's
"provably equal / unknown rejects" wording may be the wrong rule. Changing it is a
core-semantics decision beyond what this PR should take unilaterally, so I have
left the brief's rule implemented as specified and am flagging it for a call.