fix(distributed): close window-aliasing gap and O(N²) pass complexity in HOST collectives

[georgebisbas]

A deep review of the distributed collectives stack (allreduce, barrier, broadcast, reduce_scatter, allgather, all_to_all) surfaced one real correctness gap and one real complexity-budget violation, both silent today. This PR fixes both, plus the consistency issue that caused the first one.

1. Window aliasing was unchecked for most HOST collectives — a real silent-corruption risk

Before this PR, LowerHostTensorCollectives's CheckDistinctInputTargetWindows was wired into only 2 of 6 collectives (allgather, all_to_all), and even there it only checked input vs target — never vs signal. allreduce, barrier, broadcast, reduce_scatter had no aliasing check at all.

Concretely: nothing stopped a user from writing

pld.tensor.allreduce(data, signal, op=pld.ReduceOp.Sum)

where data and signal are two different pld.window() views over the same alloc_window_buffer. That's a real cross-process race — the reduce's data write and the barrier's notify/wait control write land in the same memory — and it would compile cleanly and fail silently or corrupt data at runtime instead of raising a clear compile-time error.

Fix: HostCollectiveRule::scope_buffers already enumerates every window operand for every op (it existed to resolve the enclosing CommDomainScopeStmt) — this PR reuses that same list to run a new, generic CheckPairwiseDistinctWindows helper once, in LowerCollective, covering all 6 collectives uniformly. One shared 15-line helper instead of one-off per-op checks that get missed (this is exactly the bug class a recent all_to_all_v PR needed 3 rounds of review to close for just that one op — this generalizes the fix so it doesn't need to be rediscovered per-collective).

2. O(N²) pass complexity — a direct violation of this repo's own pass-complexity.md

MaterializeCommDomainScopes's domain-clustering phase did a linear scan of the pending domains list per allocation to find a matching descriptor. For a program with many independently-scoped comm domains, this is O(allocations × domains) = O(N²).

Fix: DeviceDescriptor already has operator<, so this PR adds a std::map<DeviceDescriptor, size_t> index alongside pending, turning the lookup into O(log domains) — squarely within the project's own "ordered map/set lookups are fine" complexity allowance. No behavior change; a new 3-domain regression test pins the clustering order.

3. Signal validation was inconsistent between the public op and its HOST builtin (root cause of gap #1's blind spots)

pld.tensor.barrier/broadcast/reduce_scatter's public deducers checked signal dtype only (no rank/shape check), while their builtin counterparts required rank-1 only — and allreduce/allgather/all_to_all each separately inlined their own near-identical "rank-1-or-2" check. Net effect: a rank-2 [NR, 1] signal (the pattern used by allreduce/allgather) would type-check fine for barrier/broadcast/reduce_scatter, then fail 3 passes later inside LowerHostTensorCollectives with a confusing, late error.

Fix: one CheckSignalDistributedTensor helper (rank-1-or-2, INT32), called from all 12 signal-validation sites — every public op and every builtin. Removes ~90 lines of duplicated inline validation and makes this class of public/builtin mismatch structurally impossible going forward.

4. broadcast's root was never bounds-checked against the device count

Only root >= 0 was validated. root=100 on a 4-rank job compiled cleanly and would fail at runtime as an out-of-bounds remote read. This PR adds a root < participating_devices check on the explicit static-device-subset domain (mirroring the existing signal-capacity check's scope — the fully-dynamic "all device" domain is left as a documented limitation, same as the existing signal-capacity gap there, since there's no compile-time device count to check against in that case).

What's intentionally not in this PR

Per a broader review (docs/discussion, not code), several other findings were evaluated and deliberately deferred — a CollectiveOpBuilder abstraction (premature without an 8th collective), splitting collective.cpp into per-op files (conflicts with the in-flight #2243), a runtime-assert IR primitive for the dynamic-domain gaps mentioned above (real gap, but new IR/runtime infra, out of scope here). Happy to open tracking issues for these if useful.

Test plan

• 12 new unit tests: aliasing-rejection for allreduce/broadcast/reduce_scatter (allgather/all_to_all already had partial coverage, now exercised by the same generic path), rank-2-signal-accepted + invalid-rank-rejected for barrier/broadcast/reduce_scatter's public deducers, broadcast root-out-of-range rejection, 3-domain clustering regression for the complexity fix.
• Full unit suite (8542 tests) passes in the pypto3-hw-native-sys:sim Docker image.
• clang-format --dry-run --Werror and clang-tidy clean on all touched/new C++ files.
• ruff check / ruff format --check clean on all touched Python test files.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 34fe60f3-93e6-48d1-bae2-d2c6f95e9119

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

Collective validation and lowering

Layer / File(s)	Summary
Shared distributed signal validation src/ir/op/distributed/collective.cpp, tests/ut/ir/test_distributed_ops.py	Collective operations share INT32 signal validation. Rank-1 signals and rank-2 [world_size, 1] signals are accepted. Invalid ranks and non-unit second dimensions are rejected.
Host collective window validation include/pypto/ir/transforms/utils/window_alias_utils.h, src/ir/transforms/utils/window_alias_utils.cpp, src/ir/transforms/lower_host_tensor_collectives_pass.cpp, CMakeLists.txt, tests/ut/ir/transforms/test_lower_host_tensor_collectives.py	Named window buffers support pairwise alias checks. Host collective rules use named buffers, validate broadcast roots, and reject aliased windows.
Communication-domain grouping src/ir/transforms/materialize_comm_domain_scopes_pass.cpp, tests/ut/ir/transforms/test_materialize_comm_domain_scopes.py	Pending communication domains use descriptor-to-index lookup. Materialization preserves allocation order and creates separate scopes for distinct descriptors.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant CollectiveOp
  participant HostCollectiveLowering
  participant CommunicationDomainScopes
  CollectiveOp->>HostCollectiveLowering: collective buffers and signal
  HostCollectiveLowering->>HostCollectiveLowering: validate signal, root, and window aliases
  HostCollectiveLowering->>CommunicationDomainScopes: resolve communication scope
  CommunicationDomainScopes-->>HostCollectiveLowering: materialized scope and buffer slots

Loading

Possibly related PRs

• hw-native-sys/pypto#1997: Shares changes in collective validation, host lowering, and communication-domain handling.
• hw-native-sys/pypto#2243: Shares host collective lowering and window-alias validation changes.

Poem

A rabbit checks each window pair,
No hidden aliases linger there.
Signals take their proper shape,
Broadcast roots stay in their map.
Scopes group by device view—
Hop, hop, the checks pass through!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.13% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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 identifies the two primary fixes: window aliasing and O(N²) complexity in HOST collectives.
Description check	✅ Passed	The description directly explains the aliasing, validation, complexity, root-bound, and regression-test changes.

---

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.}

[Copilot]

Pull request overview

Improves HOST collective safety and pass scalability.

Changes:

• Adds generic window-alias rejection and broadcast root validation.
• Unifies most collective signal-shape checks.
• Replaces quadratic domain clustering lookup with a map.

Reviewed changes

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

Show a summary per file

File	Description
CMakeLists.txt	Registers alias utility source.
include/pypto/ir/transforms/utils/window_alias_utils.h	Declares window alias validation.
src/ir/transforms/utils/window_alias_utils.cpp	Implements pairwise checks.
src/ir/transforms/materialize_comm_domain_scopes_pass.cpp	Adds indexed domain lookup.
src/ir/transforms/lower_host_tensor_collectives_pass.cpp	Applies alias and root checks.
src/ir/op/distributed/collective.cpp	Centralizes most signal validation.
tests/ut/ir/transforms/test_materialize_comm_domain_scopes.py	Tests three-domain ordering.
tests/ut/ir/transforms/test_lower_host_tensor_collectives.py	Tests alias and root rejection.
tests/ut/ir/test_distributed_ops.py	Tests signal-rank validation.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: daf98aa373

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

[Copilot]

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (4)

src/ir/op/distributed/allreduce.cpp:105

• This still accepts a non-unit second dimension for the default mode="mesh", not just for ring mode. The new “ring mode” unit test calls allreduce without mode="ring", and LowerTensorAllReduceRule performs no mesh signal-shape check, so [NR, 2] passes the public mesh path while the HOST builtin rejects it. Branch on the mode kwarg here: use CheckSignalDistributedTensor for mesh and the shape-only helper for ring, then make the non-unit test explicitly select ring and add a mesh rejection case.

    CheckSignalDistributedTensorShape(As<DistributedTensorType>(args[1]->GetType()), "pld.tensor.allreduce");

src/ir/transforms/lower_host_tensor_collectives_pass.cpp:109

• This introduces a new public validation rule, but the public broadcast docstring still says only that root “Must be non-negative” (python/pypto/language/distributed/op/tensor_ops.py:688). Document that an explicit static comm domain also requires root < participating device count, while all-device domains cannot be upper-bound checked at compile time.

void CheckBroadcastRootInRange(const CallPtr& call, size_t participating_devices) {
  auto root_value = call->GetKwarg<int>("root");
  CHECK_SPAN(root_value >= 0 && static_cast<size_t>(root_value) < participating_devices, call->span_)
      << "pld.tensor.broadcast root (" << root_value << ") must be a valid rank in [0, "
      << participating_devices << ") for this explicit device subset";

src/ir/op/distributed/signal_utils.h:26

• This overview contradicts the helpers below: ring allreduce deliberately accepts [2*(NR-1), NR], so not every collective accepts only the same two shapes and the public/builtin validators can differ by mode. Describe the mesh/single-column contract separately from allreduce’s ring exception.

 * cross-rank barrier slot. All of them must accept the same two shapes
 * (rank-1 [world_size] or rank-2 [world_size, 1]) — this lives in one place,
 * shared across translation units (allreduce.cpp, collective.cpp), so the
 * public op and its builtin can never disagree on what a signal is allowed

src/ir/op/distributed/collective.cpp:131

• The added rank-2 acceptance tests stop at these public deducers; [NR, 1] already passed here before this PR because the old code checked only dtype. No test sends a rank-2 signal through HOST lowering to exercise the changed builtin validators—the exact late failure this PR intends to eliminate. Add HOST-lowering success tests with [NR, 1] for barrier, broadcast, and reduce-scatter.

  CheckSignalDistributedTensor(As<DistributedTensorType>(args[0]->GetType()), "pld.tensor.barrier");

[georgebisbas]

Addressed Copilot's second-pass suppressed findings in 5cefb5f:

1. allreduce.cpp mesh vs ring signal validation — public deducer now branches on mode: mesh uses strict CheckSignalDistributedTensor; ring uses shape-only. Added mesh [NR, 2] rejection + explicit mode="ring" acceptance tests.
2. broadcast root docstring — documented the static-subset root < participating device count check and the all-device-domain compile-time limitation.
3. signal_utils.h overview — rewritten to describe the mesh/single-column contract separately from allreduce's ring exception.
4. HOST-lowering [NR, 1] success paths — added barrier/broadcast/reduce_scatter tests that exercise both the public deducer and builtin validators through LowerHostTensorCollectives.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated no new comments.

Suppressed comments (4)

src/ir/op/distributed/allreduce.cpp:115

• Every value other than "ring" is routed through the mesh validator. For example, mode="typo" with a ring-shaped [2, 2] signal now reports shape[1] must be 1, so the intended invalid-mode diagnostic in LowerCompositeOps is never reached. Validate mode before selecting its signal contract (including the no-signal form).

    if (mode == "ring") {
      CheckSignalDistributedTensorShape(signal_type, "pld.tensor.allreduce");
    } else {
      CheckSignalDistributedTensor(signal_type, "pld.tensor.allreduce");

src/ir/transforms/lower_host_tensor_collectives_pass.cpp:341

• The pass documentation is now stale: docs/en/dev/passes/40-lower_host_tensor_collectives.md still documents only stage/data distinctness and signal capacity, omitting both the new pairwise-alias rejection and this explicit-subset root bound. Update the English ground-truth document and its Chinese mirror to match the new user-facing checks.

    CheckStaticSignalCapacity(call, rule.signal_expr(call), scope->devices_.size());
    if (rule.extra_static_check) rule.extra_static_check(call, scope->devices_.size());

tests/ut/ir/transforms/test_materialize_comm_domain_scopes.py:1167

• The test now creates four allocations after adding buf_d, so the three_allocs name no longer describes its setup. Rename it to reflect the four allocations and three descriptors/groups.

def test_three_allocs_different_descriptors_three_groups():

src/ir/op/distributed/signal_utils.h:75

• This changes the accepted public signal contract for barrier, broadcast, and reduce_scatter, but their Python API docs still only describe an INT32 barrier tensor and do not state the enforced rank-1 [NR] or rank-2 [NR, 1] shapes. Document the new contract so users can understand and fix these newly surfaced validation errors.

inline void CheckSignalDistributedTensor(const DistributedTensorTypePtr& signal_type,
                                         const std::string& op_name) {
  CheckSignalDistributedTensorShape(signal_type, op_name);
  if (signal_type->shape_.size() == 2) {