docs: fix distributed-op docstrings and the ReduceOp type-stub

[georgebisbas]

Summary

Split out of #2162 (which is being broken into smaller PRs — see that
PR's description). This piece is fully independent of the rest: pure
Python docstring / type-stub corrections found while writing the
distributed-programming and benchmarking user-manual chapters.

• ir.pyi: ReduceOp's docstring claimed only Sum is operational
  and that Max/Min/Prod "raise a C++ deducer error if passed to
  any collective." That's wrong — allreduce accepts all four (both
  InCore and HOST paths); only reduce_scatter is actually Sum-only.
  Restated per-operation.
• tensor_ops.py: per-operation signal-shape docstrings. Previously
  uniform ("rank-1 or rank-2" for every collective) — actually only
  allreduce/allgather/all_to_all accept both ranks on the HOST
  builtin path; barrier/broadcast/reduce_scatter are rank-1-only.
• system_ops.py, tile_ops.py, dsl_api.py, bench.py,
  distributed_compiled_program.py: docstring completeness fixes —
  DistributedConfig field documentation, dispatch examples, the
  spmd three-form table, and a couple of missing code examples.

Addressed review feedback (see PR discussion) in a follow-up commit:

• Added real HOST-only parser enforcement for alloc_window_buffer
  (previously only claimed by the docstring), with tests.
• Fixed the num_sub_workers=0 docstring to match actual runtime
  semantics (max(num_sub_workers, len(sub_worker_fns))).
• Corrected a wrong issue reference in the allreduce signal-reuse note.
• Added a pointer to where the mesh partial-valid / symbolic-extent
  constraints are documented (unchanged, just relocated out of this
  docstring).
• Converted two non-rendering RST tables to Markdown (the docs site
  renders docstrings as Markdown, not RST), extending the
  BenchmarkStats one to cover the per-dispatch API added by Add per-dispatch benchmark timing and resolve callable names #2223.
• Fixed DistributedConfig's docstring section header
  (Fields: → Attributes:) and stale allgather prose referring to
  a since-renamed parameter.
• Dropped an unverifiable "legacy reasons" history claim from notify's
  docstring note.

notify/wait's signal parameter type (Tensor vs. DistributedTensor)
was intentionally not narrowed here: doing so broke pyright on two
legitimate call sites where a loop-carried DistributedTensor value
loses its static type through pl.range/pl.while_. That's tracked
separately in #2234, along with the type-narrowing itself.

No behavior changes beyond the alloc_window_buffer scope check
(makes parse-time behavior match what was already enforced at runtime).

Verification

• python -c "import ast; ast.parse(...)" on every changed file
• pre-commit run --files <changed files> (ruff check/format, pyright,
  headers) — all pass
• Full pytest tests/ut suite (8433 passed, 13 skipped, 0 failed) run
  via the pypto3-hw-native-sys:sim Docker image
• Manually rendered the two converted docstrings through
  markdown.markdown(..., extensions=["tables"]) to confirm they now
  produce real <table> elements (previously collapsed to one
  paragraph under RST simple-table syntax)

[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: 83e01825-c687-447c-870a-a9812385275a

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

Documentation-only updates clarify distributed execution patterns, host-versus-InCore API constraints, collective signal shapes, SPMD forms, reduction support, remote-operation examples, and benchmark accessors. Runtime logic, public signatures, enum values, and serialization behavior remain unchanged.

Changes

Distributed Documentation

Layer / File(s)	Summary
Runtime and DSL usage documentation python/pypto/ir/distributed_compiled_program.py, python/pypto/language/distributed/op/system_ops.py, python/pypto/language/distributed/op/tile_ops.py, python/pypto/language/dsl_api.py	Distributed configuration and dispatch examples, orchestration constraints, notification terminology, remote-operation examples, and spmd usage forms are documented.
Collective and window API documentation python/pypto/language/distributed/op/tensor_ops.py	Window APIs, allreduce behavior, signal reuse constraints, chunking notes, and host/InCore signal shapes are clarified.
IR and benchmark reference documentation python/pypto/pypto_core/ir.pyi, python/pypto/runtime/bench.py	Supported reduction values and BenchmarkStats accessors are listed in reference documentation.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Possibly related PRs

• hw-native-sys/pypto#1947: Updates related pl.spmd(...) usage documentation.
• hw-native-sys/pypto#2084: Introduces behavioral changes reflected in the updated mesh allreduce documentation.

Poem

A rabbit reads the docs with glee,
From signals shaped for every tree.
Remote loads hop rank to rank,
While SPMD forms line up in sync.
No code changed—just clearer hay!

🚥 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 documentation and ReduceOp type-stub corrections in the changeset.
Description check	✅ Passed	The description accurately explains the documentation corrections, limited behavior change, and verification results.

---

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

[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 `@python/pypto/ir/distributed_compiled_program.py`:
- Around line 95-106: Update the “Persistent dispatch” example to use the
documented prepare-and-run API: prepare the compiled program before creating the
runtime, then invoke rt.run(compiled, ...) inside the loop instead of calling
the runtime directly. Keep the one-shot dispatch example unchanged.

In `@python/pypto/language/distributed/op/system_ops.py`:
- Around line 143-150: Update the documentation note around the notify signature
to say that only the first operand, the signal/target DistributedTensor, must be
passed positionally; keep the example’s peer= keyword usage and avoid implying
that all arguments must be positional.

In `@python/pypto/language/distributed/op/tile_ops.py`:
- Around line 146-150: Update the documentation example around
pld.tile.remote_store so its destination uses a data-window name rather than
signal, and do not imply that remote_store provides synchronization; if
presenting a signal handoff, show the required notify/wait protocol separately.

In `@python/pypto/runtime/bench.py`:
- Line 327: Update the documentation for stats.all_zero_device to state that it
is true only when samples exist and every device_wall_us sample is zero;
explicitly describe the empty-sample case as false.

🪄 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: ddc042b3-0ae1-49f3-b3a1-eabebbef430d

📥 Commits

Reviewing files that changed from the base of the PR and between 84d11f1 and b3d1e0f.

📒 Files selected for processing (7)

• python/pypto/ir/distributed_compiled_program.py
• python/pypto/language/distributed/op/system_ops.py
• python/pypto/language/distributed/op/tensor_ops.py
• python/pypto/language/distributed/op/tile_ops.py
• python/pypto/language/dsl_api.py
• python/pypto/pypto_core/ir.pyi
• python/pypto/runtime/bench.py

[YunjiQin]

Reviewed the whole diff against the implementation. The core corrections check out — pld.tensor.allreduce really does accept all four ReduceOp values on both rails (allreduce.cpp:105, collective.cpp:66, and all four codegen variants at distributed_ops_codegen.cpp:50-62) while reduce_scatter is Sum-only (collective.cpp:473, lower_composite_ops_pass.cpp:1582); the per-op signal ranks match the deducers exactly (rank-1-or-2 for allreduce/allgather/all_to_all at collective.cpp:106/660/745, rank-1-only via CheckSignalDistributedTensor for barrier/broadcast/reduce_scatter at collective.cpp:77,514,546,590); every BenchmarkStats accessor in the new table exists with the described semantics; and the new "all allreduce calls in loops are rejected" is more accurate than the text it replaces (CheckAllReduceLoopUse, lower_composite_ops_pass.cpp:1966-1971, fires for any allreduce, not just the signal-omitting HOST form).

Three newly introduced statements are not backed by the code, though.

Must fix

1. alloc_window_buffer is not HOST-only-enforced. The added note says calling it inside InCore "raises a parser error". There is no such check anywhere: _validate_pld_op_call only rejects non-assignment use (ast_parser.py:8001-8008), world_size is the only op with a level/scope check (ast_parser.py:8012-8023), _parse_alloc_window_buffer_assignment (ast_parser.py:1771-1821) validates LHS/uniqueness/name= only, and the deducer (memory.cpp:69-75) checks arity and the name kwarg. Either soften to "currently only supported/exercised in HOST orchestration", or add a real check mirroring world_size plus a UT.

2. num_sub_workers=0 does not mean "one sub-worker per chip process". The runtime computes num_sub = max(dc.num_sub_workers, len(sub_worker_fns)) (distributed_runner.py:1047, same at :1320) before passing it to Worker(...) (:328). So 0 means "no explicit request — use as many sub-workers as the program declares", which is 0 when none are declared.

3. Signal-reuse blocker points at the wrong tracker. The added paragraph attributes the missing self-resetting variant to "PTOAS issue #797". The constraint is tracked in this repo by #2156 ("Make composite collective barrier signals reusable"), which explicitly lists "any corresponding DSL/API documentation that describes signal lifetime or reuse" as in scope. A cross-repo issue number is also opaque to pypto readers. Worth noting: #2156's Affected areas lists allgather / reduce_scatter / all_to_all / all_to_all_v but not allreduce (which uses AtomicAdd(1)/WaitGe(1)), so please either extend that list or make the docstring say it is the same family of problem.

Question

4. The removed mesh partial-valid / symbolic-extent paragraph. That is the largest deletion in the diff, and the behaviour it documented is still live (lower_composite_ops_pass.cpp around :770 and :819). Those lines were user-visible constraints — which targets get rejected. Was this moved into one of the manual chapters? Since this PR lands independently and those chapters come later, there would be a window with no documentation for it. A condensed version, or a pointer, would avoid that.

Worth doing while you're in these docstrings

5. The two RST simple tables do not render. The docs site is mkdocs + mkdocstrings with docstring_style: google (mkdocs.yml:117-123); the body is rendered as Markdown, with no RST processing. Running the real pipeline (griffe 2.1.0 → Python-Markdown 3.10.3 with the extensions from mkdocs.yml) on the BenchmarkStats docstring gives:

<p>.. rubric:: Quick-reference</p>
<p>============================== =========================================================
Accessor                       Description
============================== ====================...
<code>stats.device_us_median</code>     Median device wall (µs).
...</p>

The whole table collapses into a single paragraph and the column structure is lost (browsers then collapse the newlines into spaces). Same for the spmd table. A Markdown table renders correctly today — the tables extension is already enabled. This is the one formatting item where information is actually lost; .. rubric::, .. seealso:: and the :func:/:class: roles also come out literal, but roles alone have ~780 pre-existing uses on main, so that is a repo-wide cleanup rather than something to pin on this PR.

6. DistributedConfig should use Attributes:, not Fields:. griffe parses an unknown Fields: block as a generic admonition (one box with all four fields inside as running text); with Attributes: it produces structured entries that attach to the dataclass members. Fields: has 0 uses in python/pypto, Attributes: has 14.

7. allgather prose still says input. Three occurrences ("Unified 3-arg form: pld.tensor.allgather(input, target, signal)", "input must be a DIFFERENT buffer", "On the HOST path input is…") while the parameter — and the Args: entry — is local_data. all_to_all is the one whose parameter is actually named input.

Separate change, not this PR

8. notify / wait annotate their signal operand as Tensor. Both deducers require a DistributedTensorType (system.cpp:80 and :114). Since DistributedTensor subclasses Tensor (distributed_tensor.py:59) the annotation is too wide rather than wrong, so a plain pl.Tensor passes static checking and only fails in C++. Narrowing it is a behaviour-visible cross-layer change (.claude/rules/cross-layer-sync.md), better as its own PR/issue than inside a docs-only one.

Minor

• The notify note attributes the target name to "legacy reasons (inherited from early signal protocol drafts)" — unverifiable history; the useful part is just that notify calls it target, wait calls it signal, and both mean the signal tensor.
• The PR description mentions adding missing benchmark() keyword-only parameters, but benchmark() itself is unchanged in the diff — only the BenchmarkStats quick-reference table was added.

[georgebisbas]

Thanks for the thorough review — addressed all of these (commits 0df025b2 and 7a7019a4).

Must fix

1. alloc_window_buffer HOST-only: went with the stronger option — added a real parser-level check in ast_parser.py mirroring world_size's existing enforcement exactly (rejects outside level=HOST and rejects nesting inside InCore/SPMD scopes), plus two new tests (test_alloc_window_buffer_rejected_outside_host_function, test_alloc_window_buffer_rejected_in_nested_device_scope_within_host_function). The docstring's "raises a parser error" claim is now actually true.
2. num_sub_workers=0: corrected to describe the real semantics — max(num_sub_workers, len(sub_worker_fns)), i.e. defers to however many sub-workers the program declares, not "one per chip process."
3. Signal-reuse blocker: replaced "PTOAS issue fix(torch_codegen): write/assemble return values, valid_shapes masking, name sanitization #797" with Make composite collective barrier signals reusable #2156, and noted allreduce isn't currently in that issue's affected-areas list (allgather/reduce_scatter/all_to_all/all_to_all_v are) but hits the identical constraint.

Question
4. Confirmed the removed mesh partial-valid/symbolic-extent content is still documented, untouched by this PR, in docs/en/dev/distributed_ops.md and docs/en/dev/passes/12-lower_composite_ops.md — no gap. Added a pointer from the docstring to those chapters.

Worth doing
5. Converted both RST simple tables (BenchmarkStats quick-reference, spmd usage-forms) to Markdown pipe tables. Verified with markdown.markdown(..., extensions=["tables"]) that they now render as real <table> elements. Also extended the BenchmarkStats table with per_dispatch/dispatch_tasks/unstable_dispatch_slots — new accessors #2223 added to main independently while this branch was open, caught during a rebase before implementing this fix so the table doesn't ship already-stale.
6. DistributedConfig: Fields: → Attributes:.
7. allgather: fixed all 4 stale input references to local_data.

Item 8 (notify/wait typing) — tried, reverted
Initially narrowed target/signal from Tensor to DistributedTensor, matching the C++ deducer. That broke pyright on two legitimate existing tests (test_tensor_view_preserves_loop_carried_distributed_metadata, ..._while_carried...): a DistributedTensor value loses its static type through pl.range/pl.while_ loop-carry, so the (correctly IR-level-distributed) view passed in there is only statically typed Tensor. Reverted the narrowing and filed #2234 tracking both the loop-carry typing gap and the narrowing itself — your original instinct to keep this as a separate PR was right.

Minor

• Dropped the unverifiable "legacy reasons" history note on notify's target param.
• Corrected the PR description (the benchmark() keyword-only-params claim wasn't actually in the diff).

Full tests/ut suite (8433 passed, 13 skipped) and pre-commit (ruff, pyright, etc.) both green after these changes.

[YunjiQin]

Re-checked the fix round — all points land, and thanks for chasing the notify/wait narrowing far enough to find the loop-carry typing gap (#2234) rather than just deferring it. On item 4 the correction is mine to make: I checked docs/en/dev/passes/12-lower_composite_ops.md:197,201 and docs/en/dev/distributed_ops.md:298-304 and the mesh partial-valid / symbolic-extent text is there verbatim, so the docstring copy really was redundant and there is no gap — my concern was unfounded.

One thing to fix in the new parser check:

raise ParserSyntaxError(
    "pld.system.alloc_window_buffer() can only be called in HOST orchestration "
    ...

pld.system.alloc_window_buffer does not exist — the op lives in the tensor namespace (pld.tensor.alloc_window_buffer, short form pld.alloc_window_buffer); grep -rn "system\.alloc_window_buffer" python/ src/ returns nothing. It looks like a leftover from adapting world_size's message, where pld.system. is correct. A user who hits this error and follows the name will look in the wrong namespace, and the rest of the message already reads correctly, so it is a one-word fix.

For what it's worth, I also scanned the enforcement for fallout: 138 pl.function bodies across tests/ and examples/ call alloc_window_buffer and every one of them is level=pl.Level.HOST. The single InCore hit (tests/st/distributed/collectives/test_l3_allreduce_ring.py:92) is a docstring mention, not a call — so the new check should not disturb any existing case, including the st suites that do not run in tests/ut.

[georgebisbas]

Fixed. Exactly right — that was a copy-paste leftover from the world_size check (where pld.system.world_size() is correct). Changed to pld.tensor.alloc_window_buffer() in 825d209.

Also rebased onto current main (672871b) — no conflicts, all pre-commit checks pass (ruff, pyright, headers, etc.). Ready for another look.