Distributed Operators (N6)¶

Overview¶

The N6 distributed op family gives the Python DSL direct, typed access to the hardware's cross-rank communication primitives. Every op operates against a window-bound DistributedTensorType — a tensor whose storage is a slice of a symmetric, per-rank communication window allocated by pld.alloc_window_buffer. Verifiers in this family generally reject plain TensorType (strict kind-trait matching — As<DistributedTensorType> does not match a plain TensorType), so a non-window-bound tensor can never be fed into a cross-rank slot by accident. Two documented exceptions: pld.tensor.put (and its lowered pld.tile.put) accepts a plain Tensor on the src side via AsTensorTypeLike — TPUT only needs a readable local GM region for the source, so kernels can push directly from host-backed inputs without first staging through a window buffer; dst still requires a window-bound DistributedTensor. pld.tensor.get (and its lowered pld.tile.get) accepts a plain Tensor on the dst side via AsTensorTypeLike — TGET only needs a writable local GM region to receive into, so kernels can TGET directly into host-backed output tensors; src still requires a window-bound DistributedTensor.

There are twelve ops and four ABI enums:

Op	Direction	Result	Hardware
`pld.tile.remote_load`	pull (read peer → local tile)	`TileType`	TLOAD
`pld.tile.remote_store`	push (write local tile → peer)	`Unknown` (side effect)	TSTORE
`pld.tensor.get`	pull (read peer → local GM)	`Unknown` (side effect)	TGET
`pld.tensor.put`	push (write local → peer)	`Unknown` (side effect)	TPUT
`pld.tensor.allreduce`	collective reduce over window slices	`DistributedTensorType` (same as src)	builtin collective
`pld.tensor.barrier`	synchronise visibility of window data across ranks	`DistributedTensorType` (same as src)	builtin collective
`pld.tensor.broadcast`	replicate root rank's data to all ranks	`DistributedTensorType` (same as src)	builtin collective
`pld.tensor.reduce_scatter`	reduce and scatter chunks across ranks	`DistributedTensorType` (same as src)	builtin collective
`pld.tensor.allgather`	gather data from all ranks via window	`DistributedTensorType` (same as src)	builtin collective
`pld.tensor.all_to_all`	push-based symmetric personalized exchange — every rank pushes its per-destination chunks to every peer's window via `pld.tensor.put` (TPUT), returns window as result	`DistributedTensorType` (same as src)	composite / HOST builtin
`pld.system.notify`	signal a peer's slot	`Unknown` (side effect)	TNOTIFY
`pld.system.wait`	block on own slot	`Unknown` (side effect)	TWAIT

The five side-effect-only ops produce UnknownType: they exist for their cross-rank effect, not for an SSA value a consumer reads.

Namespacing: why `tile.` vs `tensor.` vs `system.*`¶

The namespace encodes the IR level the op lives at, not an arbitrary grouping:

pld.tile.remote_load produces a tile (on-core SRAM region), so it is a sibling of tile.load and lives in pld.tile.
pld.tile.remote_store consumes a tile (the symmetric write companion of remote_load), so it is a sibling of tile.store and lives in pld.tile.
pld.tensor.get reads and writes tensor (GM) operands — dst may be a window-bound DistributedTensor or a plain Tensor (TGET only needs a writable local GM region to receive into) while src must be a window-bound DistributedTensor (the peer needs a window slot to read from). The VEC staging tile that TGET bounces through is materialised by ConvertTensorToTileOps as an internal pld.tile.get, never on the DSL surface. It is therefore a sibling of pld.tensor.alloc_window_buffer / pld.tensor.window, not of the tile-producing remote_load.
pld.tensor.put reads and writes tensor (GM) operands — dst is a window-bound DistributedTensor (the peer needs a window slot to receive into) while src accepts either a window-bound DistributedTensor or a plain Tensor (TPUT only needs a readable local GM region on the source side). The VEC staging tile that TPUT bounces through is materialised by ConvertTensorToTileOps as an internal pld.tile.put, never on the DSL surface. It is therefore a sibling of pld.tensor.alloc_window_buffer / pld.tensor.window, not of the tile-producing remote_load.
pld.system.notify / pld.system.wait drive the per-rank signal slot — pure control-plane synchronisation with no data operand — so they live in pld.system.

ABI enums (`include/pypto/ir/comm.h`)¶

The four enums are an append-only ABI. Their underlying int values are serialised as the op's kwarg payload (op for notify, cmp for wait, atomic for put) and cast back to the enum at codegen time. New variants may only be added at the end so existing IR and cached programs keep their meaning.

enum class NotifyOp : int { kAtomicAdd = 0, kSet = 1 };   // pld.system.notify
enum class WaitCmp  : int { kEq = 0,        kGe = 1 };     // pld.system.wait
enum class AtomicType : int { kNone = 0,    kAdd = 1 };    // pld.tensor.put
enum class ReduceOp : int { kSum = 0, kMax = 1, kMin = 2, kProd = 3 };  // pld.tensor.allreduce

Enum	Variant	Meaning
`NotifyOp`	`kAtomicAdd`	atomically add `value` into the peer's signal slot
`NotifyOp`	`kSet`	non-atomic store of `value` into the peer's signal slot
`WaitCmp`	`kEq`	block until `*signal_slot == expected`
`WaitCmp`	`kGe`	block until `*signal_slot >= expected`
`AtomicType`	`kNone`	plain remote store — overwrite the peer's dst slice
`AtomicType`	`kAdd`	atomically add the source data into the peer's dst slice
`ReduceOp`	`kSum`	sum-reduce every participating rank's window slice
`ReduceOp`	`kMax`	max-reduce every participating rank's window slice
`ReduceOp`	`kMin`	min-reduce every participating rank's window slice
`ReduceOp`	`kProd`	product-reduce every participating rank's window slice

Each enum is mirrored across three layers (C++ enum class → nb::enum_ in the bindings → .pyi stub) and surfaced to the DSL as pld.NotifyOp / pld.WaitCmp / pld.AtomicType / pld.ReduceOp. The deducer validates the packed int against the enum range so codegen can cast back without a second guard.

Op reference¶

`pld.tile.remote_load` (TLOAD)¶

pld.tile.remote_load(target, peer, offsets, shape[, valid_shape])
    -> TileType(shape, target.dtype)

Reads a region of the peer rank's slice of a window-bound DistributedTensor into a local tile. Mirrors tile.load at the IR level (positional offsets / shape tuples, TileType result) but the source is a remote slice — the address translation is realised at codegen by inline CommContext loads and offset arithmetic followed by addptr + make_tensor_view.

valid_shape is optional. With or without it, type inference intersects the requested window with the source tensor's effective valid region and checks provable physical bounds. When present, shape remains the physical UB tile allocation while valid_shape additionally limits the remote partition and the tile's valid extent. This is the fixed-width ragged-tail form used by chunked collectives.

Every symbolic source or requested valid extent that survives inference must be runtime-bound by a kernel scalar, loop variable, or physical tensor-shape parameter. A symbol that appears only in type metadata is rejected during PTO codegen.

Verifier: target must be DistributedTensorType; peer must be a ScalarType rank index; offsets / shape / optional valid_shape must each be a MakeTuple whose rank equals target.shape.size().

DSL (python/pypto/language/distributed/op/tile_ops.py) accepts positional or keyword arguments; the IR op keeps them positional, matching tile.load.

`pld.tile.remote_store` (TSTORE)¶

pld.tile.remote_store(src_tile, target, peer, offsets) -> Unknown

Writes a local tile into a region of the peer rank's slice of a window-bound DistributedTensor. Mirrors tile.store at the IR level (positional offsets tuple + side-effect-only return) but the destination is a remote slice — address translation happens at codegen via inline CommContext loads and offset arithmetic followed by addptr + make_tensor_view.

Verifier: src_tile must be TileType; target must be DistributedTensorType; peer must be a ScalarType rank index; offsets must be a MakeTuple whose rank equals target.shape.size(); src_tile.dtype must match target.dtype.

Codegen: the tile is 2-D (height × width) after the standard tile pipeline; the emitted pto.partition_view has the same rank as target, with the leading (target.rank - 2) dims set to size 1 (matching notify's one_dims(rank, "1") pattern). This lets a 2-D tile push land on the inner two dims of any N-D peer slice (N ≥ 2) without forcing the caller to reshape — and it is the regression guard against the older codegen that emitted a fixed-2D partition_view regardless of target rank.

DSL (python/pypto/language/distributed/op/tile_ops.py) exposes target / peer / offsets as keyword-only for readability; the IR op keeps them positional, matching tile.store.

`pld.tensor.put` (TPUT)¶

pld.tensor.put(dst, peer, src, *, atomic: int,
               chunk_rows: int = 0, chunk_cols: int = 0, pipeline: bool = False) -> Unknown
pld.tensor.put(dst, peer, src, dst_offsets, src_offsets, shape,
               *, atomic: int, chunk_rows: int = 0, chunk_cols: int = 0, pipeline: bool = False) -> Unknown

Synchronously writes local src data into the peer rank's slice of the window-bound dst. dst is a GM-level DistributedTensor view; src may be either a DistributedTensor view or a plain Tensor — TPUT only requires a readable local GM region on the source side, so kernels can push directly from host-backed inputs without first staging through a window buffer. The VEC staging tile is materialised by ConvertTensorToTileOps as an internal tile.create + pld.tile.put, so it flows through PyPTO's memory allocator and never appears on the DSL surface.

With no offsets/shape this writes the full local src slice to the full peer dst slice. Supplying dst_offsets, src_offsets, and shape narrows the transfer to matching subregions; all three must be provided together.

Staging-tile chunking. By default the staging tile spans the whole flattened transfer [rows, cols] extent (rows = product of the leading dims, cols = the innermost dim), so a transfer must fit in UB. The optional chunk_rows / chunk_cols attrs (0 = full) shrink the staging tile to a sub-tile of that extent; the codegen keeps the pto.comm.tput partition views at the full transfer extent and pto-isa TPUT 2-D-slides the transfer through the smaller stage. This lets a single put move data larger than UB without the caller writing an explicit chunk loop. Oversized chunk values are clamped to the transfer extent.

Double-buffering (pipeline). Setting pipeline=True makes ConvertTensorToTileOps materialise two identical VEC staging tiles (tput_stage_ping / tput_stage_pong) and thread both into pld.tile.put as a second stage operand. The codegen then emits the ping-pong form pto.comm.tput(dst_pv, src_pv, buf(%ping, %pong) : …), which PTOAS routes to pto-isa's double-buffered TPUT overload — it overlaps the TLOAD of the next chunk with the TSTORE of the previous one across the two tiles. Because the benefit only exists when the transfer is chunked into more than one piece, pipeline requires both chunk_rows and chunk_cols to be set (the deducer and the DSL both reject pipeline without a full chunk). The two tiles are distinct tile.create allocations, so the memory allocator gives them non-overlapping UB addresses (pto-isa's ping/pong requirement).

Dynamic transfer extent. The transfer may be dynamic — either the subregion shape (a runtime sub-extent of the window) or the dst / src window (DistributedTensorType) dims themselves, for a full-slice transfer. pto-isa reads the extent from the partition views at runtime, so the codegen emits a dynamic partition view (<?x…>) and chunks it. A dynamic flattened transfer dim must be bounded by the corresponding static chunk, because the VEC staging tile is statically allocated: a dynamic innermost dim requires chunk_cols, a dynamic leading dim requires chunk_rows. For a full-slice transfer the dst and src dims must match — by value when static, structurally when dynamic.

Verifier: dst must be DistributedTensorType; src must be either TensorType or DistributedTensorType (matched via AsTensorTypeLike); peer must be a ScalarType; dst and src must share element type, rank, and positive dimensions (positivity checked on static dims; dynamic dims are allowed and bounded by the chunk). Full-slice put requires matching dst / src shape; subregion put allows different full slice extents as long as the explicit transfer region is in bounds (checked on static dims). Any dynamic transfer dim requires a matching static chunk (see above). atomic selects overwrite vs atomic-add (see AtomicType). The lowered pld.tile.put verifier requires the staging tile to fit within the flattened transfer in both static dims (it may be smaller — a chunk — but never larger; dynamic dims are bounded by the chunk at runtime).

`pld.tensor.get` (TGET)¶

pld.tensor.get(dst, peer, src, *, chunk_rows: int = 0, chunk_cols: int = 0, pipeline: bool = False) -> Unknown
pld.tensor.get(dst, peer, src, dst_offsets, src_offsets, shape,
               *, chunk_rows: int = 0, chunk_cols: int = 0, pipeline: bool = False) -> Unknown

Synchronously reads the peer rank's slice of the window-bound src into the local dst. dst may be a window-bound DistributedTensor or a plain Tensor; src must be a window-bound DistributedTensor. The VEC staging tile is materialised by ConvertTensorToTileOps as an internal tile.create + pld.tile.get, so it flows through PyPTO's memory allocator and never appears on the DSL surface.

With no offsets/shape this reads the full peer src slice into the full local dst slice. Supplying dst_offsets, src_offsets, and shape narrows the transfer to matching subregions; all three must be provided together. The optional chunk_rows / chunk_cols attrs (0 = full) shrink the staging tile to a sub-tile of the flattened transfer extent so pto-isa TGET auto-chunks the full transfer through it — same contract as put above, including a dynamic transfer (the subregion shape or the full-slice dst / src window dims) bounded by a matching static chunk (dynamic innermost needs chunk_cols, dynamic leading needs chunk_rows). Setting pipeline=True double-buffers the chunked read through two staging tiles (tget_stage_ping / tget_stage_pong), emitting pto.comm.tget(…, buf(%ping, %pong) : …) for pto-isa's ping-pong TGET overload — same contract as put, and likewise requires both chunk_rows and chunk_cols.

Verifier: dst must be either TensorType or DistributedTensorType (matched via AsTensorTypeLike); src must be DistributedTensorType; peer must be a ScalarType; dst and src must share element type, rank, and positive dimensions (positivity checked on static dims; dynamic dims allowed, bounded by the chunk). Full-slice get requires matching dst / src shape; subregion get allows different full slice extents as long as the explicit transfer region is in bounds (checked on static dims); any dynamic transfer dim requires a matching static chunk. Besides chunk_rows / chunk_cols, get accepts no keyword attributes.

`pld.tensor.allreduce`¶

pld.tensor.allreduce(src, *, op: ReduceOp = ReduceOp.Sum, mode: str = "mesh") -> DistributedTensorType(src)
pld.tensor.allreduce(src, signal, *, op: ReduceOp = ReduceOp.Sum, mode: str = "mesh") -> DistributedTensorType(src)

Reduces every participating rank's window-bound src slice in place and returns the same type as src. The mode keyword selects the lowering algorithm:

Fully-valid packed mesh targets are viewed as one logical [1, N] stream and processed in UB chunks of at most 16 KiB. A statically known N smaller than that budget shrinks the physical chunk to the smallest 32-byte-aligned width that covers N; larger or dynamic inputs retain the maximum width. The last chunk keeps its selected physical width but carries valid_shape=[1, min(chunk, N-offset)], so arbitrary element counts neither read nor store past the end.

For mesh lowering, a partial TensorView.valid_shape is preserved for packed ND targets when its valid box can be represented by collapsing the leading dimensions to one 2D rectangle and a statically bounded physical tile fits within one 16-KiB chunk. A symbolic valid extent falls back to the source's physical rectangle when that rectangle fits the budget. Oversized partial rectangles, strided targets, DN partial views, and non-representable partial boxes are rejected explicitly.

Any symbolic target or partial-valid extent that survives lowering must be runtime-bound by a kernel scalar, loop variable, or physical tensor-shape parameter; a type-metadata-only symbol is rejected during PTO codegen. A fully dynamic physical target dimension is bound from that tensor parameter.

"mesh" (default) — direct all-to-all exchange with O(P) HCCL windows. Signal shape [NR, 1] (one cell per rank). An AtomicAdd 1 / wait ≥1 ready barrier precedes the chunk loop. Each chunk performs remote_load+accumulate, then AtomicAdd 1 / wait for its monotonic chunk counter before store-back, preventing write-after-read races.
"ring" — NCCL-style chunked reduce-scatter + allgather schedule with O(1) HCCL windows. Signal shape [2 * (NR − 1), NR] (one row per ring round, one cell per rank). A packed ND target is viewed as one logical [1, SIZE] stream; a partial valid box must be a contiguous row-major prefix. Lowering keeps the full physical [1, product(target.shape)] view and records the logical prefix as TensorView.valid_shape=[1, product(target.valid_shape)]. FP32 divides SIZE with balanced floor(i * SIZE / NR) boundaries. FP16 rounds each interior boundary up to 16 elements (32 bytes) and caps it at SIZE, so every non-empty segment begins at an MTE-safe address without changing the packed user-visible layout. Empty segments remain legal for very short inputs. Each segment is processed in at most 16-KiB physical subchunks; an FP16 ragged remote tail rounds only its physical read span to 32 bytes and restores the logical valid_shape before reduction or store. Every subchunk uses ready and read-complete barriers on the round's monotonic signal row before store-back, preventing write-after-read races while keeping the signal shape unchanged.

Host-orchestrator user code may omit signal outside for and while loops; the SynthesizeAllReduceSignals pass inserts a private INT32 signal window with semantic shape [world_size, 1] for that call (mesh mode only — mode="ring" requires an explicit signal). The pass binds world_size = pld.world_size() as a standalone statement and uses that variable in the synthesized buffer size and window shape. All calls in loops are rejected because the current signal protocol is single-use. Explicit signal remains the internal form used by InCore lowering and by tests that intentionally construct the internal protocol. Comm-domain materialisation then keeps the signal buffer in the same domain as src, even when it is not passed to a user chip kernel. Mesh, ring, and host-builtin paths support FP16 and FP32 with ReduceOp.Sum, Max, Min, and Prod for arbitrary positive element counts. InCore lowering uses UB-bounded chunks; the host builtin uses 256-element chunks. InCore mesh and ring round only the physical FP16 remote tail span to 32 bytes. The host builtin rounds ragged FP16 and FP32 load spans to 32 bytes. Both preserve the logical valid shape. The host builtin accepts either a rank-1 [world_size] signal or the synthesized rank-2 [world_size, 1] signal.

`pld.system.notify` (TNOTIFY)¶

pld.system.notify(target, peer, offsets, value, *, op: int) -> Unknown

Writes value into the peer rank's signal slot of target (a window-bound DistributedTensor, typically a 1-D INT32 "signal matrix"). op selects atomic-add vs set (see NotifyOp).

Verifier: target must be DistributedTensorType; peer and value must be ScalarType; offsets must be a MakeTuple of rank equal to the target rank.

`pld.system.wait` (TWAIT)¶

pld.system.wait(signal, offsets, expected, *, cmp: int) -> Unknown

Blocks until this rank's own signal slot of signal satisfies the cmp predicate against expected (see WaitCmp).

Verifier: signal must be DistributedTensorType; expected must be ScalarType; offsets must be a MakeTuple of rank equal to the signal rank.

Shared codegen infrastructure¶

All five ops lower through PTO codegen helpers in src/backend/common/pto_ops_distributed.cpp and src/codegen/pto/pto_codegen.cpp. The reusable pieces — shared so each op's lowering carries no bespoke peer arithmetic — are:

Helper	Role
`EmitCommRemoteOffset`	emits inline `CommContext` loads and converts the peer-vs-local byte delta to an element offset
`EmitCommRemoteView`	emits the inline offset calculation followed by `addptr + make_tensor_view`, yielding the peer-addressed view (used by `remote_load`, `get`'s `src`, and `put`'s `dst`)
`EmitPartitionViewPTO`	wraps a tensor view in a full-slice `partition_view` with given offsets/sizes (used by every op for both local and peer operands)
`ResolveDistTensorBinding`	resolves a `DistributedTensor` arg to its codegen binding (type + window var)
`AsTensorTypeLike`	kind-trait downcast accepting both `TensorType` and `DistributedTensorType` where a view's element/shape info is read uniformly

The local-vs-remote split is intentional: a local operand (e.g. get's dst, put's src, wait's signal) reuses the tensor view already created by EmitMakeTensorViews with no peer arithmetic, while a remote operand (e.g. remote_load's target, get's src, put's dst) goes through EmitCommRemoteView.

Pipeline integration¶

Comm domains and their slot allocations are materialised by the MaterializeCommDomainScopes pass, which wraps each host_orch body in nested CommDomainScopeStmt nodes (one per inferred comm domain) and produces the per-window WindowBuffer records that the runtime binds physical buffers to. Host-level tensor collectives are then lowered by LowerHostTensorCollectives into internal builtin chip dispatches before the final Simplify.

Testing¶

IR / parser: tests/ut/ir/parser/test_remote_load.py, tests/ut/ir/parser/test_remote_store.py, test_system_ops.py, test_get_op.py, test_put_op.py, plus the negative verifier coverage in tests/ut/ir/test_distributed_ops.py.
Codegen: tests/ut/codegen/distributed/test_distributed_pto_codegen.py.
End-to-end (ST): tests/st/distributed/test_l3_allreduce.py (mesh allreduce with dynamic rank dim NR = pl.dynamic("NR"); P=2 default, P=4 on any four devices (e.g. --device=0,1,2,3 or --device=0-3)), test_l3_allgather.py, test_l3_reduce_scatter.py, test_l3_broadcast.py (each likewise dynamic-NR, P=2/P=4), test_l3_tensor_allreduce_intrinsic.py, test_l3_tensor_allreduce_ring_intrinsic.py, test_l3_allreduce_ring.py (hand-rolled ring RS+AG), test_l3_host_tensor_allreduce.py, test_l3_ep_dispatch_combine.py, test_l3_notify_wait.py, and related L3 STs under tests/st/distributed/. Put/get canonical e2e contracts are now enabled: test_l3_put.py (ring overwrite, row-offset put, atomic-add put, and chunked/pipelined transfers ✅), test_l3_get.py (ring read, row-offset get ✅), and test_l3_remote_store.py (tile-level subview push ✅). All tests use the pld.system.notify / pld.system.wait handshake pattern established by notify/wait and collective STs.

Distributed Operators (N6)¶

Overview¶

Namespacing: why tile.* vs tensor.* vs system.*¶

ABI enums (include/pypto/ir/comm.h)¶

Op reference¶

pld.tile.remote_load (TLOAD)¶

pld.tile.remote_store (TSTORE)¶

pld.tensor.put (TPUT)¶

pld.tensor.get (TGET)¶

pld.tensor.allreduce¶

pld.system.notify (TNOTIFY)¶

pld.system.wait (TWAIT)¶