Repository navigation
feat(passes): rotate slots by the iteration counter, not the induction variable - #2922
georgebisbas wants to merge 3 commits into
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configuration
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: Advanced Run ID: 📒 Files selected for processing (5)
Included review availability: This review used your included allowance. Your plan provides up to 2 included reviews per hour; 1 remain after this review. 📝 WalkthroughWalkthroughThe pipeline-to-slots pass now accepts loops with positive constant steps and computes slot indices from the iteration count for non-unit steps. Documentation and tests describe and check the updated eligibility rules, slot rotation, and code-generation behavior. ChangesStepped Pipeline Slot Rotation
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Feature Merge Risk: ⚪ Minimal · up to The documented eligibility rules match the implementation; no identified issue remains to resolve before merging. Security Architecture ReviewSecurity architecture risk: 🟡 Moderate · up to Eligible stepped loops now share on-chip tile slots. The new index is bounded and unsupported loops retain their fallback, but synchronization of the new index has not been established for every admitted loop shape. Incorrect slot reuse could produce wrong results or stall a kernel. No new privilege boundary is evident. Retained concerns
Security review detailsSecurity Blast Radius
Trust Boundaries and Controls
Resilience and Maintainability Implications
Hardening Proposals
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 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. A rabbit counts each step with care Comment |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Codex Review
Reviewed dad0d2009acf818f708565ee8f606963c9bcd452 against base 63412dd65c7914d5444bec73c90a0dc990f3ef08.
Automated assessment of untrusted PR content; not a guarantee of correctness.
Reviewed the five-file PR diff from the verified merge-base 731191c. The supplied discussion contains no earlier findings or disputed P0–P2 threads to reassess. One P2 issue remains.
1 finding(s) posted as inline comments.
No automatic approval: Review contains findings or is incomplete.
There was a problem hiding this comment.
Codex Review
Reviewed 70f723593c421bab05e70852a53789665645e248 against base 63412dd65c7914d5444bec73c90a0dc990f3ef08.
Automated assessment of untrusted PR content; not a guarantee of correctness.
Reviewed the full five-file diff from merge-base 731191c. The earlier P2 overflow finding is fixed: the new gate rejects negative starts, and a nonnegative start with a positive step keeps iv - start within the signed index range. No actionable P0–P2 issues remain.
Approval policy: passed.
…n variable
`LowerPipelineToSlots` required `step == 1` and an aligned `start`, so any
pipeline with a different step was declined and replicated by
`LowerPipelineLoops` — F body copies, F x the IR through every later pass, and
F x the emitted code. That is not a corner: the workspace's model corpus has 86
explicit-step pipelines and not one of them uses `step == 1` (K_TILE,
QUANT_TILE, KV_TILE, 128, 64, 32, ...), across 47 files.
The gate's stated reason was that ptoas matches the *affine form* of the slot
index, so `((iv - start) / step) % F` would "risk losing exactly the analysis
this transform exists to trigger". Measured against ptoas 0.65 that is stricter
than ptoas is: a step-64 load loop slotted this way produces the same ptoas
output as the equivalent unit-step loop (171 lines, set_flag=8, wait_flag=8).
`PTOAS/lib/PTO/Transforms/Passes/SlotAffineAnalysis.cpp` shows why —
`findMultiTileSlotExpr` returns whatever `pto.multi_tile_get` carries, and
`extractSlotForm` accepts `arith.remui X, %const` for any `X`. Only the
*identity* of that value decides whether two accesses are provably the same
slot; its arithmetic shape never does.
So this admits any compile-time-constant `step > 0` and emits
`((iv - start) / step) % factor` when the induction variable is not already the
rotation counter. `step == 1` loops keep the literal `iv % factor`, so their
emitted IR is unchanged. A runtime `step` or `start` is still declined: the slot
would depend on a value the analysis cannot bound.
What is given up: the kDisjoint inference for an `(iv + c) % N` pair. The
constant peel now sits above the division, so the two sides carry different
symbols and the relation degrades to kUnknown, which the callers treat
conservatively — the dependency is kept. Precision, not correctness.
Verified in sim Docker (PTOAS planner, ptoas --pto-level=level2):
* tests/ut/ir/transforms/test_lower_pipeline_to_slots.py 22 passed
* tests/ut/codegen/test_multi_buffer_codegen.py 32 passed
(SteppedPipelinedLoad + four tests pinning the counter-based slot)
* tests/ut/codegen 1396 passed,
2 failed — both the pre-existing test_orchestration_codegen_graph.py pair
* all repo pre-commit hooks pass
Not verified: behaviour on silicon. The sync-op counts match, but per-slot event
accounting is what has to be right at runtime and the failure mode is a race,
not a crash — see the plan's Developer verification section.
`LoopShapeAllowsSlots` accepted any constant `start`, but both emitted slot forms need the induction variable non-negative. `iv % factor` is a signed remainder, so `start = -2, step = 1, factor = 2` emits `(-1) % 2 == -1` — a slot before the region — and the counter form's `iv - start` is not representable for a sufficiently negative start (`pl.pipeline(INT64_MIN, 1, 1 << 61, stage=3)` reaches `iv = 0`, where `0 - INT64_MIN` overflows and the wrapped quotient selects slot `-1` of a three-slot allocation). Requiring `start >= 0` closes both: `iv >= start >= 0` keeps the subtraction non-negative and bounded by the loop's own span, so both remainders land in `[0, factor)`. Loops that start negative now decline and keep the `LowerPipelineLoops` replication path, which handles arbitrary bounds. Tests, both in `TestFallback`: `test_a_negative_start_is_declined` pins the pre-existing `-2` hole the old `start % factor == 0` test admitted, and `test_a_start_whose_counter_would_overflow_is_declined` pins the five-iteration `INT64_MIN` example.
70f7235 to
fd2f5ae
Compare
Superseded by a new Codex review
There was a problem hiding this comment.
Codex Review
Reviewed fd2f5aed6bd7b894ebc0ec8b1953cfd4ea271638 against base be5168bbb5477493e38a244fa4af780278bdad37.
Automated assessment of untrusted PR content; not a guarantee of correctness.
Reviewed the full five-file diff. The earlier P2 counter-subtraction finding is fixed by the nonnegative-start gate. A separate terminal-increment overflow remains for newly admitted stepped loops.
1 finding(s) posted as inline comments.
No automatic approval: Review contains findings or is incomplete.
| } | ||
| // `((iv - start) / step) % factor` is well-formed: `iv >= start >= 0` keeps | ||
| // the subtraction non-negative and bounded by the loop's own span. | ||
| return true; |
There was a problem hiding this comment.
[P2] Reject stepped loops whose final induction update overflows
This gate admits pl.pipeline(INT64_MAX - 4, INT64_MAX, 3, stage=3). It has two valid iterations, at INT64_MAX - 4 and INT64_MAX - 1, but slot lowering retains a scf.for whose update after the second iteration computes INT64_MAX + 2. That overflows the signed index and can make the loop continue incorrectly or hang. Previously, LowerPipelineLoops emitted two tail copies with no loop because the trip count is below the stage count. Decline this shape or lower it with an overflow-safe iteration counter.
There was a problem hiding this comment.
Fixed in 450672f.
You are right: for step > 1 the retained scf.for computes iv + step after the last iteration, which can overshoot stop by up to step - 1. pl.pipeline(INT64_MAX - 4, INT64_MAX, 3, stage=3) would reach INT64_MAX - 1 + 3. step == 1 cannot overshoot, so that form is unchanged.
LoopShapeAllowsSlots now declines a stepped loop unless stop is a constant with stop <= INT64_MAX - (step - 1). A runtime stop cannot be bounded, so it declines as well; stepped loops were not slotted at all before this PR, so those stay on replication exactly as on main. Tests: your example is declined, the tight boundary stop == INT64_MAX - (step - 1) is still slotted, and a stepped loop with a runtime stop is declined. The English and Chinese pass docs describe the rule.
A stepped loop (step > 1) admitted by LoopShapeAllowsSlots keeps its scf.for, which computes iv + step after the last iteration. That update can overshoot stop by up to step - 1, so pl.pipeline(INT64_MAX - 4, INT64_MAX, 3, stage=3) overflows the signed index and can keep the loop running. Require a constant stop no larger than INT64_MAX - (step - 1) for step > 1. step == 1 cannot overshoot and keeps accepting a runtime stop; a stepped loop with a runtime stop was never slotted before this series and stays on replication.
There was a problem hiding this comment.
Codex Review
Reviewed 450672f2ce647ba785f862ab84ab596b5eae302a against base be5168bbb5477493e38a244fa4af780278bdad37.
Automated assessment of untrusted PR content; not a guarantee of correctness.
Reviewed the full five-file diff. Both disputed P2 findings are fixed in the current code: negative starts are declined, and stepped loops require a constant stop that leaves room for the final induction update. No actionable P0–P2 issues remain.
Approval policy: passed.
This branch has not been deployed
1 participant
What this offers
LowerPipelineToSlotsis the cheap pipeline lowering: one loop body, with each per-stage buffer becoming sloti % Fof a single allocation. It requiredstep == 1, so any pipeline with a different step was declined and handed toLowerPipelineLoops, which replicates the bodyFtimes.That is not a corner case. Across the workspace's real model corpus:
step == 1step != 1pypto-lib/modelspypto/testsEvery explicit-step pipeline in the model corpus has a step that is not 1 —
K_TILE,QUANT_TILE,KV_TILE,128,64,32,512and ~40 other values — spread over 47 files, with thedeepseek_v4_*variants carrying 5–9 each. Their bodies hold top-levelpl.loads, so they satisfy every other gate. Today all of them take the replication path.Measured, same planner, one variable changed
pl.pipeline(0, 256, 64, stage=2)(4 iterations,step = 64, 2 slots), compiled with the PTOAS planner in both runs and ptoas at--pto-level=level2; the only difference is whether the gate admits a non-unit step:main@731191cf)pto.alloc_multi_tilecount=2)pto.multi_tile_getpto.tloadpto.alloc_tilestep %c128_index← rewritten toF × stepstep %c64_index← the user's own step203 → 166 lines (−37 lines) on this shape, and the reduction scales with the pipeline depth (
F) and with how much work the body carries — because the body itself is no longer multiplied. Every later pass sees one copy instead ofF:CanonicalizeIOOrder's topological sort,InitMemRef,MemoryReuse's O(S·M²) pack,AllocateMemoryAddr, and ptoas's own parse and pass pipeline. There is also no remainder dispatch and nopipeline_membershipmachinery, since bounds, step anditer_argsare left untouched.The IR, before and after
Before — the loop runs at
step 128(F × 64), and the body holds the replicated pair: twoalloc_tiles, twotloads, withmaxsiclamps on the rewritten offsets.After — one body, the loop keeps
step 64, and the per-stage buffer is a slot selected by the rotation counter:(type suffixes elided for width; artifacts reproduced with
probe19_ab.py)Why the old gate is safe to relax
The gate was documented as protecting a real property, so this is a falsification, not an assumption:
That is stricter than ptoas is. In
PTOAS/lib/PTO/Transforms/Passes/SlotAffineAnalysis.cpp:findMultiTileSlotExprreturns whateverpto.multi_tile_getcarries — the slot value, not a required form;extractSlotForm's first case acceptsarith.remui X, %constfor anyX; onlypeelAddSubConstis shape-sensitive, and it peels at most four add/sub-of-constant steps;compareSlotSSAreasons about(innerSym + offset) mod N, and the only thing it needs is that two slots share the identity of that value. Adivsiis a perfectly good symbol.Measured against ptoas 0.65: a
step = 64loop slotted this way produces the same ptoas output as the equivalentstep == 1loop — 171 lines,set_flag=8,wait_flag=8, identical after normalising event ids.And the property the gate protected is delivered by the new form, not lost. A slot-keyed sync is an ordinary
set_flag/wait_flagwhose event-id operand is a runtime value (the EmitC lowering ofpto.set_flag_dyn; the literalset_flag_dynnever appears in generated C++):stage=2stage=2(this change)The stepped form earns the same per-slot keying as the unit-step form; the replication path still earns none.
What it costs — stated, not glossed
The
kDisjointinference degrades tokUnknownfor an(iv + c) % Npair. The constant peel now sits above the division, so the two sides carry different symbols. The callers treatkUnknownconservatively — the dependency is kept — so this is a precision loss, never a wrong answer. It is documented in both languages in this PR.What this does not do
CollectCandidatescollects top-levelpl.tile.loadonly, and anAutoTileMatmulL0K-loop bodytile.extracts from an already-loaded tile. Changing that is a separate, larger item (and it runs into the L0 budget — L0A/L0B are 64 KB each, and a two-slot region of 128×128 fp16 K-blocks is exactly 64 KB).memory_planner=PTOAS; under the defaultPyPTOplanner it returns every function untouched, so that path is unchanged.Verification (agent, sim Docker)
Run on the pushed commit after a clean
rm -rf build _skbuild && pip install --no-build-isolation -q ".[dev]"in the container that measured.tests/ut/ir/transforms/test_lower_pipeline_to_slots.pytests/ut/codegen/test_multi_buffer_codegen.pySteppedPipelinedLoad+ 4 teststests/ut/codegentest_orchestration_codegen_graph.pypair, reproduced on unmodifiedorigin/mainstep == 1loops keep the literaliv % factor, so their emitted IR is unchanged — the change is additive.Developer verification (NPU) — passed 2026-09-28
The gate this change needs is numeric, not "it did not crash": the slot rotation and the replicated fallback must compute the same numbers on silicon. Run on the frozen head
70f72359— checkout in apypto3-hw-native-sys:cann9container joined to the host queue; every job throughtask-submitonhng-atlas01(a2a3, one card per run):tests/st/runtime/ops/test_memory_planner_ptoas.pyon device (PTOAS among the planners)pl.pipeline(0, 256, 64, stage=2)andpl.pipeline(64, 320, 64, stage=2), each compiled under PTOAS (slots) and PyPTO (replicated) and run against one goldenstartshapesalloc_multi_tile count=2+ onemulti_tile_get, loop keepsstep %c64; PyPTO: no region, replicated (step %c128)Probe sources and the full command/task record:
hw-native-sys/slots-verify/and plan 125 § NPU verification record inpypto-3.0-notes.Pushed by the agent at the developer's request; per the branch-and-PR standard the agent stops at the push and the developer opens the PR. The NPU verification above was run on 2026-09-28 and passed.