fix(codegen): size the GM pipe workspace for both rings of a bidirectional pipe - #2271
fix(codegen): size the GM pipe workspace for both rings of a bidirectional pipe#2271vloncar wants to merge 2 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 Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughGM pipe workspace sizing now accepts explicit ChangesGM pipe workspace sizing
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related issues
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: c8960f9283
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| const int64_t pipe_bytes = static_cast<int64_t>(ring_count) * static_cast<int64_t>(slot_count) * | ||
| static_cast<int64_t>(slot_size); |
There was a problem hiding this comment.
When a function contains multiple frontend pipe IDs and an earlier pipe is bidirectional or has an explicit slot_num, this allocation now reserves ring_count * slot_count * slot_size, but PTOCodegen::PrepareGMSlotBufferLayout still advances gm_slot_buffer_offsets_ by only GetGMPipeSlotCount(dir_mask) * slot_size and ignores both the second ring and explicit slot_num (src/codegen/pto/pto_codegen.cpp:496-503). The next pipe is therefore handed an addptr into the previous pipe's second ring/extra slots, so two independent GM pipes can alias even though the total orchestration allocation is now large enough; please make the PTO offset layout use the same footprint calculation.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in c3dd15c — thank you, this was a real gap and I had missed it entirely.
You are right on both counts: PrepareGMSlotBufferLayout advanced by GetGMPipeSlotCount(dir_mask) * slot_size, so with more than one pipe in a function the next pipe's addptr landed inside the previous pipe's region — the second ring of a bidirectional pipe, or the tail of an explicit slot_num, whichever applied. The workspace being big enough overall does not help when the per-pipe offsets are wrong.
Rather than patch the second copy to match, I hoisted the rule into include/pypto/codegen/gm_pipe_layout.h and had both sites derive from it, since there were two independent copies of GetGMPipeSlotCount that had to agree and no way to tell when they stopped agreeing. Regression added in test_gm_slot_buffer_regions_do_not_overlap_across_pipes: pipe 0 is bidirectional and carries slot_num=8, so pipe 1's emitted offset is %c4096_index (2 rings x 8 slots x 1024 B). Before the fix it was %c1024_index — inside pipe 0.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/codegen/orchestration/orchestration_analysis.cpp`:
- Around line 129-137: Update the slot-count validation in the initialize_pipe
handling around slot_num_by_pipe so the direction-specific default is resolved
before comparison. Validate the resulting effective slot count on every call for
the same pipe_id and dir_mask, including calls that omit slot_num, while
preserving the existing consistency error for mismatched values.
- Around line 180-186: Update the validation in the pipe initialization logic
around slot_count and ring_count so dir_mask-derived values are validated before
accepting an explicit slot_num. Ensure invalid dir_mask values such as those
producing a nonpositive slot or ring count are rejected even when
slot_num_by_pipe contains a positive override, while preserving valid explicit
slot_num behavior.
- Around line 180-191: Update PTOCodegen::PrepareGMSlotBufferLayout to compute
each pipe’s byte range using the explicit slot_num when present, otherwise
GetGMPipeSlotCount(dir_mask), and multiply by GetGMPipeRingCount(dir_mask)
before advancing byte_offset. Add a regression covering two pipe IDs, including
an explicit slot_num and bidirectional ring, and verify their computed regions
do not overlap.
🪄 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: 669ac948-079c-4b46-8b1e-663b05cb1295
📒 Files selected for processing (2)
src/codegen/orchestration/orchestration_analysis.cpptests/ut/codegen/test_orchestration_tensor_rw.py
…ional pipe ComputeGMPipeWorkspaceElements sized __gm_pipe_buffer as slot_count * slot_size, one ring's worth, and re-derived slot_count from dir_mask. A bidirectional pipe is two rings — C2V at offset 0 and V2C at slot_num * slot_size, total 2 * slot_num * slot_size, per the A2A3 GM layout in pto-isa's HL_ptoisa_newfeature20260306_TPUSH_TPOP.md — so a DIR_BOTH pipe got half the memory it needs, and an explicit pl.cross_core_slot(slot_num=N) was ignored entirely. Multiply by a ring count derived from dir_mask, and honour the slot_num kwarg on initialize_pipe when it carries one, falling back to the dir_mask default when it does not. The under-allocation was masked by a matching pto-isa defect that never applied the per-direction ring offset, so both rings really did live in the first half. This must land before that fix (hw-native-sys/pto-isa#227), which otherwise writes past the end of the workspace. Fixes hw-native-sys#2269
c8960f9 to
2142c39
Compare
|
please fix ai comment |
Review of hw-native-sys#2271 found a second copy of the GM pipe layout rule. PTOCodegen::PrepareGMSlotBufferLayout advances each frontend pipe's byte offset by GetGMPipeSlotCount(dir_mask) * slot_size, ignoring both the second ring of a bidirectional pipe and an explicit slot_num. With more than one pipe in a function, the next pipe's addptr therefore lands inside the previous pipe's region and two independent pipes alias — even though the orchestration workspace is now large enough to hold both. Hoist the rule into include/pypto/codegen/gm_pipe_layout.h and derive both the workspace size and the per-pipe offsets from it, so the two can no longer drift apart. Two smaller fixes fall out of the shared helper: the effective slot count is now validated on every initialize_pipe call for a pipe (a call that omits slot_num while another states a non-default value is a mismatch, not a silent win for the explicit one), and dir_mask is validated before any explicit slot_num is honoured, so an unlayoutable direction is still rejected. Adds a regression asserting a second pipe starts past the first pipe's full two-ring, explicit-slot_num footprint.
|
Pushed c3dd15c addressing all four review comments — the Codex P1 and CodeRabbit's three. Summary for anyone reading the PR rather than the threads: The substantive one (flagged independently by both bots) was real and I had missed it: The root problem was duplication — two independent copies of The two smaller findings fall out of the shared helper: the effective slot count is validated on every New regression
|
Review of hw-native-sys#2271 found a second copy of the GM pipe layout rule. PTOCodegen::PrepareGMSlotBufferLayout advances each frontend pipe's byte offset by GetGMPipeSlotCount(dir_mask) * slot_size, ignoring both the second ring of a bidirectional pipe and an explicit slot_num. With more than one pipe in a function, the next pipe's addptr therefore lands inside the previous pipe's region and two independent pipes alias — even though the orchestration workspace is now large enough to hold both. Hoist the rule into include/pypto/codegen/gm_pipe_layout.h and derive both the workspace size and the per-pipe offsets from it, so the two can no longer drift apart. Two smaller fixes fall out of the shared helper: the effective slot count is now validated on every initialize_pipe call for a pipe (a call that omits slot_num while another states a non-default value is a mismatch, not a silent win for the explicit one), and dir_mask is validated before any explicit slot_num is honoured, so an unlayoutable direction is still rejected. Adds a regression asserting a second pipe starts past the first pipe's full two-ring, explicit-slot_num footprint.
c3dd15c to
a06c857
Compare
|
CI on a06c857 is red on four device jobs ( 1. Three of the four fail identically on 2. Every failure is the same device-bootstrap error, before any generated kernel runs: 135 occurrences in the dist job alone. This is AICPU op loading at 3. The failing tests cannot reach this code path. 4. The same commit's sim and compile jobs are green: Reads like device/runner state on the NPU runners rather than anything in this change. Could someone with permissions re-run the failed jobs? Happy to be proven wrong — if the collectives keep failing on a clean runner while For completeness, the last push itself was small: |
2 participants
What
Fixes #2269.ComputeGMPipeWorkspaceElementssized the injected__gm_pipe_bufferasslot_count * slot_size— one ring's worth — and re-derivedslot_countfromdir_mask. Two changes:GetGMPipeRingCount(dir_mask)— 2 for bidirectional, 1 otherwise — so a bidirectional pipe gets2 * slot_num * slot_size, matching the A2A3 GM layout specified in pto-isa'sdocs/HL_ptoisa_newfeature20260306_TPUSH_TPOP.md:169and drawn at:288(C2V ring at offset 0, V2C ring atslot_num * slot_size).slot_numkwarg oninitialize_pipewhen present, falling back to thedir_maskdefault when absent. Mismatchedslot_numvalues for one pipe id are aCHECK, mirroring the existingslot_sizeconsistency check.src/codegen/orchestration/orchestration_analysis.cppslot_numtests/ut/codegen/test_orchestration_tensor_rw.pyTest changes
test_submit_dispatched_pipe_group_sizes_workspace_and_resolves_calleesasserted["512"]elements for adir_mask=3, slot_size=512pipe — one ring, i.e. the bug written down as an expectation. It becomes["1024"]. Two new tests pin the rule:test_gm_pipe_buffer_bidirectional_reserves_two_rings—dir_mask=3, slot_size=512→2 * 4 * 512 B= 1024 floats. Fails onmainwith 512.test_gm_pipe_buffer_honours_explicit_slot_num—dir_mask=2, slot_size=512, slot_num=16→16 * 512 B= 2048 floats. Fails onmainwith 1024, thedir_maskdefault of 8 slots.Behaviour
The GM workspace for bidirectional pipes doubles (4 KiB → 8 KiB per pipe at the default depth) and follows
slot_numwhere given. No IR-structural change: the injected__gm_pipe_bufferparameter, its placement and the pass pipeline are untouched — only the element count on the orchestration-sidetensor.create.Why this is worth landing on its own, and first
Today this is a latent sizing bug: the under-allocation is consistent with a matching pto-isa defect where the device never applies the per-direction ring offset either, so both rings really do live in the first half and the short allocation "fits". Fixing this side alone changes nothing observable — measured, the resulting corruption is bit-identical.
But the device-side fix (hw-native-sys/pto-isa#227, for hw-native-sys/pto-isa#226) puts the V2C ring at
base + slot_num * slot_size, which with the current sizing is past the end of the allocation. So this must ship first. It is safe to merge ahead of the device change, and it is wrong on its own terms regardless — against a spec that predates both bugs.Together the two fix a silent cross-core corruption on A2A3 hardware. Full analysis, minimal repro and the prior history (including hw-native-sys/pto-isa#195, which reported the device half in July and was withdrawn) are in hw-native-sys/pto-isa#226.
A related inconsistency, deliberately left out
BuildAutomaticPipeSetup(src/ir/transforms/utils/cross_core_pipe.cpp:348,352) passes the rawslot_num_overridetoCreateInitializePipewhile sizing the reserved buffer fromeffective_slot_num, so in the default case the emittedinitialize_pipecarries noslot_numattribute even though a concrete depth was used. That is harmless — every consumer re-derives the same default fromdir_mask— and orthogonal to this fix, which readsslot_numonly when a pipe actually carries one. Passingeffective_slot_numinstead makes the attribute always explicit, but it adds a kwarg to every auto-generated pipe and so changes the expected IR in 12 structural-equality tests (test_expand_mixed_kernel_a2a3.py,_a5.py,_split_aiv.py). Happy to send that as a follow-up if you want the attribute made explicit.Validation
tests/ut/codegen+tests/ut/ir/transforms: 3195 passed.On hardware (Atlas A2/A3, fp32), with the pto-isa half applied, at the stock ring depth and with no DSL overrides: a GLA/ZeCO sequence-parallel forward that failed 5/12 configurations at C=16 and 6/12 at C=32 now passes 24/24 over C ∈ {16,32} × P ∈ {1,2,4} × N ∈ {2,4,8,16}, max abs error ≤ 3.1e-5; and a standalone cube↔vector loop repro goes from 12/21 corrupted dispatches to 21/21 clean.