feat(distributed): add all_to_all push-based composite intrinsic (InCore only)

[georgebisbas]

Summary

Adds pld.tensor.all_to_all as a push-based InCore composite intrinsic. Every rank pushes its per-destination chunks directly to every peer's window via pld.tile.put (TPUT), then synchronises with a notify/wait barrier. Returns target in-place (window-as-result — same idiom as reduce_scatter / broadcast).

Compiler infrastructure

• LowerTensorAllToAllRule in lower_composite_ops_pass.cpp: 2-phase push-based lowering using pld.tile.put (TPUT-based), emitted directly (not pld.tensor.put) since ConvertTensorToTileOps runs before LowerCompositeOps — mirroring how broadcast/allgather emit pld.tile.get

Tests

• Tier 1 test_l3_all_to_all.py: Hand-rolled push-based using pld.tensor.put (TPUT) + notify/wait barrier

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds a new distributed collective op pld.tensor.all_to_all: C++ type deduction and registration, a composite lowering rule (stage-in, barrier, exchange phases), Python IR-builder and DSL wrapper functions with exports, documentation updates listing twelve ops, and new system tests. Also adds an unrelated Pyright suppression comment.

Changes

all_to_all Collective Op

Layer / File(s)	Summary
Op registration and type deduction src/ir/op/distributed/collective.cpp	Defines pld.tensor.all_to_all with a deduce-type function validating arity, tensor/distributed-tensor shapes, and signal dtype, and registers it with REGISTER_OP.
Composite lowering rule src/ir/transforms/lower_composite_ops_pass.cpp	Adds LowerTensorAllToAllRule implementing stage-in, notify/wait barrier, and exchange phases, registered in LookupCompositeRule.
Python IR and DSL wrappers python/pypto/ir/op/distributed/tensor_ops.py, python/pypto/language/distributed/op/tensor_ops.py	Adds all_to_all IR-builder and DSL wrapper functions, exported via updated __all__ lists.
Op family documentation docs/en/dev/distributed_ops.md, docs/zh-cn/dev/distributed_ops.md	Updates op count from eleven to twelve and adds pld.tensor.all_to_all table entry in both languages.
Hand-rolled DSL all_to_all test tests/st/distributed/collectives/test_l3_all_to_all.py	Adds a manual DSL implementation of a 3-phase mesh all-to-all with golden reference and parametrized test.
Intrinsic all_to_all test tests/st/distributed/collectives/test_l3_tensor_all_to_all_intrinsic.py	Adds a test exercising the pld.tensor.all_to_all intrinsic directly with orchestration helpers and golden comparison.

Unrelated Minor Fix

Layer / File(s)	Summary
Pyright suppression python/pypto/backend/pto_backend.py	Adds a Pyright ignore comment on the Python 3.11+ Traversable import.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Rank
  participant TargetWindow
  participant SignalWindow
  participant OutTensor
  Rank->>TargetWindow: stage input[dest,:] into target[dest,0]
  Rank->>SignalWindow: notify all peers
  Rank->>SignalWindow: wait all signals
  Rank->>TargetWindow: pld.tile.get peer chunk from target[my_rank,0]
  Rank->>OutTensor: write chunk into out[src,0]

Loading

Possibly related PRs

• hw-native-sys/pypto#1568: The all_to_all exchange phase relies on remote tile push/store operations added in this PR.
• hw-native-sys/pypto#1782: Both PRs extend the same distributed collective API plumbing by adding new pld.tensor.* intrinsics wired into lowering and Python exports.

Poem

A rabbit hops from rank to rank,
Chunks exchanged, no thanks to chance—
Stage, then signal, wait in line,
All-to-all, each piece aligns! 🐇
Twelve ops now dance in docs so neat,
Tests confirm the swap complete. 🥕

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the new distributed all_to_all intrinsic and its InCore-only scope.
Description check	✅ Passed	The description is clearly about adding pl d.tensor.all_to_all, its lowering, and tests.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 052a342433

ℹ️ 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".

[gemini-code-assist]

Code Review

This pull request introduces the pld.tensor.all_to_all symmetric all-to-all collective operator as a 4-arg InCore composite, complete with C++ registration, lowering rules, Python bindings, and system tests. The review feedback highlights several critical improvements: validating element type compatibility in the type deducer, adding a post-exchange barrier in the lowering rule to prevent Write-After-Read (WAR) hazards, explicitly constructing 2D offsets instead of using MakeSignalOffsets to avoid coupling, and correcting the documentation to classify the operator as a composite rather than a hardware-supported builtin collective.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (1)

python/pypto/ir/op/distributed/tensor_ops.py (1)

321-337: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Docstring references pld.tile.get, but the actual primitive is pld.tile.remote_load.

The composite is documented as doing "peer-read via pld.tile.get" (Line 333), but the sibling hand-rolled DSL test and the intrinsic test both call pld.tile.remote_load for the exchange phase. This stale naming could mislead future readers cross-referencing the DSL primitive.

📝 Suggested docstring fix

-    3-phase mesh decomposition (stage-in → barrier → peer-read via pld.tile.get).
+    3-phase mesh decomposition (stage-in → barrier → peer-read via pld.tile.remote_load).

🤖 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/ir/op/distributed/tensor_ops.py` around lines 321 - 337, The
all_to_all docstring is using the wrong peer-read primitive name, which can
mislead readers and cross-references. Update the documentation in all_to_all to
describe the exchange phase with the actual primitive name used elsewhere,
pld.tile.remote_load, and keep the rest of the composite description unchanged.

🤖 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 `@docs/en/dev/distributed_ops.md`:
- Line 36: The `pld.tensor.all_to_all` row is mislabeled as a “builtin
collective,” which implies a HOST-dispatch builtin exists like the other
collectives. Update the wording in the distributed ops table to avoid that
classification for `all_to_all` and make it clear it is not backed by a
`builtin.tensor.all_to_all` HOST op, keeping the description aligned with the
`all_to_all` entry in `distributed_ops.md`.

In `@docs/zh-cn/dev/distributed_ops.md`:
- Line 32: The `pld.tensor.all_to_all` row is mislabeled as `builtin collective`
even though this op does not have a HOST-dispatch builtin counterpart. Update
the hardware/classification entry in the distributed ops table for
`pld.tensor.all_to_all` to reflect InCore-only composite support, matching the
intended labeling used for this operation in the docs.

In `@src/ir/op/distributed/collective.cpp`:
- Around line 303-339: The type deduction in DeduceTensorAllToAllType only
checks each argument’s own rank, but it must also verify that input, target, and
out are mutually consistent. Add cross-argument checks in
DeduceTensorAllToAllType so input_type, target_type, and out_type agree on the
expected [NR, SIZE] shape (especially the SIZE dimension) and, where applicable,
compatible dtypes before returning. Use the existing symbols
DeduceTensorAllToAllType, input_type, target_type, signal_type, and out_type to
enforce these invariants and fail fast on mismatches.

---

Nitpick comments:
In `@python/pypto/ir/op/distributed/tensor_ops.py`:
- Around line 321-337: The all_to_all docstring is using the wrong peer-read
primitive name, which can mislead readers and cross-references. Update the
documentation in all_to_all to describe the exchange phase with the actual
primitive name used elsewhere, pld.tile.remote_load, and keep the rest of the
composite description unchanged.

🪄 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

Run ID: a8baea7e-5949-4453-a35e-dd42da862814

📥 Commits

Reviewing files that changed from the base of the PR and between d598b41 and 052a342.

📒 Files selected for processing (9)

• docs/en/dev/distributed_ops.md
• docs/zh-cn/dev/distributed_ops.md
• python/pypto/backend/pto_backend.py
• python/pypto/ir/op/distributed/tensor_ops.py
• python/pypto/language/distributed/op/tensor_ops.py
• src/ir/op/distributed/collective.cpp
• src/ir/transforms/lower_composite_ops_pass.cpp
• tests/st/distributed/collectives/test_l3_all_to_all.py
• tests/st/distributed/collectives/test_l3_tensor_all_to_all_intrinsic.py

[Copilot]

Pull request overview

Adds a new distributed collective, pld.tensor.all_to_all, as an InCore-only composite intrinsic that lowers into a 3-phase mesh exchange (stage-in → notify/wait barrier → peer read via pld.tile.get/TGET-style transfer). This extends the distributed DSL/IR surface area, compiler lowering, and system tests to cover an all-to-all personalized exchange pattern.

Changes:

• Register pld.tensor.all_to_all as a new DistributedOp with a type deducer in C++.
• Add composite lowering in LowerCompositeOps implementing the 3-phase all-to-all decomposition.
• Add Python IR builder + DSL wrapper, plus ST coverage (hand-rolled reference + intrinsic test) and docs row updates.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/ir/transforms/lower_composite_ops_pass.cpp	Adds LowerTensorAllToAllRule and wires it into the composite lowering dispatch.
src/ir/op/distributed/collective.cpp	Registers pld.tensor.all_to_all and introduces DeduceTensorAllToAllType.
python/pypto/ir/op/distributed/tensor_ops.py	Adds IR builder helper to construct pld.tensor.all_to_all(...) calls.
python/pypto/language/distributed/op/tensor_ops.py	Adds the DSL wrapper all_to_all(input, target, signal, out).
tests/st/distributed/collectives/test_l3_all_to_all.py	Adds a hand-rolled 3-phase DSL reference all-to-all system test.
tests/st/distributed/collectives/test_l3_tensor_all_to_all_intrinsic.py	Adds a system test exercising the composite intrinsic end-to-end.
docs/en/dev/distributed_ops.md	Adds pld.tensor.all_to_all to the distributed ops table and updates op count.
docs/zh-cn/dev/distributed_ops.md	Mirrors the EN docs update in Chinese and updates op count.
python/pypto/backend/pto_backend.py	Adds a pyright suppression for Traversable import resolution.

[YunjiQin]

Suggestion: a push-based lowering removes the out parameter (and is simpler + less memory)

Not a blocker — a design alternative worth considering as a follow-up. The current lowering is pull-based (Phase 1 stage-in → barrier → per-peer pld.tile.get → Phase 3.5 post-exchange barrier), mirroring simpler's all_to_all_distributed. A push-based lowering built on pld.tensor.put would be structurally simpler and drop the out argument entirely.

Idea: every rank pushes its per-destination chunks directly into the peers' windows at the "received" slot; the window itself becomes the result and is returned in-place (same idiom as reduce_scatter / broadcast).

# each rank r: put input[dest, :] -> peer `dest`'s window row `r`
for dest in 0..NR-1:
    pld.tensor.put(dst=target, peer=dest, src=input,
                   dst_offsets=[my_rank, 0], src_offsets=[dest, 0], shape=[1, SIZE])
# barrier: each peer notifies signal[peer] after its put; wait on all peers
return target        # in-place rebind; the window IS the result

After the puts, target_r[src, :] == input_src[r, :] == out[r][src, :] — bit-identical to the current golden.

Why it's cleaner than the pull version:

• No Phase 1 stage-in — push straight from input.
• No Phase 3 pull loop and no recv_stage tile.
• No separate out parameter — return target in-place (pld.tensor.all_to_all(input, target, signal)).
• No write-after-read hazard on target during the exchange: target is write-only from peers and input is the read-only source, so the Phase 3.5 read-after barrier is unnecessary; only the entry barrier remains.

Memory: the HCCL window and the out tensor both live in HBM, so the account is total bytes: pull needs (NR+1)*SIZE (a 1-row staging window + NR*SIZE out), push needs NR*SIZE (the window holds the result, no out). all_to_all already requires an NR-row window regardless, so push adds no window cost — it's a net saving plus fewer phases.

The primitives are already registered: pld.tensor.put / pld.tile.put / pld.tile.remote_store.

Trade-offs: the main cost is diverging from the simpler all_to_all_distributed reference this PR mirrors (the Tier-1 hand-rolled test would need rewriting), plus re-validating the pld.tile.put codegen path on hardware (the current pull path is sim-verified P=2).

---

pld.tensor.allgather could get the same treatment. Each rank pushes its single chunk to every peer's slot; the window becomes the [NR, SIZE] gathered result and out drops. One nuance: allgather currently uses a 1-row window (pull only stages each rank's own chunk), so push would grow its window to NR rows. But since window and out share HBM, total HBM still drops from (NR+1)*SIZE to NR*SIZE. The residual trade-offs are the same (reference parity) plus window-reuse lifetime — with push the result occupies the window until consumed, whereas pull frees the 1-row window immediately since the bulk result already lives in out. So for allgather it's a genuine judgement call; for all_to_all (window is NR rows either way) push looks like a clear win.

[georgebisbas]

Implemented. The lowering has been fully rewritten to push-based using pld.tile.remote_store (TSTORE-based). Key changes:

• API: 4-arg → 3-arg (input, target, signal), returns target in-place (window-as-result, same idiom as reduce_scatter/broadcast)
• Lowering: 2-phase (push → barrier) instead of 5-phase — no stage-in, no pull loop, no WAR barrier
• Memory: NR×SIZE (window only) instead of (NR+1)×SIZE (window + out)
• Self-rank: local tile.store for efficiency

[georgebisbas]

pld.tensor.allgather could get the same treatment. Each rank pushes its single chunk to every peer's slot; the window becomes the [NR, SIZE] gathered result and out drops. One nuance: allgather currently uses a 1-row window (pull only stages each rank's own chunk), so push would grow its window to NR rows. But since window and out share HBM, total HBM still drops from (NR+1)*SIZE to NR*SIZE. The residual trade-offs are the same (reference parity) plus window-reuse lifetime — with push the result occupies the window until consumed, whereas pull frees the 1-row window immediately since the bulk result already lives in out. So for allgather it's a genuine judgement call; for all_to_all (window is NR rows either way) push looks like a clear win.

tracked this for a seaprate follow up