[RFC] feat(distributed): ergonomic pld.* collective wrappers (auto-managed signals) - #2275
[RFC] feat(distributed): ergonomic pld.* collective wrappers (auto-managed signals)#2275georgebisbas wants to merge 1 commit 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:
📝 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)
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
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".
7ec7dfc to
c2463d2
Compare
1 participant
[RFC] feat(distributed): ergonomic
pld.*collective wrappers — auto-managed signalsStatus: RFC — sim-verified; NPU verification pending.
Tracks: #1189 (Scheme A, orchestration-level collectives).
Branch:
feat/ergonomic-collective-api— 6 commits, rebased onorigin/mainwith #2175 (self-clearing credit-barrier protocol,e10ce252).Summary
This PR adds ergonomic short forms over the existing HOST
pld.tensor.*collectives that auto-allocate a fresh, correctly-shaped INT32 signal window per call, removing the per-callalloc_window_buffer+windowsignal boilerplate. The wrapper then delegates to the exact samepld.tensor.*HOST builtin — same lowering, same semantics, same goldens.Why this matters
The host-collective infrastructure landed across #1798 / #1829 / #1997 / #2094, but every call still required hand-rolled signal plumbing:
This PR is the user-facing layer of #1189's Scheme A: it turns raw, powerful builtins into something distributed kernels can actually call — the foundation the ergonomic roadmap builds on.
What this PR offers — in detail
1. The 7 wrappers (
python/pypto/language/distributed/op/collective_api.py)pld.all_reduce(target, *, op, mode="mesh"|"ring", nranks)[world_size, 1]); ring: auto[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)All signal shapes match the corresponding HOST builtin's requirement exactly, encoded in code so users cannot get them wrong.
2. Auto-managed, fresh signals — correctness by construction
__auto_<op>_<n>INT32 window; no stale/reused-buffer bugs, no collisions.3. Early validation — errors where you typed them, not deep in the compiler
pld.all_reduceisLiteral["mesh","ring"]+@overload—mode="tree"andmode="ring"withoutnranksare type errors.ValueError): invalidmode;nranksrequired for ring / rejected for mesh; ring isReduceOp.Sum+ FP32 only;reduce_scatterisReduceOp.Sumonly.4. Exports & discoverability
Wired through
op/unified_ops.py→op/__init__.py→distributed/__init__.py, sopld.all_reduce(...)resolves inside@pl.programhost bodies (parser expansion topld.tensor.*verified).5. Documentation (EN + ZH, parity-checked)
docs/{en,zh}/dev/distributed_ops.md: API table, a worked "publish → collective → consume" example, mesh-vs-ring guidance, and an honest constraints list.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).allreduceis not self-clearing, soSynthesizeAllReduceSignalsstill rejects an allreduce insidefor/whileloops. fix(ir): make composite collective barrier signals reusable #2175's loop-safety applies to the InCore composite rail only (not targeted here).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 (window(alloc_window_buffer(...))) and dropsname=, whichparse_programrejects. Captured as a strict-xfail UT +KNOWN_ISSUES; the fix is apython_printerhoisting change (own PR). Does not affect compile/run.Testing (all green on the rebased base)
tests/ut/language/test_collective_api.py: 15 passed + 1 xfail — signal shape/name generation, fresh-per-call uniqueness, kwarg passthrough, parser resolution, all validation guards, and the documented printer-gap xfail.tests/st/distributed/test_l3_ergonomic_api.py: 5 passed on a2a3sim — mesh AR, ring AR, broadcast, all_gather, barrier (cross-rank goldens).test_l3_host_tensor_allreduce.py: 11 passed / 1 skipped.Design decisions
mode="auto"will automate selection.Follow-ups (tracked, non-blocking)
pld.all_reduce(plain_tensor)to hide the data side too.mode="auto"default, async collectives,group=parameter, 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).