feat(st): relocate collective examples into tests/st scene tests

[georgebisbas]

Summary

Move the L3 collective algorithm corpus from examples/workers/l3/ into tests/st/a2a3/tensormap_and_ringbuffer/collectives/ as @scene_test(level=3) classes.

• 5 collective ops: allreduce (5 modes: onephase, twophase, ring, bidirectional_ring, ibing), allgather, reduce_scatter, broadcast, all_to_all
• 19 test classes: one per mode×rank-count (P2 + P4), plus ibing nranks≠2 negative test
• Shared helper (_helpers.py): DRY comm-window domain-allocation and orch-function building
• Kernel sources preserved via git mv; algorithm documentation in per-op READMEs
• Minimal demo: examples/workers/l3/allreduce/ (onephase only) as "how to do a collective"
• Framework: add src/common/ to KernelCompiler.get_orchestration_include_dirs() for platform_comm/comm_context.h

Testing

• 19/19 scene tests pass in pypto3-hw-native-sys:sim Docker (P2 + P4 + ibing error)
• Minimal allreduce example passes in sim Docker
• pre-commit clean (11/11 hooks)
• NPU golden match (a2a3, devices 0-3) — developer gate

Files changed

49 files: +894 / −2032 (net −1138). Bulk of churn is kernel source renames + old example deletions. Net-new logic: ~280 lines in _helpers.py + ~200 lines across 5 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

Run ID: 43446f52-bfb9-4b9b-a193-72260fdf699f

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

This PR removes legacy L3 distributed collective example scripts/tests (allreduce_distributed, allreduce_ring_distributed, allgather_distributed, broadcast_distributed, reduce_scatter_distributed, all_to_all_distributed), rewrites the allreduce example as a minimal onephase implementation, adds shared collective test helpers, and introduces new scene test suites under tests/st/a2a3/tensormap_and_ringbuffer/collectives with corresponding docs. A kernel compiler include-path fix is also included.

Changes

Collective examples consolidation

Layer / File(s)	Summary
Shared test helpers and build fix tests/st/a2a3/tensormap_and_ringbuffer/collectives/_helpers.py, simpler_setup/kernel_compiler.py	Adds sizing constants, orchestration functions (allreduce_orch_fn, generic_collective_orch_fn), golden-output helpers, and argument builders; adds a src/common include directory to orchestration compilation.
Minimal allreduce example rewrite examples/workers/l3/allreduce/main.py, examples/workers/l3/allreduce/README.md	Converts the allgather example into a onephase allreduce example with new sizing constants, kernel/orchestration wiring, golden-check logic, and updated README/run instructions.
Allreduce scene tests tests/st/.../collectives/allreduce/*	Adds test suite covering onephase/twophase/ring/bidirectional_ring/ibing modes for 2 and 4 ranks plus an ibing negative test, with README.
Allgather scene tests tests/st/.../collectives/allgather/*	Adds allgather scene test module with orchestration, args helper, P2/P4 cases, and README.
Broadcast scene tests tests/st/.../collectives/broadcast/*	Adds broadcast scene test module with orchestration, callable config, args helper, P2/P4 cases, and README.
Reduce-scatter scene tests tests/st/.../collectives/reduce_scatter/*	Adds reduce-scatter scene test module with orchestration, callable config, args helper, P2/P4 cases, and README.
All-to-all scene tests tests/st/.../collectives/all_to_all/*	Adds all-to-all scene test module with orchestration, callable config, args helper, P2/P4 cases, and README.
Removal of legacy examples and docs update examples/workers/l3/allreduce_distributed/*, examples/workers/l3/allreduce_ring_distributed/*, examples/workers/l3/allgather_distributed/*, examples/workers/l3/broadcast_distributed/*, examples/workers/l3/reduce_scatter_distributed/*, examples/workers/l3/all_to_all_distributed/*, docs/comm-domain.md, examples/workers/l3/README.md	Deletes legacy distributed example scripts/tests/READMEs and updates references to point at the new scene tests directory.

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

Possibly related PRs

• hw-native-sys/simpler#1092: Overlaps at the same allreduce_distributed implementation and mode-parameterized tests being removed here.
• hw-native-sys/simpler#1271: Modifies the same allreduce_distributed entrypoint that this PR deletes entirely.
• hw-native-sys/simpler#975: Adds the ring allreduce example that this PR removes from allreduce_ring_distributed.

Poem

Old distributed burrows, dug and gone,
New scene-test warrens carry the dawn. 🐇
Onephase allreduce hops in light and lean,
Golden checks glimmer, tidy and clean.
Helpers shared, README trails anew—
This rabbit thumps twice for the tests that came through!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: moving collective examples into scene tests.
Description check	✅ Passed	The description matches the changeset and clearly describes the relocation, helpers, demo, and include-path update.
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.

---

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

[gemini-code-assist]

Code Review

This pull request consolidates and moves several collective algorithm examples (allreduce, allgather, reduce_scatter, broadcast, and all_to_all) from examples/workers/l3/ to tests/st/a2a3/tensormap_and_ringbuffer/collectives/ as scene tests, while leaving a minimal allreduce example in the examples directory. Feedback on the new shared test helpers highlights a lack of validation for nranks and mode_id parameters in allreduce_orch_fn and generic_collective_orch_fn, which could lead to out-of-bounds memory allocation or access on the device.

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: 4

🤖 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/comm-domain.md`:
- Around line 153-154: The docs link in comm-domain.md points too far up the
directory tree and resolves outside the repo when clicked. Update the allreduce
tests link target used in the bullet so it stays relative to the docs directory
and lands on the intended allreduce scene tests, using the same link text but
the corrected path from the markdown entry.

In `@examples/workers/l3/allreduce/README.md`:
- Line 9: The relative link in the allreduce README points to the wrong location
because it does not go far enough up to the project root. Update the markdown
link under the allreduce example so it uses the correct number of
parent-directory hops from the README location, ensuring it resolves to the
tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce target instead of
an examples/tests path.

In `@tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/README.md`:
- Line 89: The link in the allreduce README is relative to the wrong depth and
resolves to a missing target. Update the markdown reference in the README so it
climbs back to the repository root before pointing to the minimal allreduce
example, and verify the corrected path from this README location still reaches
examples/workers/l3/allreduce/.

In `@tests/st/a2a3/tensormap_and_ringbuffer/collectives/reduce_scatter/README.md`:
- Around line 12-19: The README description for reduce_scatter does not match
the helper’s integer-based contract and input sizing, so update the
“Input/Output” and “Golden Check” text to use the same `nranks * COUNT_PER_RANK`
per-rank input size and the exact integer formula used by the shared helper.
Refer to the reduce_scatter example text in the README and align the closed-form
result wording with the implementation so it no longer implies a conflicting
chunk size or arithmetic model.

🪄 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: 34964106-4b1a-42ae-9d23-694839b2fa61

📥 Commits

Reviewing files that changed from the base of the PR and between a5895d8 and 1488f13.

📒 Files selected for processing (56)

• docs/comm-domain.md
• examples/workers/l3/README.md
• examples/workers/l3/all_to_all_distributed/README.md
• examples/workers/l3/all_to_all_distributed/main.py
• examples/workers/l3/all_to_all_distributed/test_all_to_all.py
• examples/workers/l3/allgather_distributed/README.md
• examples/workers/l3/allgather_distributed/test_allgather.py
• examples/workers/l3/allreduce/README.md
• examples/workers/l3/allreduce/main.py
• examples/workers/l3/allreduce_distributed/README.md
• examples/workers/l3/allreduce_distributed/main.py
• examples/workers/l3/allreduce_distributed/test_allreduce.py
• examples/workers/l3/allreduce_ring_distributed/main.py
• examples/workers/l3/broadcast_distributed/README.md
• examples/workers/l3/broadcast_distributed/main.py
• examples/workers/l3/broadcast_distributed/test_broadcast.py
• examples/workers/l3/reduce_scatter_distributed/README.md
• examples/workers/l3/reduce_scatter_distributed/__init__.py
• examples/workers/l3/reduce_scatter_distributed/main.py
• examples/workers/l3/reduce_scatter_distributed/test_reduce_scatter.py
• simpler_setup/kernel_compiler.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/__init__.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/_helpers.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/all_to_all/README.md
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/all_to_all/__init__.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/all_to_all/kernels/aiv/all_to_all_kernel.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/all_to_all/kernels/orchestration/all_to_all_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/all_to_all/test_all_to_all.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allgather/README.md
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allgather/__init__.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allgather/kernels/aiv/allgather_kernel.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allgather/kernels/orchestration/allgather_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allgather/test_allgather.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/README.md
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/__init__.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/kernels/aiv/allreduce_bidirectional_ring_kernel.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/kernels/aiv/allreduce_ibing_kernel.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/kernels/aiv/allreduce_onephase_kernel.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/kernels/aiv/allreduce_ring_kernel.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/kernels/aiv/allreduce_twophase_kernel.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/kernels/orchestration/allreduce_bidirectional_ring_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/kernels/orchestration/allreduce_ibing_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/kernels/orchestration/allreduce_onephase_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/kernels/orchestration/allreduce_ring_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/kernels/orchestration/allreduce_twophase_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/test_allreduce.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/broadcast/README.md
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/broadcast/__init__.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/broadcast/kernels/aiv/broadcast_kernel.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/broadcast/kernels/orchestration/broadcast_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/broadcast/test_broadcast.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/reduce_scatter/README.md
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/reduce_scatter/__init__.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/reduce_scatter/kernels/aiv/reduce_scatter_kernel.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/reduce_scatter/kernels/orchestration/reduce_scatter_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/reduce_scatter/test_reduce_scatter.py

💤 Files with no reviewable changes (20)

• examples/workers/l3/all_to_all_distributed/README.md
• examples/workers/l3/allgather_distributed/test_allgather.py
• examples/workers/l3/broadcast_distributed/README.md
• examples/workers/l3/reduce_scatter_distributed/init.py
• examples/workers/l3/allreduce_distributed/test_allreduce.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/all_to_all/init.py
• examples/workers/l3/allreduce_distributed/README.md
• examples/workers/l3/reduce_scatter_distributed/README.md
• examples/workers/l3/all_to_all_distributed/test_all_to_all.py
• examples/workers/l3/allreduce_ring_distributed/main.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allgather/init.py
• examples/workers/l3/broadcast_distributed/test_broadcast.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/init.py
• examples/workers/l3/allreduce_distributed/main.py
• tests/st/a2a3/tensormap_and_ringbuffer/collectives/allreduce/init.py
• examples/workers/l3/reduce_scatter_distributed/test_reduce_scatter.py
• examples/workers/l3/all_to_all_distributed/main.py
• examples/workers/l3/broadcast_distributed/main.py
• examples/workers/l3/reduce_scatter_distributed/main.py
• examples/workers/l3/allgather_distributed/README.md

[ChaoWao]

Two items from review — both small but worth resolving before merge.

1. Orphaned dead doc: examples/workers/l3/allreduce_distributed/README.md

That directory's main.py, test_allreduce.py, and __init__.py were deleted, but README.md was left behind — it's now the only file in the dir (git ls-tree -r pr-head examples/workers/l3/allreduce_distributed/ returns just the README). It still documents the deleted main.py/test_allreduce.py layout and the old -m onephase|twophase|ring CLI, and the parent examples/workers/l3/README.md link to it was removed, so it's unreachable, stale documentation.

Per the repo's doc-consistency rule ("delete docs when their referent is gone"), please git rm examples/workers/l3/allreduce_distributed/README.md.

2. Silent a5 / a5sim test-coverage regression

On main, the collective tests parametrized:

• allreduce P2 + ibing-negative: @pytest.mark.platforms(["a2a3sim", "a2a3", "a5sim", "a5"])
• allreduce P4, allgather, reduce_scatter, broadcast, all_to_all: ["a2a3sim", "a2a3", "a5sim"]

Every new scene-test CASES entry in tests/st/a2a3/tensormap_and_ringbuffer/collectives/** lists only ["a2a3sim", "a2a3"], and there's no tests/st/a5/.../collectives/ mirror. Net effect: a5sim coverage for all 5 ops is dropped, and a5-onboard allreduce coverage is dropped — a real reduction vs main that isn't called out in the PR body.

This may be intentional (collectives treated as a2a3-only), but it's currently silent. Please either restore the a5sim cases (a mirrored tests/st/a5/.../collectives/) or state explicitly in the PR body / a README that collectives are a2a3-only now.

(Minor, separate: the PR body says "19 test classes" — the actual count is 18: allreduce 10, allgather/reduce_scatter/broadcast/all_to_all 2 each.)

[ChaoWao]

examples/workers/l3/allreduce/ should be self-contained — give it its own kernel

Keep the example — a minimal onephase allreduce is a good demo of the L3 Worker collective interface, which is exactly what examples/workers/ is for. But right now it isn't self-contained: main.py reaches four levels up into the test tree for its kernel sources —

_COLLECTIVES_DIR = os.path.join(HERE, "..","..","..","..",
    "tests","st","a2a3","tensormap_and_ringbuffer","collectives","allreduce")
_KERNEL_AIV  = os.path.join(_COLLECTIVES_DIR, "kernels","aiv","allreduce_onephase_kernel.cpp")
_KERNEL_ORCH = os.path.join(_COLLECTIVES_DIR, "kernels","orchestration","allreduce_onephase_orch.cpp")

This inverts the dependency direction: an examples/ demo now depends on tests/, so any future rename/reorg of the collectives test tree silently breaks the example, and a plain grep won't surface the coupling. examples/workers/ is meant to hold standalone demos of how to use a worker interface / a simpler feature — each example should stand on its own.

Please make it self-contained:

1. Add the example's own kernel copy under examples/workers/l3/allreduce/kernels/:

   • kernels/aiv/allreduce_onephase_kernel.cpp
   • kernels/orchestration/allreduce_onephase_orch.cpp

   (The onephase sources only include from pto-isa / src/common / the runtime include dirs — no reach back into tests/ — so a local copy compiles unchanged. The small duplication of one onephase kernel is the accepted cost of a standalone example.)

2. Point _KERNEL_AIV / _KERNEL_ORCH at the local kernels/ dir and drop the ../../../../tests/... _COLLECTIVES_DIR reach.

README should reflect the examples/workers/ principle

The purpose of every examples/workers/ entry is to show how to use the worker interface / demonstrate a simpler feature, and the README should say so — not read as an algorithm reference that points at tests/ as the source of truth. Please reframe examples/workers/l3/allreduce/README.md around what the demo teaches: e.g. the L3 Worker(level=3) flow — register → worker.init() → orch.allocate_domain(...) with a CommBufferSpec comm window → submit_next_level per rank → golden check. Keep the pointer to the scene tests as "for the full algorithm corpus", but the example + its README stand on their own as an interface-usage demo, not as documentation hanging off the test tree.

[georgebisbas]

@ChaoWao thanks for the thorough review — all four items addressed on feat/collectives-scene-tests (rebased on latest main, pushed to fork at 2c01e88):

1. Orphaned allreduce_distributed/README.md
git rm'd. The rest of that directory was already deleted in the original relocation; the README was the only leftover.

2. Silent a5 / a5sim test-coverage regression
Added tests/st/a2a3/tensormap_and_ringbuffer/collectives/README.md documenting that collective scene tests are scoped to a2a3 / a2a3sim only. a5 coverage is not yet exercised at this level. I'll also call this out in the PR body.

3. examples/workers/l3/allreduce/ is now self-contained
Copied the onephase kernel + orchestration shim locally into examples/workers/l3/allreduce/kernels/. main.py now uses os.path.join(HERE, "kernels/...") instead of reaching four levels up into tests/. The example compiles and runs standalone — verified in sim Docker, all ranks matched golden ✅.

4. README reframed around the examples/workers/ principle
Rewritten to focus on what the demo teaches: the L3 Worker flow (register → init → allocate_domain with CommBufferSpec → submit_next_level per rank → golden check). Still points to scene tests for the full algorithm corpus, but the example + README stand on their own as an interface-usage demo.

Also fixed the PR body test-class count (18, not 19). All 18 scene tests pass in sim Docker, pre-commit is clean (11/11).