[RFC] feat(distributed): ergonomic pld.* collective wrappers — auto window+signal; typing, validation, docs - #2275
Conversation
📝 WalkthroughWalkthroughThis PR adds HOST ergonomic collective wrappers with automatic signal allocation, public exports, documentation, unit tests, and distributed end-to-end tests for all-reduce, broadcast, all-gather, and barrier. ChangesCollective API
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant HostOrchestration
participant CollectiveAPI as collective_api.all_reduce
participant SignalWindow
participant TensorCollective as HOST tensor allreduce
HostOrchestration->>CollectiveAPI: invoke all_reduce(target, mode)
CollectiveAPI->>SignalWindow: allocate fresh INT32 signal
CollectiveAPI->>TensorCollective: delegate target and parameters
TensorCollective-->>HostOrchestration: return distributed target
Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Title checkExplanation The title accurately identifies the main change: ergonomic distributed
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: a9c2c42947
ℹ️ 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".
| Returns: | ||
| The ``target`` :class:`pld.DistributedTensor` (window-as-result). | ||
| """ | ||
| signal = _fresh_signal("all_to_all_v", [world_size(), 1]) |
There was a problem hiding this comment.
For any call to the new pld.all_to_all_v(...) wrapper, this line builds the hidden signal as [pld.world_size(), 1], but the underlying pld.tensor.all_to_all_v type deducer requires the signal's first dimension to be a compile-time ConstInt so it can derive MAX_RECV = target.shape[0] // NR. That means the advertised short form raises during IR construction before it can reach the documented host rejection path; either require a static nranks here (like ring all-reduce) or avoid exposing this wrapper until a dynamic/world-size signal is supported.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in aec3b80: pld.all_to_all_v now requires a positive static nranks (mirroring ring allreduce) and builds the signal as [nranks, 1]; a dynamic world_size() or non-positive value raises a context-rich ValueError at the wrapper call. Added delegation + guard unit tests. Thanks — this was a real construction-time failure.
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docs/en/dev/distributed_ops.md (1)
182-186: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winEN and ZH docs list different ops for the credit-barrier protocol. The English doc names
all_to_all_vas one of the collectives synchronized by the self-clearing credit barrier; the Chinese doc's equivalent sentence omits it. Align both lists to name the same set of ops.
docs/en/dev/distributed_ops.md#L182-L186: authoritative list; keepall_to_all_vhere and confirm it is intentional.docs/zh/dev/distributed_ops.md#L164-L165: addall_to_all_vto match the English list.As per coding guidelines, "English documentation in
docs/en/dev/is authoritative and corresponding Chinese documentation indocs/zh/dev/must remain aligned."🤖 Prompt for 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. In `@docs/en/dev/distributed_ops.md` around lines 182 - 186, Keep all_to_all_v in the authoritative collective list in docs/en/dev/distributed_ops.md at lines 182-186, confirming it remains intentional; update the corresponding list in docs/zh/dev/distributed_ops.md at lines 164-165 to add all_to_all_v so both documentation sets name the same credit-barrier operations.Source: Coding guidelines
🧹 Nitpick comments (2)
tests/ut/language/test_collective_api.py (1)
125-142: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd direct
all_to_all_vcoverage.This test module does not call
pld.all_to_all_v. Add a delegation test that checks the argument orderinput,target,signal,send_counts,recv_counts, plus the generated INT32[world_size(), 1]signal.Add this routine contract test in this module. As per coding guidelines, use unit tests in
tests/ut/for routine testing.🤖 Prompt for 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. In `@tests/ut/language/test_collective_api.py` around lines 125 - 142, Add a unit test in this module that calls pld.all_to_all_v and verifies delegation to the expected IR operation, argument order input, target, signal, send_counts, recv_counts, and the generated signal’s INT32 [world_size(), 1] shape. Follow the existing collective API test patterns and keep the coverage focused on this routine contract.Source: Coding guidelines
tests/st/distributed/test_l3_ergonomic_api.py (1)
240-244: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winValidate all all-gather rows.
consume_steploads only row 0. The expected value at Line 376 checks only rank 0's chunk on each rank. An implementation that broadcasts rank 0 or leaves later rows invalid passes this test. Copy the complete[NR, SIZE]target into each rank output and compare all rank-indexed input chunks.Proposed coverage change
def consume_step( self, target: pld.DistributedTensor[[NR, SIZE], pl.FP32], - out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], - ) -> pl.Tensor[[1, SIZE], pl.FP32]: - return pl.store(pl.load(target, [0, 0], [1, SIZE]), [0, 0], out) + out: pl.Out[pl.Tensor[[NR, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, SIZE], pl.FP32]: + return pl.store(pl.load(target, [0, 0], [NR, SIZE]), [0, 0], out) @@ def consume_orch( self, target: pld.DistributedTensor[[NR, SIZE], pl.FP32], - out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], - ) -> pl.Tensor[[1, SIZE], pl.FP32]: + out: pl.Out[pl.Tensor[[NR, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, SIZE], pl.FP32]: @@ - outputs: pl.Out[pl.Tensor[[NR, 1, SIZE], pl.FP32]], - ) -> pl.Tensor[[NR, 1, SIZE], pl.FP32]: + outputs: pl.Out[pl.Tensor[[NR, NR, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, NR, SIZE], pl.FP32]: @@ - outputs = torch.zeros_like(inputs) + outputs = torch.zeros((NR, NR, SIZE), dtype=inputs.dtype, device=inputs.device) @@ - expected = torch.stack([inputs[0]] * NR) + expected = torch.stack([inputs[:, 0, :]] * NR)🤖 Prompt for 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. In `@tests/st/distributed/test_l3_ergonomic_api.py` around lines 240 - 244, Update consume_step to load and store the complete [NR, SIZE] gathered target into out rather than only row 0. Extend the assertions around the existing expected-value check to validate every rank-indexed input chunk, ensuring later all-gather rows are compared on each rank.
🤖 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 `@python/pypto/language/distributed/op/collective_api.py`:
- Around line 12-15: Update the module documentation describing signal
allocation to explicitly state that all_reduce with mode="mesh" uses a
compiler-synthesized signal and barrier requires a caller-provided signal, while
preserving the existing fresh-signal behavior for other wrappers.
- Around line 67-69: Update the printer/serialization path for
_fresh_signal-generated buffers so alloc_window_buffer() is emitted as a named
assignment before the window() call, preserving a parser-visible binding that
parse_program() can consume. Ensure signal-bearing wrappers round-trip through
as_python() and make the strict collective API xfail pass.
- Around line 123-141: Update the nranks validation in the mode == "ring" branch
to require a static integer greater than zero, rejecting zero and negative
values before calculating the signal shape. Preserve the existing bool/type
validation and error handling, and add regression coverage for both non-positive
inputs.
---
Outside diff comments:
In `@docs/en/dev/distributed_ops.md`:
- Around line 182-186: Keep all_to_all_v in the authoritative collective list in
docs/en/dev/distributed_ops.md at lines 182-186, confirming it remains
intentional; update the corresponding list in docs/zh/dev/distributed_ops.md at
lines 164-165 to add all_to_all_v so both documentation sets name the same
credit-barrier operations.
---
Nitpick comments:
In `@tests/st/distributed/test_l3_ergonomic_api.py`:
- Around line 240-244: Update consume_step to load and store the complete [NR,
SIZE] gathered target into out rather than only row 0. Extend the assertions
around the existing expected-value check to validate every rank-indexed input
chunk, ensuring later all-gather rows are compared on each rank.
In `@tests/ut/language/test_collective_api.py`:
- Around line 125-142: Add a unit test in this module that calls
pld.all_to_all_v and verifies delegation to the expected IR operation, argument
order input, target, signal, send_counts, recv_counts, and the generated
signal’s INT32 [world_size(), 1] shape. Follow the existing collective API test
patterns and keep the coverage focused on this routine contract.
🪄 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: a1458aa0-907e-4e89-95a1-21cb9bd8eb1a
📒 Files selected for processing (8)
docs/en/dev/distributed_ops.mddocs/zh/dev/distributed_ops.mdpython/pypto/language/distributed/__init__.pypython/pypto/language/distributed/op/__init__.pypython/pypto/language/distributed/op/collective_api.pypython/pypto/language/distributed/op/unified_ops.pytests/st/distributed/test_l3_ergonomic_api.pytests/ut/language/test_collective_api.py
| name = f"__auto_{op_name}_{next(_SIGNAL_COUNTER)}" | ||
| buf = _tensor.alloc_window_buffer(shape, dtype=DataType.INT32, name=name) | ||
| return _tensor.window(buf, shape, dtype=DataType.INT32) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Make generated signal allocations round-trip.
_fresh_signal() embeds alloc_window_buffer() directly in window(). The printer emits that allocation inline, but parse_program() requires a named simple assignment. The strict xfail in tests/ut/language/test_collective_api.py, Lines 282-297, proves that signal-bearing wrappers cannot round-trip through as_python().
Hoist generated allocations into named statements in the printer or preserve an equivalent parser-visible binding. Make the strict xfail pass before exporting this API.
🤖 Prompt for 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.
In `@python/pypto/language/distributed/op/collective_api.py` around lines 67 - 69,
Update the printer/serialization path for _fresh_signal-generated buffers so
alloc_window_buffer() is emitted as a named assignment before the window() call,
preserving a parser-visible binding that parse_program() can consume. Ensure
signal-bearing wrappers round-trip through as_python() and make the strict
collective API xfail pass.
There was a problem hiding this comment.
Agreed this is a real gap. It is the python_printer not hoisting generated alloc_window_buffer into a named assignment (and dropping name=), which parse_program requires. Tracked as a strict-xfail UT (proves the gap) + KNOWN_ISSUES, and scoped as a separate printer fix rather than expanding this Python-only PR into a C++ printer change — it does not affect compile/run. I will open a follow-up for the printer hoist.
There was a problem hiding this comment.
@georgebisbas, acknowledged. The finding remains valid, but the printer fix is reasonably out of scope for this PR.
The strict expected-failure test and KNOWN_ISSUES document the current limitation. Please link the follow-up issue here when available. If you want, I can create the GitHub issue and assign it to you.
🐇 ✅
✏️ Learnings added
Learnt from: georgebisbas
Repo: hw-native-sys/pypto PR: 2275
File: python/pypto/language/distributed/op/collective_api.py:74-76
Timestamp: 2026-08-04T08:08:02.027Z
Learning: In `python/pypto/language/distributed/op/collective_api.py`, generated signal buffers from `_fresh_signal()` compile and run correctly, but `as_python()` does not round-trip them through `parse_program()`. The C++ `python_printer` emits `alloc_window_buffer()` inline and drops its `name=` argument, while `parse_program()` requires the allocation in a named simple assignment. The repository tracks this limitation with a strict expected-failure unit test and `KNOWN_ISSUES`; the printer hoist fix is intentionally scoped to a separate follow-up.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
You are interacting with an AI system.
|
Addressed the AI review feedback in Codex (P2): CodeRabbit:
CodeRabbit (Major, printer round-trip): real gap — Verified: UT 18 passed + 1 xfail, ST 5 passed on a2a3sim, pre-commit clean (incl. pyright). |
aec3b80 to
7ec7dfc
Compare
c2463d2 to
11c1a64
Compare
|
Heads-up for after #2279 merges: the "Fresh per call" bullet added here in Currently it says:
#2175 only shipped self-clearing for the InCore composite rail ( Once #2279 merges, the "Fresh per call" claim becomes true — but the citation should point to #2279 (or credit both), and the two bullets should be reconciled so they don't contradict each other. No functional change needed: the wrappers always allocate a fresh signal per call regardless, so this is docs-only. Worth a small follow-up commit here after rebasing onto post-#2279 |
Convenience wrappers over the host collective builtins that auto-manage signal windows, replacing the manual alloc_window_buffer + window + signal boilerplate with a single call. - all_reduce(target, op=Sum, mode="mesh"|"ring", nranks): mesh uses a compiler-synthesized signal (no window); ring validates positive static nranks and rejects non-Sum ops / non-FP32 inputs, allocating a fresh INT32 signal window. - all_gather / reduce_scatter / broadcast / all_to_all / all_to_all_v: allocate a fresh INT32 signal window sized from world_size()/nranks; all_to_all_v requires a positive static nranks (signal [nranks, 1]). - barrier(signal): caller-provided, comm-domain-covered signal. - Each wrapper emits an __auto_<op>_<n> named window; self-clearing credit-barrier signals (merged hw-native-sys#2175) make repeated calls loop-safe. - Literal/overload typing on all_reduce; context-rich ValueError guards (mesh nranks, ring nranks<=0, non-static all_to_all_v nranks). Tests: tests/ut/language/test_collective_api.py; ST tests/st/distributed/test_l3_ergonomic_api.py (a2a3sim). Docs: EN/ZH docs/en|zh/dev/distributed_ops.md "Ergonomic collective API".
11c1a64 to
50b1841
Compare
georgebisbas
changed the title
1 participant
[RFC] feat(distributed): ergonomic
pld.*collective wrappers — auto window+signal; typing, validation, docsStatus: RFC — rebased onto current
origin/main(aa773212) + sim-verified, PRMERGEABLE; NPU verification pending (developer gate P=2/P=4). Head50b18410.Tracks: #1189 (Scheme A, orchestration-level collectives).
Branch:
feat/ergonomic-collective-api— one squashed commit onorigin/main.TL;DR — what this PR is
A pure-Python ergonomic short-form API over the HOST
pld.tensor.*collectives. It removes the per-callalloc_window_buffer+windowsignal boilerplate (and, where needed, the window boilerplate) and adds static typing, early validation, and EN/ZH docs, so distributed host code can be written as:Every wrapper delegates to the exact same HOST builtin — same lowering, same semantics, same goldens. Zero C++/IR/codegen surface.
What this PR offers — and what it deliberately does not
✅ Offers
all_gather,reduce_scatter,broadcast,all_to_all,all_to_all_v,barrier, and ringall_reduce. These have no compiler-side signal synthesis, so the wrapper is the only ergonomic path; signal shapes are encoded so users cannot get them wrong.all_reducerides the compiler's shared-signal synthesis (plan 88 / feat(distributed): share synthesized allreduce signals per data-buffer lineage #2504, merged 2026-08-28): the wrapper delegates to the implicit form, so it gets one shared, hoisted, loop-safe signal with zero Python-side allocation.pld.all_reduceisLiteral["mesh","ring"]+@overload; invalidmode, missingnranksfor ring, ring with non-Sum/FP32, andreduce_scatterwith non-Sum ops are type errors / context-richValueErrors at the call site, not deep in the compiler.pld.all_*resolves inside@pl.programhost bodies (parser expansion verified).❌ Deliberately NOT (kept out of scope)
mode="auto"— mesh vs ring stays explicit until plan 42 (auto-select) lands.The 7 wrappers (
python/pypto/language/distributed/op/collective_api.py)pld.all_reduce(target, *, op, mode="mesh"|"ring", nranks)[2*(NR-1)+1, NR](requires staticnranks)pld.all_gather(local, target)[world_size(), 1]pld.reduce_scatter(target, *, op)[world_size()](rank-1, per builtin requirement)pld.broadcast(target, *, root)[world_size()](rank-1)pld.all_to_all(input, target)[world_size(), 1]pld.all_to_all_v(input, target, send_counts, recv_counts)[world_size(), 1]pld.barrier(signal)Each call allocates a fresh
__auto_<op>_<n>INT32 window — no stale/reused-buffer bugs, no collisions. The self-clearing protocol (#2175/#2279) makes these reusable across calls and loop iterations; fresh-per-call is the safe default (a shared/pooled default is a tracked follow-up — plan 82 I6).Constraints / limitations (documented, not hidden)
pld.barrierrequires an explicit, comm-domain-covered signal — a barrier has no data buffer from whichMaterializeCommDomainScopescan inherit coverage. Zero-argpld.barrier()lands with the coverage fallback (feat(distributed): add HOST-orchestrated builtin.tensor.all_to_all_v #2243).all_reduceuses the compiler's shared signal (feat(distributed): share synthesized allreduce signals per data-buffer lineage #2504 / plan 88) — implicit-signal allreduce insidefor/whileis accepted.pld.all_to_all_vHOST rail is gated on feat(distributed): add HOST-orchestrated builtin.tensor.all_to_all_v #2243 — wrapper delegation is unit-tested; the HOST path is rejected until it merges.as_python()round-trip (IR print→reparse) is currently broken: the Python printer emits the auto signal inline and dropsname=. Captured as a strict-xfail UT +KNOWN_ISSUES; fix is apython_printerhoisting change (own PR). Does not affect compile/run.Testing (re-verified on the rebased
main)tests/ut/language/test_collective_api.py: 18 passed + 1 xfail — signal shape/name generation, fresh-per-call uniqueness, kwarg passthrough, parser resolution, all validation guards, printer-gap xfail.tests/st/distributed/test_l3_ergonomic_api.py: 5/5 passed on a2a3sim — mesh AR, ring AR, broadcast, all_gather, barrier (cross-rank goldens).test_l3_host_tensor_allreduce.py: 9 passed / 3 skipped / 4 pre-existing sim baselines (max,fp16[8193],fp16 Max/Min— Max/Min are NPU-only gates), identical to cleanmain.Design decisions
mode="auto"is plan 42.Follow-ups (tracked, non-blocking)
python_printerhoisting, own PR).pld.all_reduce(plain_tensor)(plan 82 I5).mode="auto"(plan 42), async collectives,group=, FP16/FP8 dtypes.How to review / verify
collective_api.py+tests/ut/language/test_collective_api.py(signals, validation, parser resolution).tests/st/distributed/test_l3_ergonomic_api.py(end-to-end goldens).