[RFC] L3 Distributed Programming Interface Design

[YunjiQin]

Summary

Propose the L3 (HOST-level) distributed programming interface for PyPTO. This RFC defines a three-level execution hierarchy (HOST → CHIP → CORE_GROUP) with two complementary writing styles: functional (separate @pl.function methods) and inline (pl.at() context manager). The goal is to collect feedback on API ergonomics, naming, and design trade-offs before the implementation lands.

Motivation

PyPTO currently supports single-device programming (InCore kernels + Chip Orchestration). To enable multi-device and HOST-level workflows — such as dispatching chip tasks, verifying results on CPU, or aggregating outputs from multiple devices — we need a distributed programming interface.

Key requirements:

• Express multi-level hierarchies (HOST → CHIP → CORE_GROUP) in the DSL
• Support SubWorkers: HOST-level pure-Python tasks that run after chip tasks complete (verification, reduction, post-processing)
• Provide both a modular (functional) and a compact (inline) writing style
• Integrate with the existing ir.compile() pipeline and simpler's Worker(level=3) runtime

Design

Hierarchy Model

Three execution levels, each with an optional role:

Level	Keyword	Role	Description
L0–L1	pl.Level.CORE_GROUP	—	Tile-level vector operations (InCore)
L2	pl.Level.CHIP	Orchestrator	Device-level task dispatch
L3	pl.Level.HOST	Orchestrator / Worker	HOST-level distributed scheduling

Roles (L3+ only):

• pl.Role.Orchestrator — Builds DAG, dispatches chip tasks and SubWorkers
• pl.Role.Worker — Executes pure-Python compute/validation (SubWorker)

---

Style 1: Functional (Separate Functions)

Each hierarchy level is a separate @pl.function-decorated method. Explicit, modular, reusable.

import pypto.language as pl
import torch
from pypto import ir
from pypto.ir.distributed_compiled_program import DistributedConfig

@pl.program
class L3DependencyProgram:
    """L3: HOST orch → CHIP worker (a + b) → SubWorker (verify)."""

    @pl.function(type=pl.FunctionType.InCore)
    def tile_add(
        self,
        a: pl.Tensor[[128, 128], pl.FP32],
        b: pl.Tensor[[128, 128], pl.FP32],
        f: pl.Out[pl.Tensor[[128, 128], pl.FP32]],
    ) -> pl.Tensor[[128, 128], pl.FP32]:
        tile_a = pl.load(a, [0, 0], [128, 128])
        tile_b = pl.load(b, [0, 0], [128, 128])
        tile_f = pl.add(tile_a, tile_b)
        out_f = pl.store(tile_f, [0, 0], f)
        return out_f

    @pl.function(type=pl.FunctionType.Orchestration)
    def chip_orch(
        self,
        a: pl.Tensor[[128, 128], pl.FP32],
        b: pl.Tensor[[128, 128], pl.FP32],
        f: pl.Out[pl.Tensor[[128, 128], pl.FP32]],
    ) -> pl.Tensor[[128, 128], pl.FP32]:
        out_f = self.tile_add(a, b, f)
        return out_f

    @pl.function(level=pl.Level.HOST, role=pl.Role.Worker)
    def verify(self, f: pl.Tensor[[128, 128], pl.FP32]):
        expected = torch.full((128, 128), 5.0, dtype=torch.float32)
        if not torch.allclose(f, expected, rtol=1e-5, atol=1e-5):
            raise AssertionError(
                f"SubWorker verify failed: expected 5.0, "
                f"got max={f.max().item()}, min={f.min().item()}"
            )

    @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator)
    def host_orch(
        self,
        a: pl.Tensor[[128, 128], pl.FP32],
        b: pl.Tensor[[128, 128], pl.FP32],
        f: pl.Out[pl.Tensor[[128, 128], pl.FP32]],
    ) -> pl.Tensor[[128, 128], pl.FP32]:
        out_f: pl.Tensor[[128, 128], pl.FP32] = self.chip_orch(a, b, f)
        self.verify(out_f)
        return out_f

Key characteristics:

• Each level is a separate method — clear, testable in isolation
• level= / role= on @pl.function specifies the hierarchy position
• SubWorker body (verify) is pure Python — can use torch, numpy, etc.
• SubWorker inputs/outputs defined via standard function signature with pl.Tensor type annotations
• HOST Orchestrator (host_orch) dispatches chip work via self.chip_orch(...) and SubWorker via self.verify(...)

---

Style 2: Inline (pl.at() Context Manager)

All hierarchy levels inlined into a single @pl.function via with pl.at(). Compact, self-contained.

@pl.program
class L3DependencyInlineProgram:
    """L3: all levels inlined into host_orch via pl.at() using tensor API."""

    @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator)
    def host_orch(
        self,
        a: pl.Tensor[[128, 128], pl.FP32],
        b: pl.Tensor[[128, 128], pl.FP32],
        f: pl.Out[pl.Tensor[[128, 128], pl.FP32]],
    ) -> pl.Tensor[[128, 128], pl.FP32]:
        # CHIP orchestration scope
        with pl.at(level=pl.Level.CHIP, role=pl.Role.Orchestrator):
            # CORE_GROUP scope (tile-level computation)
            with pl.at(level=pl.Level.CORE_GROUP):
                out = pl.add(a, b)
                out_f = pl.assemble(f, out, [0, 0])

        # SubWorker scope (verification on HOST)
        with pl.at(level=pl.Level.HOST, role=pl.Role.Worker):
            expected = torch.full((128, 128), 5.0, dtype=torch.float32)
            if not torch.allclose(out_f, expected, rtol=1e-5, atol=1e-5):
                raise AssertionError(
                    f"SubWorker verify failed: expected 5.0, "
                    f"got max={out_f.max().item()}, min={out_f.min().item()}"
                )

        return out_f

Key characteristics:

• Single method — all hierarchy levels visible at a glance
• Nesting reflects the execution hierarchy: HOST → CHIP → CORE_GROUP
• SubWorker body via with pl.at(level=pl.Level.HOST, role=pl.Role.Worker) — pure Python, source captured from AST
• SubWorker inputs: variables referenced from parent scope are auto-detected as inputs
• SubWorker outputs: declared via annotated assignments with pl.Tensor type (e.g., result: pl.Tensor[[128, 128], pl.FP32] = ...). This is how the parser identifies output variables from the SubWorker scope
• Variables flow across scope boundaries (e.g., out_f used in SubWorker block)

---

Comparison: Functional vs Inline

Aspect	Functional	Inline (pl.at())
Structure	Separate methods	Nested with blocks
Reusability	Methods reusable across programs	Scopes are local
Readability	Modular, each level self-contained	Compact, full picture in one function
SubWorker declaration	@pl.function(level=HOST, role=Worker)	with pl.at(level=HOST, role=Worker)
SubWorker inputs	Function parameters with pl.Tensor types	Auto-detected from parent scope variable references
SubWorker outputs	Function parameters with pl.Out[pl.Tensor[...]]	Annotated assignments: var: pl.Tensor[...] = expr
Best for	Complex programs with shared kernels	Simple pipelines, prototyping

Both styles produce identical compiled output and runtime behavior.

---

Advanced Pattern: Parallel Chip Tasks + SubWorker Reduce

Multiple independent chip tasks dispatched by the HOST orchestrator, with a SubWorker aggregating results.

@pl.program
class L3ParallelReduceProgram:
    """L3: HOST orch → 2 independent chip workers → SubWorker (reduce)."""

    @pl.function(type=pl.FunctionType.InCore)
    def tile_add(
        self,
        a: pl.Tensor[[128, 128], pl.FP32],
        b: pl.Tensor[[128, 128], pl.FP32],
        f: pl.Out[pl.Tensor[[128, 128], pl.FP32]],
    ) -> pl.Tensor[[128, 128], pl.FP32]:
        tile_a = pl.load(a, [0, 0], [128, 128])
        tile_b = pl.load(b, [0, 0], [128, 128])
        tile_f = pl.add(tile_a, tile_b)
        return pl.store(tile_f, [0, 0], f)

    @pl.function(type=pl.FunctionType.InCore)
    def tile_sub(
        self,
        a: pl.Tensor[[128, 128], pl.FP32],
        b: pl.Tensor[[128, 128], pl.FP32],
        f: pl.Out[pl.Tensor[[128, 128], pl.FP32]],
    ) -> pl.Tensor[[128, 128], pl.FP32]:
        tile_a = pl.load(a, [0, 0], [128, 128])
        tile_b = pl.load(b, [0, 0], [128, 128])
        tile_f = pl.sub(tile_a, tile_b)
        return pl.store(tile_f, [0, 0], f)

    @pl.function(type=pl.FunctionType.Orchestration)
    def chip_orch_add(self, a, b, f):
        return self.tile_add(a, b, f)

    @pl.function(type=pl.FunctionType.Orchestration)
    def chip_orch_sub(self, a, b, f):
        return self.tile_sub(a, b, f)

    @pl.function(level=pl.Level.HOST, role=pl.Role.Worker)
    def reduce_sum(
        self,
        sum_ab: pl.Tensor[[128, 128], pl.FP32],
        diff_ab: pl.Tensor[[128, 128], pl.FP32],
        f: pl.Out[pl.Tensor[[128, 128], pl.FP32]],
    ):
        result = sum_ab + diff_ab  # f = (a+b) + (a-b) = 2a
        f[:] = result

    @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator)
    def host_orch(self, a, b, f):
        sum_ab = pl.create_tensor([128, 128], dtype=pl.FP32)
        diff_ab = pl.create_tensor([128, 128], dtype=pl.FP32)
        out_sum = self.chip_orch_add(a, b, sum_ab)
        out_diff = self.chip_orch_sub(a, b, diff_ab)
        self.reduce_sum(out_sum, out_diff, f)
        return f

Key points:

• chip_orch_add and chip_orch_sub are independent → runtime can parallelize them
• reduce_sum (SubWorker) runs after both chip tasks complete
• pl.create_tensor() allocates intermediate HOST tensors
• SubWorker receives multiple chip outputs and writes the final result

---

API Surface

@pl.function Decorator Parameters

@pl.function(
    type: pl.FunctionType = ...,      # InCore, Orchestration
    level: pl.Level | None = None,    # HOST, CHIP, CORE_GROUP, ...
    role: pl.Role | None = None,      # Orchestrator, Worker
)

For L3 distributed, level= and role= replace type=:

• level=pl.Level.HOST, role=pl.Role.Orchestrator → HOST entry point
• level=pl.Level.HOST, role=pl.Role.Worker → SubWorker

pl.at() Context Manager

pl.at(
    level: pl.Level,                          # Target hierarchy level
    role: pl.Role | None = None,              # Function role (L3+ only)
    *,
    optimizations: list[Optimization] | None = None,
    name_hint: str = "",                      # Outlined function name
) -> AtContext

DistributedConfig

from pypto.ir.distributed_compiled_program import DistributedConfig

DistributedConfig(
    device_ids: list[int] = [0],        # Device IDs to use
    num_sub_workers: int = 0,           # SubWorker thread count
    runtime: str = "tensormap_and_ringbuffer",
    block_dim: int = 1,
    aicpu_thread_num: int = 1,
)

Compilation and Execution

compiled = ir.compile(
    L3DependencyProgram,
    platform="a2a3sim",
    distributed_config=DistributedConfig(
        device_ids=[7],
        num_sub_workers=1,
        block_dim=3,
        aicpu_thread_num=4,
    ),
)

a = torch.full((128, 128), 2.0, dtype=torch.float32)
b = torch.full((128, 128), 3.0, dtype=torch.float32)
f = torch.zeros((128, 128), dtype=torch.float32)

compiled(a, b, f)          # In-place style
# or
f = compiled(a, b)         # Return style (auto-allocates outputs)

---

SubWorker Mechanism

SubWorkers (level >= HOST, role = Worker) are pure-Python functions executed on the HOST after chip tasks complete.

Source capture:

• Functional style: inspect.getsource() at @pl.program creation time
• Inline style: AST-level source extraction from with pl.at() body

Input/output detection in inline style:

• Inputs: AST-analyzed references to parent-scope variables
• Outputs: Annotated assignments with pl.Tensor type (e.g., result: pl.Tensor[[128, 128], pl.FP32] = expr)

Generated output structure:

output_dir/
├── orchestration/
│   └── host_orch.py           # Generated HOST orchestrator
├── next_levels/
│   └── chip_orch/             # Compiled chip-level task
│       ├── kernels/
│       └── kernel_config.py
└── sub_workers/
    └── verify.py              # Generated SubWorker callable

Runtime execution:

1. Compile chip-level tasks → simpler.ChipCallable
2. Load generated SubWorker Python modules
3. Create simpler.Worker(level=3), register SubWorkers
4. Execute HOST orchestration closure via worker.run(Task(...))

Affected Files

Key implementation files:

File	Purpose
python/pypto/language/dsl_api.py	pl.at() context manager
python/pypto/language/parser/decorator.py	@pl.function, @pl.program, SubWorker registry
python/pypto/language/parser/ast_parser.py	AST parsing, _parse_sub_worker_scope()
python/pypto/ir/compile.py	Distributed program detection, compilation entry
python/pypto/ir/distributed_compiled_program.py	DistributedConfig, DistributedCompiledProgram
python/pypto/backend/pto_backend.py	Distributed code generation
python/pypto/runtime/distributed_runner.py	Distributed execution via simpler Worker

Open Questions

1. Naming: Is pl.Role.Worker clear enough for SubWorker semantics? Should it be pl.Role.SubWorker to avoid confusion with the runtime Worker class?

2. SubWorker constraints: Should SubWorkers be restricted to pure-Python (no DSL ops), or should they support a subset of DSL operations for HOST-level tensor manipulation?

3. Variable flow in inline style: Variables cross with pl.at() boundaries (e.g., out_f computed inside CHIP scope, used in SubWorker scope). Is implicit data flow clear enough, or should we require explicit declarations?

4. Output annotation convention: In the inline style, SubWorker outputs use annotated assignments (var: pl.Tensor[...] = expr). Is this convention intuitive, or should there be a more explicit pl.output(...) marker?

5. Error handling: How should errors in SubWorkers be surfaced? Currently they raise Python exceptions — should there be a structured error reporting mechanism?

6. Multi-device: The current design uses a single device_ids list. How should we express per-chip-task device affinity for multi-device workflows?

7. Inline style limitations: Should the inline style support the advanced pattern (multiple chip tasks + SubWorker reduce), or is the functional style sufficient for complex workflows?

[lyfne123]

1. rename pl.Role.Worker to pl.Role.SubWorker

[YunjiQin]

Open Question: Explicit torch.Tensor ↔ pl.Tensor Conversion Interface

Currently SubWorker bodies operate directly on pl.Tensor-annotated parameters, but the actual runtime values are torch.Tensor (or similar host-side tensor objects). This raises a design question:

Should the front-end provide an explicit conversion interface between torch.Tensor and pl.Tensor, so that SubWorker bodies don't operate on pl.Tensor directly?

Sub-questions

1. Frontend semantics: Should SubWorker code be written entirely in terms of torch.Tensor (host-native types), with pl.Tensor only appearing in the function signature as a type contract? This would make SubWorker bodies "look like" plain Python/PyTorch code without any DSL-typed values flowing through.

2. Conversion location: If explicit conversion is required, where should it happen?

   • At the Orchestrator — the HOST orchestrator converts pl.Tensor to torch.Tensor before dispatching to the SubWorker. SubWorker signature uses torch.Tensor types directly.
   • Inside the SubWorker — the SubWorker receives pl.Tensor and explicitly converts to torch.Tensor (e.g., t = pl.to_torch(x)) before computation, and converts back (pl.from_torch(...)) when writing outputs.
   • Implicit (current behavior) — the framework auto-unpacks ContinuousTensor arguments to host-native tensors transparently; SubWorker code uses pl.Tensor-annotated values as if they were torch.Tensor.

Trade-offs

Approach	Pro	Con
Implicit (current)	Less boilerplate; SubWorker bodies are concise	Type annotation lies — declared pl.Tensor actually behaves like torch.Tensor; confusing for users reading the code
Convert at Orchestrator	SubWorker is "pure PyTorch"; clean separation of concerns	Orchestrator gains conversion responsibility; signature/dispatch pairing becomes more involved
Convert in SubWorker	Explicit and auditable; no hidden magic	More boilerplate per SubWorker; users must remember to convert

Why this matters

• Readability: A SubWorker annotated with pl.Tensor but using torch.allclose(...) is misleading — the value is not actually a pl.Tensor.
• Type checkers: Static analysis tools can't reason about implicit conversions.
• Generated code clarity: Generated sub_workers/*.py files currently insert _tensor_from_continuous(...) wrappers — surfacing this in the user-facing API would make the conversion model first-class.
• Future extensibility: If SubWorkers later need to support NumPy, JAX, or other host-side tensor libraries, an explicit conversion interface scales better than implicit auto-unpacking.

Need to decide: do we keep the implicit model and document the "pl.Tensor in signature is really a torch.Tensor at runtime" convention, or do we introduce explicit conversion APIs (and if so, where)?

[YunjiQin]

Communication Interface Design (rev. 2026-05-12)

Extends this RFC with the inter-rank communication interface needed to express simpler/examples/workers/l3/allreduce_distributed/ (and ffn_tp_parallel) in PyPTO. The original RFC covers HOST orchestration and SubWorker dispatch; this comment proposes how to layer HCCL window buffers, multi-comm-group support, and signal-matrix synchronization on top.

Supersedes rev. 2026-05-08. The central change is alloc / view decoupling: pld.alloc_window_buffer is now a pure byte-sized allocation (mirroring tile.alloc / MemRef / TileType.memref_), and the DistributedTensor view is materialized at the use site by a separate pld.window op.

Principle: user does not construct CommGroup; alloc / view strictly mirror the MemRef triplet

Per-rank window buffers live in independent address spaces on different chips, with no unified offset; rank is also a runtime value, so codegen can't implicitly distinguish local vs remote in plain pl.load/store. We introduce four key abstractions, structured exactly like the tile.alloc / MemRef / TileType.memref_ triplet:

• pld.alloc_window_buffer(size) — host-orch op for pure address-space allocation; size is in bytes. Returns a singleton PtrType — the same allocation-identity token as tile.alloc. The parser binds the LHS as a plain Var(PtrType) (NOT a WindowBuffer subclass). The buffer's name is not user-supplied — the parser extracts it from the LHS variable name (carried via Var.name_hint) and enforces program-wide uniqueness.
• pld.window(buf, [shape], *, dtype) — at the use site, materializes the Ptr handle into a DistributedTensor view. Shape and dtype first enter the type system here (mirroring how TileType[[shape], dtype, memref] combines MemRef at use sites). Result type is DistributedTensorType; its window_buffer_ back-reference is nullopt at parse time.
• pld.DistributedTensor — a pl.Tensor subclass (IR: DistributedTensorType : TensorType). Used as formal-param type in chip_orch / InCore. Verifier requires it for pld.tile.remote_load / pld.system.notify / pld.system.wait. Param annotations have window_buffer == nullopt.
• device=r dispatch kwarg on self.chip_orch(..., device=r) — r must be ConstInt or the induction var of an enclosing pl.range(...) loop. This lets the CollectCommGroups pass derive the device set per buffer at compile time.

The CollectCommGroups pass constructs a WindowBuffer (Var subclass mirroring MemRef) for each alloc and threads the same shared_ptr<const WindowBuffer> into both Program.comm_groups_[i].slots_ and every consuming pld.window result type's window_buffer_ — exactly how InitMemRef constructs a MemRef for tile.alloc and shares one pointer across many TileType.memref_ consumers.

Top-level look

import pypto.language as pl
import pypto.language.distributed as pld

@pl.program
class AllReduceProgram:
    # No class-attribute CommGroup. Comm groups inferred by CollectCommGroups pass.

    @pl.function(type=pl.FunctionType.InCore)
    def allreduce_tile(
        self,
        inp:    pl.Tensor[[ALLREDUCE_COUNT], pl.FP32],
        out:    pl.Out[pl.Tensor[[ALLREDUCE_COUNT], pl.FP32]],
        data:   pld.DistributedTensor[[ALLREDUCE_COUNT], pl.FP32],
        signal: pld.DistributedTensor[[pld.world_size], pl.INT32],
    ) -> pl.Tensor[[ALLREDUCE_COUNT], pl.FP32]:
        # CommCtx derived per-tensor (multi-group safe)
        my_rank = data.comm.rank
        nranks  = data.comm.nranks
        ...
        pld.system.notify(target=signal, peer=peer, offsets=[my_rank], value=1, op=pld.NotifyOp.AtomicAdd)
        pld.system.wait(signal=signal, offsets=[peer], expected=1, cmp=pld.WaitCmp.GE)
        recv = pld.tile.remote_load(data, peer=peer, offsets=[0], shape=[ALLREDUCE_COUNT])

    @pl.function(type=pl.FunctionType.Orchestration)
    def chip_orch(
        self,
        inp, out,
        data:   pld.DistributedTensor[[ALLREDUCE_COUNT], pl.FP32],
        signal: pld.DistributedTensor[[pld.world_size], pl.INT32],
    ): return self.allreduce_tile(inp, out, data, signal)

    @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator)
    def host_orch(self, inputs, outputs):
        # alloc lives OUTSIDE the dispatch loop. size is BYTES — exactly like
        # tile.alloc(memspace, size_in_bytes). alloc doesn't know shape/dtype;
        # it's a pure address-space allocation. LHS name ("data_buf" / "signal_buf")
        # becomes the runtime-unique buffer identifier (program-wide), carried via
        # Var.name_hint — no separate field.
        data_buf   = pld.alloc_window_buffer(ALLREDUCE_COUNT * 4)   # 256 FP32 = 1024 bytes
        signal_buf = pld.alloc_window_buffer(pld.world_size * 4)    # world_size INT32

        for r in pl.range(pld.world_size):
            # pld.window materializes Ptr handle -> DistributedTensor view; shape
            # and dtype enter the type system here. DistributedTensorType.window_buffer
            # is None at parse, filled in by CollectCommGroups pass.
            data   = pld.window(data_buf,   [ALLREDUCE_COUNT], dtype=pl.FP32)
            signal = pld.window(signal_buf, [pld.world_size], dtype=pl.INT32)
            self.chip_orch(inputs[r], outputs[r], data, signal, device=r)
        return outputs

Multi-comm-group support: dist_t.comm

A chip may participate in multiple comm groups (different rank in each). A single ctx: pld.CommCtx formal param can't express this, so we drop the formal ctx param entirely. Instead, every DistributedTensor exposes .comm.rank / .comm.nranks that resolve to its own group's CommContext.

Lowering chain (extended for the alloc / view split):

%buf:  Var(PtrType)             = pld.alloc_window_buffer(size)         (plain Var, type = singleton PtrType)
%data: DistributedTensorType    = pld.window(%buf, [shape], dtype=...)  (window_buffer = None at parse)
       --(def-use, dispatch device=r)-->
                                      [after CollectCommGroups pass]
                                                       v
WindowBuffer (Var; base=%buf, size, host flags)  -- shared shared_ptr -->
   ├─ Program.comm_groups_[i].slots_[k]
   └─ each consuming pld.window's DistributedTensorType.window_buffer_
                                                       v
                                                CommContext (per-device)

data.comm and signal.comm resolve to the same CommContext if they share a device set; tensors from different groups stay independent. Resolution: dist_t.type.window_buffer → WindowBuffer → program.comm_groups_ → per-device CommContext.

Parser rules

• pld.alloc_window_buffer(size) must be the RHS of a single-target assignment (buf = pld.alloc_window_buffer(size)). Multi-target / tuple-unpack / subscript / attribute targets are rejected. Parser auto-injects the LHS name as the op's name kwarg — user-supplied kwargs are rejected. The size arg must be a scalar byte count (int / ir.Expr); list / tuple is rejected. LHS is bound as plain Var(PtrType) (NOT a WindowBuffer subclass) — exact mirror of mem_vec_7: Ptr = tile.alloc(...).
• pld.alloc_window_buffer must appear inside a HOST Orchestrator function body (nested for / if / with is allowed). SubWorker bodies reject it. LHS names are program-wide unique.
• pld.window(buf, [shape], *, dtype): buf must evaluate to a Var with .type == PtrType; shape must be a list / tuple literal; only dtype kwarg is accepted. Result type DistributedTensorType(shape, dtype) with window_buffer_ = nullopt. Not restricted to host_orch loops — may materialize anywhere, including inside dispatch arg expressions.
• device=r on a dispatch must be ConstInt or the induction var of an enclosing pl.range(...) loop; otherwise ParserSyntaxError. Only permitted on Orchestration callees.

CollectCommGroups pass

Input: parsed program where alloc LHSs are plain Var(PtrType), every pld.window result type has window_buffer_ = nullopt, and program.comm_groups_ = []. No WindowBuffer Vars exist yet.

For each pld.alloc_window_buffer, walk def-use: alloc LHS → all pld.window(ptr, ...) consumers → all self.chip_orch(..., dist_t, ..., device=r) dispatches. Derive a device-set descriptor from each r:

• r = ConstInt c → subset = {c}
• r = loop var of pl.range(pld.world_size) → ALL_DEVICES (sentinel; resolved at runtime to DistributedConfig.device_ids)
• r = loop var of pl.range(ConstInt N) → subset = {0..N-1}

Then:

1. For each alloc, construct a new WindowBuffer Var (base = alloc LHS, size = alloc's byte-count arg, load_from_host / store_to_host flags; no name / dtype fields — name flows through inherited Var.name_hint).
2. In-place rewrite each consuming pld.window Call op so its result type goes from DistributedTensorType(shape, dtype, nullopt) to DistributedTensorType(shape, dtype, wb) — the same shared_ptr<const WindowBuffer> is shared across all consumers of one alloc.
3. Cluster WindowBuffers by device descriptor: same descriptor → one CommGroup (one HCCL window with multiple ChipBufferSpec slots); different descriptors → independent groups. Each CommGroup.slots_[k] shares its shared_ptr with the corresponding pld.window result type's window_buffer_.
4. Attach { group_idx, slot_idx } attrs to each alloc Call op (lowering index for dist_t.comm). Dispatch calls do NOT get any extra attr — each DistributedTensor arg resolves its own group via dist_t.type.window_buffer.

This mirrors InitMemRef exactly: pass-constructed MemRef shared across TileType.memref_ slots ↔ pass-constructed WindowBuffer shared across DistributedTensorType.window_buffer_ slots and CommGroup.slots_.

Important: pld.world_size stays symbolic in IR throughout. It denotes the runtime count of cards (len(DistributedConfig.device_ids)); the pass does NOT instantiate it. Driver substitutes it when materializing chip_bootstrap_configs.

WindowBuffer maps 1:1 to simpler.task_interface.ChipBufferSpec (name = wb.name_hint_ / nbytes = int(wb.size) / dtype recovered by the driver from the consuming pld.window result type).

IR transformation (before / after)

Before pass:

# program.comm_groups_ = []  (empty)

@host_orch [level=HOST, role=Orchestrator]
    (%inputs:  Tensor<[world_size, 256] f32>,
     %outputs: Tensor<[world_size, 256] f32, dir=Out>) -> ...
{
    # alloc is a pure allocation; LHS is plain Var(PtrType), name carried via Var.name_hint.
    # Exact mirror of `mem_vec_7: Ptr = tile.alloc(memspace, size)`.
    %data_buf:   Var<Ptr, name_hint="data_buf">
        = pld.alloc_window_buffer(1024, name="data_buf")
    %signal_buf: Var<Ptr, name_hint="signal_buf">
        = pld.alloc_window_buffer(world_size * 4, name="signal_buf")

    pl.for %r in pl.range(0, world_size) {
        # pld.window materializes into DistributedTensorType; shape/dtype enter the type system here.
        # window_buffer = None — filled back by the pass.
        %data:   DistributedTensor<[256] f32, window_buffer=None>
            = pld.window(%data_buf,   [256],        dtype=f32)
        %signal: DistributedTensor<[world_size] i32, window_buffer=None>
            = pld.window(%signal_buf, [world_size], dtype=i32)

        pl.call @chip_orch (..., %data, %signal) {device = %r}
    }
}

After pass:

# program.comm_groups_ = [
#     CommGroup#0 {
#         devices_ = [],                              # empty list = ALL_DEVICES; runtime = DistributedConfig.device_ids
#         slots_   = [
#             wb_data    = WindowBuffer(base=%data_buf,   size=1024,         load=false, store=false),
#             wb_signal  = WindowBuffer(base=%signal_buf, size=world_size*4, load=false, store=false),
#         ],
#     },
# ]
# wb_data / wb_signal are newly constructed by the pass; name_hint inherited from base.

@host_orch [level=HOST, role=Orchestrator]
    (%inputs:  Tensor<[world_size, 256] f32>,         # unchanged
     %outputs: Tensor<[world_size, 256] f32, dir=Out>) -> ...
{
    # alloc Call op gains group_idx / slot_idx attrs; LHS is still a plain Var(PtrType).
    %data_buf:   Var<Ptr, name_hint="data_buf">
        = pld.alloc_window_buffer(1024, name="data_buf")
            { group_idx = 0, slot_idx = 0 }
    %signal_buf: Var<Ptr, name_hint="signal_buf">
        = pld.alloc_window_buffer(world_size * 4, name="signal_buf")
            { group_idx = 0, slot_idx = 1 }

    pl.for %r in pl.range(0, world_size) {            # still symbolic
        # pld.window result type's window_buffer_ is rewritten in-place:
        # None -> shared_ptr pointing to program.comm_groups_[0].slots_[0/1].
        %data:   DistributedTensor<[256] f32, window_buffer=&wb_data>
            = pld.window(%data_buf,   [256],        dtype=f32)
        %signal: DistributedTensor<[world_size] i32, window_buffer=&wb_signal>
            = pld.window(%signal_buf, [world_size], dtype=i32)

        # dispatch call carries NO group attr — each DistributedTensor arg's
        # type.window_buffer resolves directly to the WindowBuffer in program.comm_groups_.
        pl.call @chip_orch (..., %data, %signal) {device = %r}
    }
}

Pass changes summary:

1. program.comm_groups_ populated with [CommGroup#0]. devices_ = [] is a descriptor — driver instantiates it.
2. WindowBuffer instances are constructed by the pass (they do NOT exist at parse), with base pointing back to the alloc's plain Var(PtrType) LHS.
3. pld.window result-type window_buffer_ is rewritten in-place: nullopt at parse → shared_ptr to the WindowBuffer in program.comm_groups_[i].slots_[k] after the pass. Same alloc materialized by N pld.windows → all N result types' window_buffer_ point to the same WindowBuffer. Structurally isomorphic to InitMemRef populating TileType.memref_.
4. Each pld.alloc_window_buffer Call op gets a { group_idx, slot_idx } attr (lowering index for dist_t.comm). Dispatch calls do NOT get a group attr — each DistributedTensor arg resolves its own group via dist_t.type.window_buffer.
5. world_size stays symbolic. Runtime value = len(DistributedConfig.device_ids).
6. Function bodies (including allreduce_tile) are textually unchanged.

New ops (categorized to align with existing namespaces)

Op	Namespace	Why
pld.alloc_window_buffer(size, *, name)	host-orch op	Pure address-space allocation; size in bytes; returns singleton PtrType (same marker as tile.alloc). Parser binds LHS as plain Var(PtrType).
pld.window(buf, [shape], *, dtype)	host-orch op	Materializes Ptr → DistributedTensor view; result DistributedTensorType.window_buffer_ filled by pass.
pld.tile.remote_load(target, peer, offsets, shape)	pld.tile.* (data movement)	Mirrors single-device pl.load / pl.store
pld.system.notify(target, peer, offsets, value, op)	pld.system.* (synchronization)	Mirrors pl.system.tpush/tpop/tfree
pld.system.wait(signal, offsets, expected, cmp)	pld.system.* (synchronization)	Same

Verifier: target/signal must be DistributedTensorType. Plain pl.load/store does not check — operates on the local rank's slice for any tensor type.

Driver / runtime translation (distributed_runner.py)

1. Read program.comm_groups_. For each CommGroup:
   • Resolve devices_: empty list → DistributedConfig.device_ids; subset → use as-is.
   • Resolve symbolic shapes / sizes (substitute world_size with len(DistributedConfig.device_ids)).
   • Build one ChipBootstrapConfig per device with buffers = [ChipBufferSpec(name=wb.name_hint, dtype=<from consuming view>, count=..., nbytes=int(wb.size), load_from_host=wb.load_from_host, store_to_host=wb.store_to_host) for wb in slots_]. window_size = sum(spec.nbytes).
2. Pass cfgs to Worker(level=3, chip_bootstrap_configs=cfgs).
3. Inject worker.chip_contexts into the generated entry() closure as contexts=.
4. Manage rootinfo_path lifecycle.
5. load_from_host host tensors (first dim = world_size) sliced per-rank into HostBufferStaging, attached to cfg.host_inputs. store_to_host analogous.

Note: comm_window_size is no longer user-supplied on DistributedConfig — the driver computes it from sum(int(wb.size) for wb in CommGroup.slots_) after the pass.

Generated host_orch (codegen target)

def entry(orch, _args, cfg, *, tensors, callables, sub_ids, _keep, contexts):
    inputs = tensors["inputs"]
    outputs = tensors["outputs"]
    chip_orch_callable = callables["chip_orch"]

    for r in range(len(contexts)):                # pld.world_size -> len(contexts)
        ctx = contexts[r]
        chip_args = TaskArgs()
        chip_args.add_tensor(make_tensor_arg(inputs[r]),  TensorArgType.INPUT)
        chip_args.add_tensor(make_tensor_arg(outputs[r]), TensorArgType.OUTPUT_EXISTING)
        # Each DistributedTensor formal -> resolved via dist_t.type.window_buffer -> WindowBuffer -> group/slot.
        # Buffer name read from dist_t.type.window_buffer.name_hint (pass has filled window_buffer).
        chip_args.add_tensor(
            ContinuousTensor.make(
                data=ctx.buffer_ptrs["data_buf"],
                shapes=(ALLREDUCE_COUNT,), dtype=DataType.FLOAT32, child_memory=True,
            ), TensorArgType.INOUT,
        )
        chip_args.add_tensor(
            ContinuousTensor.make(
                data=ctx.buffer_ptrs["signal_buf"],
                shapes=(len(contexts),), dtype=DataType.INT32, child_memory=True,
            ), TensorArgType.INOUT,
        )
        chip_args.add_scalar(ctx.device_ctx)       # CommContext, for dist_t.comm.* lowering
        orch.submit_next_level(chip_orch_callable, chip_args, cfg, worker=r)

Why this design

• alloc / view decoupling isomorphic to MemRef: identical mental model for cross-rank windows vs single-device tile allocations — the same compiler infrastructure (InitMemRef-style pass populating a back-reference, shared_ptr sharing across consumers) extends straight to distributed.
• No user-facing CommGroup: comm groups are derived, not declared. Users only write alloc_window_buffer + pld.window + device=r dispatch — natural and minimal.
• Multi-group safe: dist_t.comm per-tensor; no single ambient context, so a chip participating in multiple groups is naturally expressible.
• Allocation flags don't pollute kernels: load_from_host / store_to_host live on WindowBuffer only.
• Symbolic world_size survives the IR: pass-time scope is "device set descriptor + group attribution"; runtime substitution at chip_bootstrap_configs materialization.
• comm_window_size no longer user-supplied: driver computes it from sum(int(wb.size) for wb in CommGroup.slots_) — one less thing for users to keep in sync.

Out-of-scope (deferred)

• High-level collective ops (pld.tile.allreduce, etc.) — get the three primitives working first.
• pld.for_each_rank — host_orch uses plain for r in pl.range(pld.world_size).
• Host-staging modes other than load_from_host / store_to_host.
• Final user-facing surface for load_from_host / store_to_host (default False for now).

[YunjiQin]

Implementation Plan & Status (rev. 2026-05-12)

Implementation plan for the Communication Interface Design. Each Nx is a separate PR; everything is described as an increment over origin/main.

Supersedes rev. 2026-05-11. The major shift: N2 landed with alloc / view decoupling that strictly mirrors tile.alloc / MemRef / TileType.memref_ — pld.alloc_window_buffer(size_in_bytes) is now a pure address-space allocation returning a singleton PtrType (shared with tile.alloc); the parser binds the LHS as a plain Var(PtrType) (NOT a WindowBuffer). WindowBuffer is now a Var subclass (mirroring MemRef) constructed by the CollectCommGroups pass (N4) — not at parse. DistributedTensorType gained a back-reference window_buffer_: optional<WindowBufferPtr> (mirroring TensorType.memref_), nullopt at parse and filled by N4. N3 also merged (pld.world_size op + device= kwarg).

Stage table

Stage	Increment	Depends on	Status
N1	New IR types (DistributedTensorType subclass + WindowBuffer / CommGroup IRNodes + Program.comm_groups_ field) + DSL pld.DistributedTensor + AOT comm-manifest sidecar (comm_manifest.json v2 + runtime bridge)	—	✅ Merged in #1297
N2	pld.alloc_window_buffer op (returns singleton PtrType, same marker as tile.alloc) + pld.window op (materializes DistributedTensorType(shape, dtype, window_buffer=None)) + DSL + parser; IR type layer: WindowBufferType collapsed to singleton marker, WindowBuffer upgraded from IRNode to Var subclass with base / size / load / store fields (no name / dtype fields), DistributedTensorType.window_buffer_ back-ref field added	N1	✅ Merged in #1346
N3	pld.world_size() host-only op + dispatch device=r kwarg + parser/IR plumbing	N2	✅ Merged in #1359
N4	CollectCommGroups pass: constructs WindowBuffer Var instances (parse has none), in-place rewrite of each pld.window result type's window_buffer_ to share the same shared_ptr<const WindowBuffer> with program.comm_groups_[i].slots_ (mirrors InitMemRef populating TileType.memref_); plus visitor-framework wire-up (VisitExpr_(WindowBufferPtr) on functor / visitor / mutator and 11 downstream functor subclasses)	N2, N3	✅ Merged in #1369
N5	CommCtxType + pld.get_comm_ctx(dist_t) op + dist_t.comm.rank/.nranks desugar + verifier	N1, N4	Planned
N6	pld.tile.remote_load + pld.system.notify + pld.system.wait ops + DistributedTensorType verifier + NotifyOp / WaitCmp enums; P0: comm_layout.h compile-time CommContext layout self-check	N1	Planned
N7	Codegen adaptation (add_tensor(ContinuousTensor.make(child_memory=True))) + driver host-staging materialization + final user-facing API for load_from_host / store_to_host	N4, N5, N6	Planned
N8	End-to-end: AllReduceProgram system test + doc migration	N7	Planned

Key design decisions (carried into the implementation)

The cross-cutting decision is strict isomorphism with the tile.alloc / MemRef / TileType.memref_ triplet:

MemRef side	Window side
tile.alloc(memspace, size_in_bytes) returns singleton PtrType	pld.alloc_window_buffer(size_in_bytes) returns the same singleton PtrType (shared marker)
mem_vec_7: Var(PtrType) = tile.alloc(...) — LHS is a plain Var via _assign_or_let	buf: Var(PtrType) = pld.alloc_window_buffer(...) — LHS is a plain Var via the same _assign_or_let (NOT a WindowBuffer)
MemRef extends Var with base / byte_offset / size; type = singleton MemRefType	WindowBuffer extends Var with base / size / load_from_host / store_to_host (no name / dtype fields); type = singleton WindowBufferType
InitMemRef pass constructs MemRef, fills back TileType.memref_	CollectCommGroups pass (N4) constructs WindowBuffer, fills back DistributedTensorType.window_buffer_ + Program.comm_groups_.slots_
Same base Ptr referenced by many views → shared shared_ptr<const MemRef>	Same base Ptr materialized by many pld.window → shared shared_ptr<const WindowBuffer> (slots_ and every consuming view's window_buffer_ share one instance)
TensorType.memref_ records the allocation source	DistributedTensorType.window_buffer_ records the source WindowBuffer — same shape / dtype but different source ⇒ structurally NOT equal

Item-by-item:

1. DistributedTensorType : TensorType subclass — exact ObjectKind dispatch, verifiers use As<DistributedTensorType> to reject plain Tensor. DSL is fully symmetric with pl.Tensor: pld.DistributedTensor[[shape], dtype] / [..., layout] / [..., memref] are all accepted; only the IR ObjectKind differs. Extra field window_buffer_: optional<WindowBufferPtr> — nullopt at parse (including pld.window result), pointing to the source WindowBuffer after N4; user parameter annotations always carry nullopt.
2. WindowBufferType is a singleton marker — no fields, isomorphic to MemRefType. Never appears on any SSA edge — alloc's SSA-edge type is PtrType; WindowBufferType is only the GetType() of the WindowBuffer Var subclass itself.
3. WindowBuffer extends Var (N1 placeholder IRNode → N2 Var subclass → N4 actual instantiation):
   • base_: VarPtr — points to the alloc op LHS's plain Var(PtrType) (exact mirror of MemRef.base_).
   • size_: ExprPtr — byte count (ConstInt or symbolic Expr, may contain pld.world_size).
   • load_from_host_ / store_to_host_: bool — host-staging flags (default false, wired up in N4+).
   • No name_ field — name flows through inherited Var::name_hint_ (IgnoreField); structural eq / hash depend on base + size + host flags only.
   • No dtype_ field — dtype lives on the view-side DistributedTensorType; the driver recovers it from the consuming view.
4. pld.alloc_window_buffer(size, *, name) op — pure address-space allocation:
   • Positional size is a scalar ExprPtr (bytes; NOT a shape list).
   • name kwarg is auto-injected by the parser from the LHS name_hint; users cannot supply kwargs (including dtype).
   • Return type: singleton PtrType (shared with tile.alloc); parser routes through standard _assign_or_let, binding the LHS as a plain Var(PtrType). No specialized Var subclass at parse.
   • load_from_host / store_to_host and other host-staging inputs are not wired up yet; deferred to N4+ with CollectCommGroups.
5. pld.window(buf, [shape], *, dtype) op — view materialization:
   • Positional buf must have type PtrType (i.e. output of pld.alloc_window_buffer); positional shape is a MakeTuple literal.
   • kwarg: dtype: DataType.
   • Returns DistributedTensorType(shape, dtype, /*window_buffer=*/nullopt); the deducer is self-contained — no parser cooperation or def-use trace needed. window_buffer_ is filled by N4.
6. CommGroup.devices_: vector<int64_t> (sorted device list; empty list = "covers all DistributedConfig.device_ids"). CommGroup.slots_: vector<WindowBufferPtr> — N4 pass constructs new WindowBuffer instances and the same shared_ptr is in slots_ and every consuming DistributedTensorType.window_buffer_.
7. pld.world_size is implemented as a host-only IR op — not a sentinel/DynVar (N3).
8. init_tensor / dispatch args do not enforce "must be a host_orch param" — SSA use-def already guarantees referential validity; neither op verifier nor parser repeats this check.
9. N1 was the "land on main" boundary: a self-contained IR/runtime infrastructure PR with no new user-facing DSL entry points (those follow in N2+); existing comm-less tests (test_l3_distributed.py, test_l3_parallel_reduce.py) did not regress.
10. N2 is the alloc / view decoupling milestone — landed in feat(distributed): pld.alloc_window_buffer / pld.window ops; remodel WindowBuffer as a Var #1346; the redesign below describes the post-merge state.

N1 — Merged in #1297

PR #1297 shipped the IR/runtime infrastructure baseline:

• DistributedTensorType IR subclass + pld.DistributedTensor DSL with full 6-arg ctor parity vs. TensorType (shape / dtype / layout / memref / view combinations), ObjectKind + KindTrait registration, and parser support in type_resolver.py.
• WindowBuffer (initially as IRNode placeholder, upgraded to Var subclass in N2) / CommGroup IRNodes + Program.comm_groups_ field — schema target for the future CollectCommGroups pass. Reflection visitors learned the new std::vector<int64_t> leaf field; deserializers cover WindowBuffer / CommGroup / Program(.., comm_groups, ..) overload.
• AOT comm-manifest sidecar (comm_manifest.json v2): lift_comm_manifest + emit_comm_manifest produce output_dir/orchestration/comm_manifest.json for comm-bearing programs (comm-less programs skip the file). Runtime bridge in distributed_runner.py reads the manifest, builds chip_bootstrap_configs, passes them to Worker(level=3, ...), and probes entry_fn's contexts= parameter so existing comm-less tests stay on the original code path.
• Test coverage: test_distributed_tensor_type.py, test_comm_group_schema.py, test_distributed_tensor_annotation.py, test_chip_bootstrap_configs.py. Existing tests/st/distributed/test_l3_distributed.py and test_l3_parallel_reduce.py continue to pass.

N2 — pld.alloc_window_buffer op + pld.window op + DSL + parser — Merged in #1346

Landed the alloc / view decoupling that mirrors tile.alloc / MemRef / TileType.memref_. Two ops, single source file src/ir/op/distributed/memory.cpp (modelled on tile_ops/memory.cpp):

pld.alloc_window_buffer(size, *, name) — pure address-space allocation, size in bytes; returns singleton PtrType (shared with tile.alloc). Parser routes through standard _assign_or_let, binding the LHS as a plain Var(PtrType) (NOT a WindowBuffer).

pld.window(buf, [shape], *, dtype) — view materialization; the deducer checks buf.type is PtrType and returns DistributedTensorType(shape, dtype, /*window_buffer=*/nullopt). The N4 pass fills back window_buffer_.

IR type-layer redesign (also part of the N2 PR):

• WindowBufferType collapsed to a singleton marker (no fields, GetWindowBufferType() returns the shared singleton); folded into the MemRefType / UnknownType / PtrType singleton disjunction in structural_equal.cpp / structural_hash.cpp / serialization.
• WindowBuffer upgraded from IRNode to Var subclass (mirroring MemRef), with base / size / load_from_host / store_to_host fields and no name / dtype fields. Its SSA-edge type is the singleton WindowBufferType (never on any SSA edge — only the GetType() of WindowBuffer itself). N4 pass is the only constructor.
• DistributedTensorType gained window_buffer_: optional<WindowBufferPtr> as a UsualField (participates in structural eq / hash). Same shape / dtype but different window_buffer → structurally NOT equal. Two views sharing the same WindowBuffer shared_ptr are structurally equal.

DSL (python/pypto/language/distributed/alloc.py): adds both alloc_window_buffer and window sentinels; parser intercepts and lifts to OpExpr.

Parser:

• buf = pld.alloc_window_buffer(size) requires single-target LHS, globally-unique name (parser-extracted from LHS, injected as op kwarg; users cannot supply any kwargs including dtype). size is a scalar (int / ir.Expr); list / tuple is rejected. LHS is bound as plain Var(PtrType) — exact mirror of mem_vec_7: Ptr = tile.alloc(...).
• dist_t = pld.window(buf, [shape], *, dtype=...) requires buf to have PtrType at parse time and only the dtype kwarg. Lifts to OpExpr(op="pld.window", inputs=[buf, MakeTuple(shape)], attrs={dtype}) with output type DistributedTensorType(shape, dtype) and window_buffer_ = nullopt.
• "host_orch only" is enforced by the op's host-only flag (uniform with tensor.create), not duplicated in parser.

Tests (all passing on merge):

• tests/ut/ir/test_distributed_ops.py — 9 cases covering WindowBufferType singleton, alloc returning the shared singleton PtrType, name non-empty, LHS being a plain Var(PtrType) (not WindowBuffer), pld.window returning DistributedTensorType(shape, dtype) with window_buffer is None, structural equality between views sharing a WindowBuffer instance, and structural inequality between views with different window_buffer.
• tests/ut/ir/parser/test_alloc_window_buffer.py — 8 cases: LHS is a plain ir.Var of ir.PtrType (not WindowBuffer); Call carries name kwarg matching the LHS, no dtype kwarg; multiple allocs share the singleton PtrType; non-Name LHS / duplicate name / any user kwarg / bare call / list as size are each rejected.
• tests/ut/ir/parser/test_window_op.py — 8 cases: result is DistributedTensorType(shape, dtype) with window_buffer is None at parse; multi-dim shape pass-through; non-Ptr arg / unknown kwarg / missing shape or dtype rejected; nested for OK; cross-function name uniqueness.
• N1 schema regression (test_comm_group_schema.py / test_chip_bootstrap_configs.py / test_distributed_tensor_type.py / test_distributed_tensor_annotation.py) — 16 cases re-run green.

N3 — pld.world_size() + dispatch device=r — Merged in #1359

PR #1359 landed:

• IR op pld.world_size() — host-only registration; returns ScalarType(INT64). Codegen lowers to len(contexts) in the generated host_orch Python. Multiple call sites within one function body are distinct OpExpr instances; a later simplify pass CSEs them to one SSA value when needed.
• Usage note: design doc §1 reads pl.range(pld.world_size) (attribute style); the implementation uses pl.range(pld.world_size()) (function call) to match the IR-op model. Equivalent for readers; if the attribute style is later preferred, the DSL can wrap pld.world_size to auto-lift in host_orch context.
• Call.attrs["device"] for dispatch — reuses the existing attrs: dict[str, ExprPtr] slot on Call. Parser only allows it on FunctionType.Orchestration callees; the device= value must be ConstInt or the induction var of an enclosing for r in pl.range(...) loop. Other shapes (function calls, expressions) raise ParserSyntaxError.
• Tests: positive device=0 / device=r in pl.range(pld.world_size()); negative on InCore callees / non-induction-var r / device=foo().

N4 — CollectCommGroups pass — Merged in #1369

This is where WindowBuffer Var instances finally come into existence — the schema landed in N1 and was upgraded to a Var subclass in N2, but no WindowBuffer exists at parse.

Pass input (post-parse state):

• alloc LHS is plain Var(PtrType); alloc Call op has a name kwarg.
• pld.window result type DistributedTensorType.window_buffer_ == nullopt.
• program.comm_groups_ == [].

Algorithm (mirrors design doc §2.3):

1. Enumerate every pld.alloc_window_buffer Call op — its AssignStmt LHS is a plain Var(PtrType) (call it ptr_var).
2. From each ptr_var, walk def-use to find all pld.window(ptr_var, ...) Call ops.
3. From each pld.window result dist_t, walk def-use to all self.<orch>(..., dist_t, ..., device=r) consumers, threading through chip_orch param bindings (so the pass must run before InlineFunctions).
4. Build a DeviceDescriptor per consumer (ConstInt N → subset = {N} / kAll if r is the induction var of pl.range(pld.world_size()) / subset = {0..N-1} if it's the induction var of pl.range(ConstInt N)). Merge by "any kAll → kAll; otherwise union".
5. Construct one WindowBuffer Var per alloc (base = ptr_var, size = alloc.args[0] bytes, load_from_host = false, store_to_host = false at N4; name_hint inherits from base->name_hint_ via the Var ctor). Host-staging flags default to false; final user-facing API for them is deferred to N7.
6. In-place rewrite every consuming pld.window Call op's result type so window_buffer_ goes from nullopt to wb — the same shared_ptr<const WindowBuffer> is threaded across all consumers of one alloc.
7. Cluster WindowBuffers by device descriptor → CommGroups; each CommGroup.slots_[k] shares its shared_ptr with the corresponding pld.window result type's window_buffer_.
8. Attach { group_idx, slot_idx } attrs to each alloc Call op (lowering index for dist_t.comm). Dispatch calls do NOT get a group attr — each DistributedTensor arg resolves its own group via dist_t.type.window_buffer.
9. Validate: every alloc has at least one pld.window → dispatch consumer (otherwise dead alloc → ValueError); per-group name_hint uniqueness (parser already enforces program-wide uniqueness; pass sanity-checks).

Visitor framework wire-up — N2 made WindowBuffer a Var subclass (ObjectKind::WindowBuffer), but ExprFunctor::VisitExpr's dispatch table only has IterArg → MemRef → Var branches. Once N4 actually constructs WindowBuffer and attaches it into Program.comm_groups_ / DistributedTensorType.window_buffer_, any visitor / mutator calling VisitExpr(window_buffer_ptr) would fall through and raise TypeError("Unknown expression type") (since As<Var> doesn't match the subclass kind). The N4 PR must therefore first wire up dispatch:

• include/pypto/ir/transforms/base/functor.h — add virtual R VisitExpr_(const WindowBufferPtr& op, ...) and an EXPR_FUNCTOR_DISPATCH(WindowBuffer) entry alongside MemRef.
• include/pypto/ir/transforms/base/visitor.h + src/ir/transforms/visitor.cpp — add a default VisitExpr_(const WindowBufferPtr&) empty implementation (mirroring MemRef).
• include/pypto/ir/transforms/base/mutator.h + src/ir/transforms/mutator.cpp — add VisitExpr_(const WindowBufferPtr&) that recursively mutates base_ / size_ and mints a fresh WindowBuffer if either changed.
• 11 downstream functor subclasses (the ones grep'd from VisitExpr_(const MemRefPtr) — each needs a WindowBuffer overload. Most can defer to the base / treat it identically to MemRef; python_printer.cpp needs a print form.

This "framework wire-up" must be the first commits of the N4 PR (independently verifiable) before the CollectCommGroups algorithm itself is written.

Pass registration: python/pypto/ir/pass_manager.py default strategy inserts CollectCommGroups before InlineFunctions. PassProperties: required = (), produced = IRProperty::kCommGroupsCollected, invalidated = (). Doc: docs/en/dev/passes/<NN>-collect_comm_groups.md numbered per .claude/rules/pass-doc-ordering.md.

N5 — CommCtxType + dist_t.comm.rank/.nranks desugar

• New CommCtxType : Type (no shape/dtype, marker only). Ops pld.get_comm_ctx(dist_t) / pld.comm_ctx.rank(ctx) / pld.comm_ctx.nranks(ctx) with strict type verifiers.
• Parser desugars data.comm.rank → pld.comm_ctx.rank(pld.get_comm_ctx(data)); t.comm on a plain Tensor is rejected by the type verifier.

N6 — Cross-rank ops + CommContext layout self-check (P0)

• Enums NotifyOp { kAtomicAdd, kSet }, WaitCmp { kEq, kGe } in include/pypto/ir/comm.h, exposed as pld.NotifyOp / pld.WaitCmp.
• pld.tile.remote_load(target, peer, offsets, *, shape) -> TileType(shape, target.dtype)
• pld.system.notify(target, peer, offsets, value, *, op: NotifyOp)
• pld.system.wait(signal, offsets, expected, *, cmp: WaitCmp)
• All three reject targets/signals whose IR type isn't DistributedTensorType.
• P0 task: compile-time CommContext layout self-check — codegen emits CommRemotePtr against fixed offsets into the runtime's CommContext struct (rankId, windowsIn, windowsOut). To prevent silent miscompiles when runtime adds/reorders fields, add src/codegen/distributed/comm_layout.h with constexpr field offsets pulled via offsetof from runtime's platform_comm/comm_context.h, guarded by static_assert(kRankIdOffset == 16, ...) etc. Codegen emits these comm_layout::k* values into the kernel string instead of magic numbers. Failure mode: runtime ABI drift breaks PyPTO cmake --build immediately. Requires CMake target_include_directories to add runtime's common/ include path.

N7 — Codegen + driver host-staging + host-staging API

• DistributedCodegen dispatches per callee param type: DistributedTensorType formals emit chip_args.add_tensor(ContinuousTensor.make(data=ctx.buffer_ptrs["<name>"], shapes=(...,), dtype=..., child_memory=True), TensorArgType.INOUT). The <name> is read from dist_t.type.window_buffer.name_hint (filled by N4). Plain TensorType formals unchanged. One add_scalar(ctx.device_ctx) per group (single group simplification for the demo; multi-group picks per-formal-idx ctx).
• Call.attrs["device"] lowers to submit_next_level(..., worker=<expr>). pld.world_size() lowers to len(contexts).
• InCore: dist_t.comm.rank/.nranks lower to CommContext->rankId/rankNum reads (using comm_layout::k* offsets from N6); pld.tile.remote_load → CommRemotePtr + TLOAD; notify → CommRemotePtr + TNOTIFY; wait → TWAIT.
• Ctx scalar passing: InCore signatures don't declare a CommCtx* formal — codegen flattens one CommContext* scalar per DistributedTensorType formal at the end of the kernel arg list (and at every level above: L2 orch / kernel_wrapper / L3 host_orch chip_args.add_scalar). Even when two DistributedTensor formals share a group, two scalars are emitted (InCore is group-unaware; per-formal-idx consumption keeps lowering uniform).
• comm_window_size is no longer user-supplied on DistributedConfig — driver computes it from sum(int(wb.size) for wb in CommGroup.slots_) after N4.
• Driver host-staging: load_from_host host tensors (first dim = world_size) are sliced per-rank into HostBufferStaging, attached to cfg.host_inputs; store_to_host analogous for cfg.host_outputs.
• Final user-facing surface for load_from_host / store_to_host is decided here (N4 defaults them to false; the surface — alloc kwarg, separate op, manifest-only — is TBD).

N8 — End-to-end + docs

• tests/st/distributed/test_l3_allreduce.py (2× a2a3 cards): golden outputs[r,i] == nranks*i + 100*nranks*(nranks-1)/2.
• Migrate the user-facing slice of the design doc into docs/en/dev/distributed/l3_allreduce.md (and docs/zh-cn/).

Risks / open items

1. CollectCommGroups before InlineFunctions: the def-use traversal needs chip_orch param bindings intact; pass position must come before inlining. Inline must not drop the device=r attr.
2. Visitor framework wire-up is a prerequisite for N4 — the WindowBuffer Var-subclass dispatch must be added to functor / visitor / mutator (and 11 downstream subclasses) before CollectCommGroups runs; otherwise any visitor touching N4's product hits Unknown expression type. Treat as the first commits of the N4 PR.
3. Multi-group chip_args layout: design doc §4.1 simplifies to a single add_scalar for the demo. N7 lays out per-formal-idx ctx for multi-group; two formals sharing a group still get two scalars (InCore is group-unaware).
4. pld.world_size syntax: function call (pld.world_size()) shipped in N3 vs. attribute (pld.world_size) in design doc §1. Function-call form lands now for IR symmetry; revisit if attribute form is preferred.
5. Host-staging API surface: N4 defaults load_from_host / store_to_host to false; the final user-facing surface (alloc kwarg, separate op, manifest-only) is open until N7.
6. CommContext ABI lock (N6 P0): runtime submodule pin drift is the only thing that breaks the static_asserts in comm_layout.h. If runtime later upstreams these constexprs into comm_context.h, PyPTO's comm_layout.h degenerates to a re-export — fine, but requires coordination with the runtime team given the Directory Ownership rule.