diff --git a/docs/conceptual/triggered-sdma.md b/docs/conceptual/triggered-sdma.md new file mode 100644 index 000000000..6e89772da --- /dev/null +++ b/docs/conceptual/triggered-sdma.md @@ -0,0 +1,127 @@ +# Triggered SDMA + +Triggered SDMA moves data that a running kernel produces without spending compute units on the +transfer. The host pre-arms copy-engine (SDMA) chains that wait on *gates*; the kernel opens a +gate when its data is written, and the copy engine copies the data to every peer and signals its +arrival there. + +``` +kernel writes a batch ──► publish(batch) opens the gate + │ + copy engine: POLL_REGMEM(ready >= epoch) ─► COPY × segments ─► ATOMIC(peer arrival += 1) + │ +peer kernel: wait(batch) ◄────┘ +``` + +The API is `iris.triggered`; see the [reference](../reference/host/triggered.md). + +## Model + +- A **segment** is one copy: a byte range, or a 2D sub-window (rows of a strided tensor). It is + written at the same offset in each peer's buffer. A segment is independent of how the producer + tiles its output: `waves()` makes one per tile, `chunks()` one per row block (more above 1 GiB), + and a hand-built schedule may cover part of a tile or several adjacent tiles. +- A **batch** is an ordered group of segments behind one gate, with one arrival signal per peer. + Batching amortizes the poll and the signal over many copies. +- A batch has **contributors**: the number of `publish` calls that complete it each epoch. One for + a batch written by a single CTA; the number of CTAs for a batch they write together. +- A **schedule** lists a rank's batches in the order its producer completes them. The plan arms + them in that order on each peer's queue. + +`iris.triggered.schedule` builds schedules for common producers: `whole()` (one batch, published +once per program at the end), `chunks()` (row blocks completed in order) and `waves()` (one batch +per wave of a persistent tiled kernel, such as a GEMM). Any list of batches works. + +```python +from iris.triggered import TriggeredPlan, TriggeredView, waves + +ctx = iris.iris(heap_size, allocator_type="vmem_chunked", engines=("sdma",)) +out = ctx.zeros(world_size * M, N, dtype=torch.bfloat16) +schedules = [waves(out[r * M : (r + 1) * M], (BM, BN), num_programs, base=out) for r in range(world_size)] +plan = TriggeredPlan(ctx, out, schedules) # collective: allocates gates, arms epoch 1 + +for step in range(steps): + view = plan.begin_device_epoch() + producer[(num_programs,)](..., view) # TriggeredView.initialize(view).publish(batch) + plan.end_device_epoch() # every batch arrived and drained + ... + plan.next_epoch() # barrier, re-arm, barrier +plan.destroy() # collective: peers may still touch this rank's memory +``` + +## The kernel author's contract + +1. **Batches complete in schedule order.** All of a rank's chains to one peer share one queue and + run strictly in arming order, so a batch completed early waits for every batch armed before it. + Out-of-order completion is still correct, only slower. A persistent kernel whose programs each + work through their tiles in order completes waves in order by construction. +2. **Each batch gets exactly `contributors` publishes per epoch,** each after the caller's writes + to that batch. `publish` must be called by the whole CTA, outside divergent branches. +3. **A producer never waits on its own plan's arrivals before it has published everything it + contributes to.** Its wait could depend on a chain parked behind one of its own unpublished + batches. +4. **Segment sizes are fixed when the plan is created.** +5. **`arrived`, `wait` and `reusable` are called uniformly across the CTA:** every wave makes the + same number of calls. Each is one atomic whose result reaches the other waves through a CTA + barrier, so a polling loop must not exit on a per-wave condition (such as a clock read): + waves would read the result before it is written, or deadlock. Count polls instead. + +Plans that share queues must begin, and be aborted, in the order they were armed. While chains +are armed on a queue, host `Iris.put()` and `quiet()` to that peer raise rather than queue behind +them. + +Creating a plan and `next_epoch()` arm on every rank or on none: if one rank cannot (for example +its queue is full), every rank raises, and a failed construction drains what it armed. After a +failed `next_epoch()` the plan is either still complete or `FAILED`; `abort()` a failed plan. +`destroy()` is collective and waits for every rank, since a peer's chains and polls can still touch +a rank's buffer and gates after its own work is done. + +## Memory ordering + +Measured on 8× MI350X (gfx950) and 8× MI300X (gfx942). + +**Producer.** Copy engines read memory behind the GPU's L2, so a batch's data must be written +back before its gate opens. A release fence only covers the issuing wave, so every wave of the +CTA must first complete its stores. `publish` does this: + +```python +wait_cnt() # s_waitcnt vmcnt(0): this wave's stores are done +tl.debug_barrier() # every wave's stores are done +tl.atomic_add(counter, 1, sem="acq_rel", scope="sys") # the last contributor then: +tl.atomic_xchg(ready, epoch, sem="release", scope="sys") +``` + +| Release recipe | Stale data at the peer (256 KiB and 4 MiB) | +| ----------------------------------------------- | ------------------------------------------ | +| `wait_cnt` + barrier + release (`gpu` or `sys`) | none in 350 epochs | +| release without the barrier | 80–100% of epochs | +| `.wt` store of the gate | every element, every epoch | +| plain store | none, but the gate stays in L2 until kernel exit | + +**Consumer.** `arrived` / `wait` read the arrival gate with an acquire at system scope; data read +after it in the same kernel is fresh (no stale elements over 9,600 tile-epochs). + +## Gates and epochs + +- Each batch has a 64-byte gate row: `ready`, `arrival`, `completion` (64-bit), `error`, and a + contributor counter. The copy engine polls the low 32 bits of `ready` with `>=`. +- Epochs start at 1 and add 1 per `next_epoch()`. Gates are never reset, so a plan is limited to + 2³² − 1 epochs. +- `end_device_epoch()` checks completion with a kernel: this rank's batches published and arrived + at every peer (each chain ends with that signal), and every other batch arrived here. + +## Requirements and limits + +- **Stable peer mappings.** Armed chains hold peer addresses, so the heap must not re-map peers + while chains are armed. Create the context with `allocator_type="vmem_chunked"` (or `"vmem"`): + both keep peer addresses fixed, and the chunked allocator maps new chunks into each peer's + reserved range as the heap grows. The default torch allocator re-maps peer heaps on every + allocation and is rejected. +- **SDMA engine.** Create the context with `engines=("sdma",)` (needs `pip install iris[sdma]`). +- **One queue per peer.** Plans use the host-initiated copy-engine queue to each peer, shared with + host `Iris.put()` and with the plans of every Iris context in the process. Kernel-initiated + copies (`put(..., use_copy_engine=True)`) use a separate queue and are unaffected. +- **Ring capacity.** Armed chains may use at most half of each queue's 8 MiB ring. +- **Single node only.** +- **Linear segments are at most 1 GiB** (SDMA 4.4's count field is 30 bits). The schedule builders + split larger ranges. diff --git a/docs/index.md b/docs/index.md index 721e7b0ba..373e4fa4b 100644 --- a/docs/index.md +++ b/docs/index.md @@ -209,6 +209,7 @@ For other setup methods, see the [Installation Guide](getting-started/installati - **[Programming Model](conceptual/programming-model.md)**: How Iris works - **[Fine-grained Overlap](conceptual/finegrained-overlap.md)**: GEMM & communication overlap - **[Taxonomy](conceptual/taxonomy.md)**: Multi-GPU programming patterns +- **[Triggered SDMA](conceptual/triggered-sdma.md)**: Copy-engine transfers released by kernels ### 📖 **Reference** - **[Talks and Papers](reference/talks-and-papers.md)**: Publications, conference talks, and videos diff --git a/docs/reference/host/triggered.md b/docs/reference/host/triggered.md new file mode 100644 index 000000000..d37aeac7f --- /dev/null +++ b/docs/reference/host/triggered.md @@ -0,0 +1,25 @@ +# Triggered SDMA + +Pre-armed copy-engine transfers released by kernels. See +[Triggered SDMA](../../conceptual/triggered-sdma.md) for the model and its ordering rules. + +## TriggeredPlan +```{eval-rst} +.. autoclass:: iris.triggered.TriggeredPlan + :members: begin_device_epoch, end_device_epoch, next_epoch, run, abort, destroy, batch_id +``` + +## Schedules +```{eval-rst} +.. autoclass:: iris.triggered.Segment +.. autoclass:: iris.triggered.Batch +.. autofunction:: iris.triggered.whole +.. autofunction:: iris.triggered.chunks +.. autofunction:: iris.triggered.waves +``` + +## TriggeredView (device) +```{eval-rst} +.. autoclass:: iris.triggered.TriggeredView + :members: initialize, publish, arrived, wait, reusable, can_publish, can_consume +``` diff --git a/docs/sphinx/_toc.yml b/docs/sphinx/_toc.yml index 49ce51d06..8c32e31ec 100644 --- a/docs/sphinx/_toc.yml +++ b/docs/sphinx/_toc.yml @@ -11,6 +11,7 @@ subtrees: - file: conceptual/programming-model.md - file: conceptual/finegrained-overlap.md - file: conceptual/taxonomy.md + - file: conceptual/triggered-sdma.md - caption: Reference entries: - file: reference/talks-and-papers.md @@ -20,6 +21,7 @@ subtrees: - file: reference/host/tensor-creation.md - file: reference/host/ccl.md - file: reference/host/copy-engine.md + - file: reference/host/triggered.md - file: reference/triton/overview.md entries: - file: reference/triton/device-context.md diff --git a/docs/sphinx/_toc.yml.in b/docs/sphinx/_toc.yml.in index 49ce51d06..8c32e31ec 100644 --- a/docs/sphinx/_toc.yml.in +++ b/docs/sphinx/_toc.yml.in @@ -11,6 +11,7 @@ subtrees: - file: conceptual/programming-model.md - file: conceptual/finegrained-overlap.md - file: conceptual/taxonomy.md + - file: conceptual/triggered-sdma.md - caption: Reference entries: - file: reference/talks-and-papers.md @@ -20,6 +21,7 @@ subtrees: - file: reference/host/tensor-creation.md - file: reference/host/ccl.md - file: reference/host/copy-engine.md + - file: reference/host/triggered.md - file: reference/triton/overview.md entries: - file: reference/triton/device-context.md diff --git a/iris/__init__.py b/iris/__init__.py index 38125e1ed..46fd8dd15 100644 --- a/iris/__init__.py +++ b/iris/__init__.py @@ -83,6 +83,7 @@ from iris.host.platform import hip from . import experimental from . import ops +from . import triggered from iris.host.memory import tensors as tensor_creation from . import bench from iris.host.tracing import kernel_artifacts # noqa: F401 # triggers _init() at import time @@ -127,6 +128,7 @@ "hip", "experimental", "ops", + "triggered", "tensor_creation", "bench", "set_logger_level", diff --git a/iris/host/iris.py b/iris/host/iris.py index a399c7e65..24d7d55b3 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -54,6 +54,7 @@ ) from iris.host.memory.symmetric_heap import SymmetricHeap +from iris.host import queue_ledger import numpy as np from typing import Any import torch @@ -1113,6 +1114,7 @@ def put( if to_tensor is None: to_tensor = from_tensor + self._check_host_queue_idle(to_rank) from_rank = self.get_rank() from_ptr = from_tensor.data_ptr() to_ptr = self.heap.translate(to_tensor.data_ptr(), from_rank, to_rank) @@ -1203,6 +1205,7 @@ def put_tile( ... wait_flag=wait_ptr, wait_value=256, signal_flag=signal_ptr) """ sdma_ep = self._require_copy_engine() + self._check_host_queue_idle(to_rank) from_rank = self.get_rank() wait_ptr, wait_bits = self._flag_pointer_and_bits(wait_flag, default_bits=32) @@ -1269,6 +1272,7 @@ def put_tiles( channel: SDMA channel to use """ sdma_ep = self._require_copy_engine() + self._check_host_queue_idle(to_rank) from_rank = self.get_rank() if len(tiles) != len(to_ptrs) or len(tiles) != len(to_strides): @@ -1328,12 +1332,18 @@ def quiet(self, to_rank: int = None, channel: int = 0): sdma_ep = self._require_copy_engine() src_rank = self.get_rank() if to_rank is not None: + self._check_host_queue_idle(to_rank) sdma_ep.quiet(src_rank, to_rank, channel) else: # Quiet to all ranks for rank in range(self.get_num_ranks()): + self._check_host_queue_idle(rank) sdma_ep.quiet(src_rank, rank, channel) + def _check_host_queue_idle(self, to_rank): + # Chains parked on a host queue block everything queued behind them, quiet() included + queue_ledger.ledger().check_idle(self.get_rank(), to_rank) + def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Tensor]: """ Allocate a symmetric tensor and return it with its peer-base table. diff --git a/iris/host/queue_ledger.py b/iris/host/queue_ledger.py new file mode 100644 index 000000000..b765000f7 --- /dev/null +++ b/iris/host/queue_ledger.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +What iris.triggered plans have armed on this process's host copy-engine queues. + +The queues belong to the process, not to an Iris instance, so the ledger does too. +""" + +from collections import defaultdict + +from iris.device.sdma_utils import bind_sdma_ep + + +class QueueLedger: + """What is armed on each of this process's host queues, in arming order.""" + + def __init__(self, capacity): + self.capacity = capacity + self.entries = defaultdict(list) # (src, dst) -> [(plan, epoch, nbytes)] + + def reserve(self, plan, epoch, src, nbytes): + for peer, n in nbytes.items(): + used = sum(e[2] for e in self.entries[(src, peer)]) + if used + n > self.capacity: + raise RuntimeError( + f"arming {n} bytes on the queue to rank {peer} would exceed {self.capacity} armed bytes " + f"({used} already armed); complete or destroy other plans first" + ) + for peer, n in nbytes.items(): + self.entries[(src, peer)].append((plan, epoch, n)) + + def release(self, plan): + for key in self.entries: + self.entries[key] = [e for e in self.entries[key] if e[0] is not plan] + + def check_order(self, plan, allowed, action): + for (_, peer), entries in self.entries.items(): + for other, epoch, _ in entries: + if other is plan: + break + if other.state not in allowed: + raise RuntimeError( + f"plan armed earlier on the queue to rank {peer} (epoch {epoch}) is {other.state.value}; " + f"plans sharing queues must {action} in the order they were armed" + ) + + def check_idle(self, src, dst): + if self.entries.get((src, dst)): + raise RuntimeError( + f"triggered chains are armed on the copy-engine queue to rank {dst}; " + "complete, abort or destroy those plans before using it" + ) + + +_ledger = None + + +def ledger(): + global _ledger + if _ledger is None: + # Leave half the ring for wrap padding and work queued before the plan + _ledger = QueueLedger(bind_sdma_ep().SDMA_QUEUE_SIZE // 2) + return _ledger diff --git a/iris/mem/triton/triggered.py b/iris/mem/triton/triggered.py new file mode 100644 index 000000000..b6acf6fff --- /dev/null +++ b/iris/mem/triton/triggered.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Device side of triggered SDMA plans (see ``iris.triggered``). + +The host arms one chain per (batch, peer) on the copy engine: poll the batch's ``ready`` gate, +copy the batch's segments to the peer, then add 1 to the batch's ``arrival`` gate on the peer. +A kernel opens the gate with ``publish`` and observes arrivals with ``arrived`` / ``wait``. + +View layout (int64 words): an 8-word header, then one 4-word record per batch. Gate table: one +64-byte row per batch. +""" + +import triton +import triton.language as tl +from triton.language.core import _aggregate as aggregate + +from iris.mem.utils import wait_cnt + +ABI_VERSION = 1 + +# Header words +HEADER_WORDS = tl.constexpr(8) +H_VERSION = tl.constexpr(0) +H_NUM_BATCHES = tl.constexpr(1) +H_RANK = tl.constexpr(2) +H_EPOCH = tl.constexpr(3) +H_GATES = tl.constexpr(4) +H_WORLD_SIZE = tl.constexpr(5) + +# Batch record words +RECORD_WORDS = tl.constexpr(4) +R_ROLE = tl.constexpr(0) +R_CONTRIBUTORS = tl.constexpr(1) +R_OWNER = tl.constexpr(2) + +ROLE_PUBLISH = tl.constexpr(1) +ROLE_CONSUME = tl.constexpr(2) + +# Gate row words +GATE_WORDS = tl.constexpr(8) +G_READY = tl.constexpr(0) +G_ARRIVAL = tl.constexpr(1) +G_COMPLETION = tl.constexpr(2) +G_ERROR = tl.constexpr(3) +G_COUNTER = tl.constexpr(4) + + +@aggregate +class TriggeredView: + """ + Device handle for one epoch of a triggered plan. + + Usage:: + + @triton.jit + def producer(out, view, ...): + tv = TriggeredView.initialize(view) + ... # write this program's part of batch b + tv.publish(b) # whole CTA, after its writes + + @triton.jit + def consumer(out, view, ...): + tv = TriggeredView.initialize(view) + tv.wait(b) # batch b's data is now visible + x = tl.load(out + ...) + """ + + view: tl.tensor + gates: tl.tensor + epoch: tl.tensor + rank: tl.tensor + + @triton.constexpr_function + def __init__(self, view, gates, epoch, rank): + self.view = view + self.gates = gates + self.epoch = epoch + self.rank = rank + + @staticmethod + @triton.jit + def initialize(view): + """Read the plan's view tensor (from ``TriggeredPlan.begin_device_epoch()``).""" + gates = tl.cast(tl.load(view + H_GATES), tl.pointer_type(tl.int64)) + return TriggeredView(view, gates, tl.load(view + H_EPOCH), tl.load(view + H_RANK)) + + @triton.jit + def _record(self, batch, word): + return tl.load(self.view + HEADER_WORDS + batch * RECORD_WORDS + word) + + @triton.jit + def can_publish(self, batch): + return (self._record(batch, R_ROLE) & ROLE_PUBLISH) != 0 + + @triton.jit + def can_consume(self, batch): + return (self._record(batch, R_ROLE) & ROLE_CONSUME) != 0 + + @triton.jit + def publish(self, batch): + """ + Contribute to batch ``batch``; the last of its contributors opens the gate. + + Call from the whole CTA, outside divergent branches, after the CTA's writes to the batch. + """ + # Copy engines read behind L2: every wave's stores must land before the release writes back + wait_cnt() + tl.debug_barrier() + gate = self.gates + batch * GATE_WORDS + contributors = self._record(batch, R_CONTRIBUTORS) + ticket = tl.atomic_add(gate + G_COUNTER, 1, sem="acq_rel", scope="sys") + if ticket % contributors == contributors - 1: + tl.atomic_xchg(gate + G_READY, self.epoch, sem="release", scope="sys") + + @triton.jit + def arrived(self, batch): + """ + True once batch ``batch`` is visible here this epoch (published, for this rank's own). + + Call uniformly across the CTA: the result reaches the other waves through a CTA barrier, so + every wave must make the same number of calls (true of ``wait`` and ``reusable`` too). + """ + gate = self.gates + batch * GATE_WORDS + word = tl.where(self._record(batch, R_OWNER) == self.rank, G_READY, G_ARRIVAL) + return tl.atomic_add(gate + word, 0, sem="acquire", scope="sys") >= self.epoch + + @triton.jit + def wait(self, batch): + """Spin until ``arrived(batch)``.""" + while self.arrived(batch) == 0: + pass + + @triton.jit + def reusable(self, batch): + """True once batch ``batch`` is complete this epoch (set by ``end_device_epoch()``).""" + gate = self.gates + batch * GATE_WORDS + return tl.atomic_add(gate + G_COMPLETION, 0, sem="acquire", scope="sys") >= self.epoch diff --git a/iris/triggered/__init__.py b/iris/triggered/__init__.py new file mode 100644 index 000000000..e5ff52148 --- /dev/null +++ b/iris/triggered/__init__.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Triggered SDMA: the host pre-arms copy-engine chains, a kernel releases them per batch. + +Host side: ``TriggeredPlan`` and the schedule builders in ``iris.triggered.schedule``. +Device side: ``TriggeredView`` (``publish``, ``arrived``, ``wait``). +""" + +from iris.mem.triton.triggered import TriggeredView +from iris.triggered.plan import State, TriggeredPlan +from iris.triggered.schedule import Batch, Segment, chunks, waves, whole + +__all__ = ["TriggeredPlan", "TriggeredView", "State", "Batch", "Segment", "whole", "chunks", "waves"] diff --git a/iris/triggered/armer.py b/iris/triggered/armer.py new file mode 100644 index 000000000..5afbc0bd7 --- /dev/null +++ b/iris/triggered/armer.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Arms triggered chains on the host-initiated copy-engine queue through rocm-xio. + +Per (batch, peer): POLL_REGMEM(ready >= epoch) -> one copy per segment -> ATOMIC ADD_RTN_64(arrival += 1). +""" + +import ctypes +from dataclasses import dataclass + +from iris.device.sdma_utils import bind_sdma_ep + +CHANNEL = 0 +POLL_REGMEM_BYTES = 24 + +# nanobind passes void* as a PyCapsule named "nb_handle" +_NB_HANDLE = b"nb_handle" +_capsule_new = ctypes.pythonapi.PyCapsule_New +_capsule_new.restype = ctypes.py_object +_capsule_new.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + + +def _void_ptr(addr): + return _capsule_new(addr, _NB_HANDLE, None) + + +@dataclass(frozen=True) +class Copy: + src: int + dst: int + width: int + height: int + pitch: int + + +def chain_bytes(batch): + sdma_ep = bind_sdma_ep() + copy_bytes = sdma_ep.COPY_LINEAR_SUB_WINDOW_COMMAND_BYTES if batch.is_2d else sdma_ep.COPY_LINEAR_COMMAND_BYTES + return POLL_REGMEM_BYTES + len(batch.segments) * copy_bytes + sdma_ep.ATOMIC_COMMAND_BYTES + + +def arm_chain(src_device, dst_device, ready_addr, epoch, copies, arrival_addr, is_2d): + """Arm one chain on the host queue from ``src_device`` to ``dst_device``.""" + sdma_ep = bind_sdma_ep() + if is_2d: + tiles = [] + for c in copies: + tile = sdma_ep.Tile() + tile.data = _void_ptr(c.src) + tile.pid_m = tile.pid_n = 0 + tile.block_m, tile.block_n = c.height, c.width + tile.elem_size = 1 + tile.src_stride = c.pitch + tiles.append(tile) + sdma_ep.wait_flag_then_put_tiles( + src_device, + dst_device, + CHANNEL, + ready_addr, + epoch, + tiles, + [c.dst for c in copies], + [c.pitch for c in copies], + ) + else: + first, *rest = copies + sdma_ep.wait_flag_then_put( + src_device, dst_device, CHANNEL, ready_addr, epoch, first.src, first.dst, first.width, 32 + ) + for c in rest: + sdma_ep.put(src_device, dst_device, CHANNEL, c.src, c.dst, c.width) + sdma_ep.signal(src_device, dst_device, CHANNEL, arrival_addr, 1, 64) diff --git a/iris/triggered/kernels.py b/iris/triggered/kernels.py new file mode 100644 index 000000000..e06f32709 --- /dev/null +++ b/iris/triggered/kernels.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""Kernels a triggered plan launches itself: releasing its own batches and checking completion.""" + +import triton +import triton.language as tl + +from iris.mem.triton.context import __translate +from iris.mem.triton.triggered import ( + G_ARRIVAL, + G_COMPLETION, + G_READY, + GATE_WORDS, + HEADER_WORDS, + R_OWNER, + RECORD_WORDS, + TriggeredView, +) + + +@triton.jit +def release_kernel(view, batches): + """Open the gate of batch ``batches[pid]``, as its last contributor would.""" + tv = TriggeredView.initialize(view) + batch = tl.load(batches + tl.program_id(0)) + tl.atomic_xchg(tv.gates + batch * GATE_WORDS + G_READY, tv.epoch, sem="release", scope="sys") + + +@triton.jit +def _reached(ptr, epoch, budget: tl.constexpr): + i = 0 + while (tl.atomic_add(ptr, 0, sem="acquire", scope="sys") < epoch) & (i < budget): + i += 1 + return tl.atomic_add(ptr, 0, sem="acquire", scope="sys") >= epoch + + +@triton.jit +def progress_kernel( + view, heap_bases, armed, status, BUDGET: tl.constexpr, WORLD_SIZE: tl.constexpr, OWN_ONLY: tl.constexpr +): + """ + Mark batch ``pid`` complete once it is done this epoch: for this rank's own batches, published + and arrived at every peer it was armed for (``armed[batch * WORLD_SIZE + peer]``; each chain + ends with that signal, so it has fully run); for the others, arrived here. Writes 1 to ``status[pid]`` when complete, 0 if ``BUDGET`` ran out. + """ + batch = tl.program_id(0) + tv = TriggeredView.initialize(view) + gate = tv.gates + batch * GATE_WORDS + own = tl.load(view + HEADER_WORDS + batch * RECORD_WORDS + R_OWNER) == tv.rank + done = tl.full((), True, tl.int1) + if own: + done = _reached(gate + G_READY, tv.epoch, BUDGET) + for peer in tl.static_range(WORLD_SIZE): + if done & (peer != tv.rank) & (tl.load(armed + batch * WORLD_SIZE + peer) != 0): + done = _reached(__translate(gate + G_ARRIVAL, tv.rank, peer, heap_bases), tv.epoch, BUDGET) + elif not OWN_ONLY: + done = _reached(gate + G_ARRIVAL, tv.epoch, BUDGET) + if done: + tl.atomic_xchg(gate + G_COMPLETION, tv.epoch, sem="release", scope="sys") + tl.store(status + batch, done.to(tl.int32)) diff --git a/iris/triggered/plan.py b/iris/triggered/plan.py new file mode 100644 index 000000000..764f1ea5a --- /dev/null +++ b/iris/triggered/plan.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Triggered SDMA plans: the host arms copy-engine chains, a kernel releases them. + +Lifecycle: + + plan = TriggeredPlan(ctx, buffer, schedules) # collective: allocate gates, arm epoch 1 + view = plan.begin_device_epoch() + producer[grid](..., view) # TriggeredView.publish(batch) + plan.end_device_epoch() # every batch complete, or TimeoutError + plan.next_epoch() # barrier -> re-arm -> barrier + plan.destroy() # collective + +All of a rank's chains to one peer share that peer's host-initiated copy-engine queue, and run in +the order they were armed. Plans armed on the same queues must therefore be released in the order +they were armed, and nothing else may use those queues while chains are armed on them. +""" + +import enum +import hashlib +import time + +import numpy as np +import torch + +from iris.host import queue_ledger +from iris.host.distributed.helpers import distributed_allgather +from iris.host.memory.allocators.torch_allocator import TorchAllocator +from iris.mem.triton import triggered as dv +from iris.triggered import armer, kernels +from iris.triggered.schedule import validate + +UINT32_MAX = (1 << 32) - 1 +_PROGRESS_BUDGET = 100_000 + + +class State(enum.Enum): + ARMED = "armed" + ACTIVE = "active" + COMPLETE = "complete" + FAILED = "failed" + ABORTED = "aborted" + DESTROYED = "destroyed" + + +def _digest(buffer_bytes, schedules): + text = repr((buffer_bytes, schedules)).encode() + return np.frombuffer(hashlib.sha256(text).digest()[:8], dtype=np.int64).copy() + + +def view_words(rank, world_size, gates_ptr, schedules): + """The device view (see ``iris.mem.triton.triggered``) as a list of int64 words.""" + num_batches = sum(len(s) for s in schedules) + words = [0] * (dv.HEADER_WORDS.value + num_batches * dv.RECORD_WORDS.value) + words[dv.H_VERSION.value] = dv.ABI_VERSION + words[dv.H_NUM_BATCHES.value] = num_batches + words[dv.H_RANK.value] = rank + words[dv.H_GATES.value] = gates_ptr + words[dv.H_WORLD_SIZE.value] = world_size + rec = dv.HEADER_WORDS.value + for owner, schedule in enumerate(schedules): + for batch in schedule: + words[rec + dv.R_ROLE.value] = dv.ROLE_PUBLISH.value if owner == rank else dv.ROLE_CONSUME.value + words[rec + dv.R_CONTRIBUTORS.value] = batch.contributors + words[rec + dv.R_OWNER.value] = owner + rec += dv.RECORD_WORDS.value + return words + + +class TriggeredPlan: + """ + Copies ``buffer`` regions between ranks with pre-armed copy-engine chains. + + Rank ``r`` owns the batches in ``schedules[r]`` and sends each of them to every other rank, at + the same offsets in the peer's ``buffer``. Batch ids are global: rank ``r``'s batches come after + those of ranks ``0..r-1`` (see ``batch_id``). + + Collective: every rank passes the same ``schedules`` for the same symmetric ``buffer``. + + Args: + ctx: Iris instance. + buffer: Contiguous tensor on the symmetric heap. The heap must keep peer mappings stable while + chains are armed, which rules out the default torch allocator: use ``allocator_type="vmem_chunked"`` (or ``"vmem"``). + schedules: One list of ``Batch`` per rank, each in the order that rank completes them. + timeout: Seconds ``end_device_epoch()`` waits before raising. + """ + + def __init__(self, ctx, buffer, schedules, timeout=30.0): + self.ctx = ctx + self.rank = ctx.get_rank() + self.world_size = ctx.get_num_ranks() + self.timeout = timeout + self.state = None + ctx._require_copy_engine() + + # Armed chains hold peer addresses; the torch allocator re-maps peer heaps on every allocation + if isinstance(ctx.heap.allocator, TorchAllocator): + raise ValueError( + "triggered plans need stable peer mappings: create the Iris context with allocator_type='vmem_chunked' (or 'vmem')" + ) + if not ctx.heap.is_symmetric(buffer) or not buffer.is_contiguous(): + raise ValueError("buffer must be a contiguous tensor on the symmetric heap") + buffer_bytes = buffer.numel() * buffer.element_size() + if len(schedules) != self.world_size: + raise ValueError(f"need one schedule per rank ({self.world_size}), got {len(schedules)}") + self.schedules = [list(validate(s, buffer_bytes)) for s in schedules] + digests = distributed_allgather(_digest(buffer_bytes, self.schedules)) + if not (digests == digests[0]).all(): + raise ValueError("ranks passed different buffers sizes or schedules") + + self._first = np.cumsum([0] + [len(s) for s in self.schedules]).tolist() + self.num_batches = self._first[-1] + if self.num_batches == 0: + raise ValueError("schedules contain no batches") + + self.buffer = buffer + self.gates = ctx.zeros(self.num_batches, dv.GATE_WORDS.value, dtype=torch.int64) + # FIXME: remove once symmetric allocation is fixed. The vmem_chunked heap frees on each rank's + # GC, so its free lists can diverge and place the same allocation at different offsets + heap_base = int(ctx.heap.heap_bases_cpu[self.rank]) + offsets = np.array([buffer.data_ptr() - heap_base, self.gates.data_ptr() - heap_base], dtype=np.int64) + if not (distributed_allgather(offsets) == offsets).all(): + raise RuntimeError("buffer or gates sit at different heap offsets on different ranks") + words = view_words(self.rank, self.world_size, self.gates.data_ptr(), self.schedules) + self.view = torch.tensor(words, dtype=torch.int64, device=buffer.device) + self._own = torch.arange( + self._first[self.rank], self._first[self.rank + 1], dtype=torch.int64, device=buffer.device + ) + self._status = torch.zeros(self.num_batches, dtype=torch.int32, device=buffer.device) + self._armed = torch.zeros(self.num_batches * self.world_size, dtype=torch.int32, device=buffer.device) + self._ledger = queue_ledger.ledger() + self.epoch = 0 + try: + self._arm(1) + except Exception as error: + self._unwind_construction(error) + ctx.barrier() + + def batch_id(self, rank, index): + """Global id of batch ``index`` in ``schedules[rank]``.""" + return self._first[rank] + index + + def _agree(self, error, what): + """Collective: raise on every rank if any rank has an error, so none is left waiting for the others.""" + failed = distributed_allgather(np.array([error is not None], dtype=np.int64))[:, 0].nonzero()[0].tolist() + if failed: + if error is not None: + raise error + raise RuntimeError(f"rank(s) {failed} could not {what}") + + def _unwind_construction(self, error): + """Raise ``error`` from the constructor, after draining the chains this rank armed.""" + if self.state is not None: + # No plan reaches the caller to abort(), so drain here + try: + self.abort() + except Exception as drain_error: + raise RuntimeError(f"{error}; the chains it armed could not be drained: {drain_error}") from error + finally: + # Peers may still be draining into this rank's buffer and gates + self.ctx.barrier() + raise error + + def _arm(self, epoch): + """Collective: every rank arms ``epoch``, or every rank raises (``FAILED`` if chains are left armed).""" + peers = [p for p in range(self.world_size) if p != self.rank] + nbytes = sum(armer.chain_bytes(b) for b in self.schedules[self.rank]) + error, reserved = None, False + try: + if epoch > UINT32_MAX: + raise RuntimeError("epoch exceeds 32 bits; destroy the plan and create a new one") + self._ledger.reserve(self, epoch, self.rank, {p: nbytes for p in peers}) + reserved = True + except Exception as e: + error = e + try: + self._agree(error, "reserve queue space") + except Exception: + # Nothing is armed anywhere + if reserved: + self._ledger.release(self) + raise + self.epoch = epoch + self.state = State.ARMED + + try: + self._submit(peers) + except Exception as e: + error = e + try: + self._agree(error, "arm its chains") + except Exception: + # Chains armed so far stay parked: abort() releases them + self.state = State.FAILED + raise + + def _submit(self, peers): + heap = self.ctx.heap + self._heap_bases = heap.heap_bases_cpu.copy() + base = self.buffer.data_ptr() + gate_row = self.gates.element_size() * dv.GATE_WORDS.value + # Which chains exist, so a failed arm drains only those + armed = np.zeros(self.num_batches * self.world_size, dtype=np.int32) + try: + for i, batch in enumerate(self.schedules[self.rank]): + batch_id = self.batch_id(self.rank, i) + gate = self.gates.data_ptr() + batch_id * gate_row + for peer in peers: + copies = [ + armer.Copy( + base + s.offset, + heap.translate(base + s.offset, self.rank, peer), + s.width, + s.height, + s.row_pitch, + ) + for s in batch.segments + ] + arrival = heap.translate(gate + 8 * dv.G_ARRIVAL.value, self.rank, peer) + armer.arm_chain(self.rank, peer, gate, self.epoch, copies, arrival, batch.is_2d) + armed[batch_id * self.world_size + peer] = 1 + finally: + self._armed.copy_(torch.from_numpy(armed)) + + def _require(self, *states): + if self.state not in states: + names = " or ".join(s.value for s in states) + raise RuntimeError(f"plan is {self.state.value if self.state else 'uninitialized'}, needs {names}") + + def begin_device_epoch(self): + """Start the armed epoch and return the view tensor its kernels take.""" + self._require(State.ARMED) + if not np.array_equal(self.ctx.heap.heap_bases_cpu, self._heap_bases): + raise RuntimeError("symmetric heap was re-mapped since the plan was armed") + self._ledger.check_order(self, (State.ACTIVE,), "begin") + self.view[dv.H_EPOCH.value] = self.epoch + self.state = State.ACTIVE + return self.view + + def _wait(self, own_only, timeout): + deadline = time.monotonic() + (self.timeout if timeout is None else timeout) + while True: + kernels.progress_kernel[(self.num_batches,)]( + self.view, + self.ctx.get_heap_bases(), + self._armed, + self._status, + BUDGET=_PROGRESS_BUDGET, + WORLD_SIZE=self.world_size, + OWN_ONLY=own_only, + ) + status = self._status.cpu() + mine = slice(self._first[self.rank], self._first[self.rank + 1]) + pending = (status[mine] if own_only else status) == 0 + if not pending.any(): + return + if time.monotonic() > deadline: + missing = (pending.nonzero().flatten() + (mine.start if own_only else 0)).tolist() + raise TimeoutError(f"epoch {self.epoch}: batches {missing[:16]} not complete") + + def end_device_epoch(self, timeout=None): + """Wait until every batch is complete this epoch: own ones arrived at every peer, others here.""" + self._require(State.ACTIVE) + try: + self._wait(own_only=False, timeout=timeout) + except TimeoutError: + self.state = State.FAILED + raise + self._ledger.release(self) + self.state = State.COMPLETE + + def next_epoch(self): + """Collective: re-arm every chain for the next epoch. If any rank cannot, every rank raises.""" + self._require(State.COMPLETE) + self.ctx.barrier() + self._arm(self.epoch + 1) + self.ctx.barrier() + + def run(self, timeout=None): + """Run one epoch with no producer kernel: publish every own batch as it stands, then wait.""" + self.begin_device_epoch() + if len(self._own): + kernels.release_kernel[(len(self._own),)](self.view, self._own) + self.end_device_epoch(timeout) + + def abort(self, timeout=None): + """ + Release this rank's armed chains and wait for them to drain. The plan can only be destroyed + afterwards. Peers' chains to this rank are theirs to abort. Plans sharing queues must be + aborted in the order they were armed: later chains wait behind earlier ones. + """ + if self.state not in (State.ARMED, State.ACTIVE, State.FAILED): + return + self._ledger.check_order(self, (State.ACTIVE, State.ABORTED), "abort") + self.view[dv.H_EPOCH.value] = self.epoch + if len(self._own): + kernels.release_kernel[(len(self._own),)](self.view, self._own) + self.state = State.ABORTED + self._wait(own_only=True, timeout=timeout) + self._ledger.release(self) + + def destroy(self): + """ + Collective: wait until every rank is done with the plan, then forget it. Its chains must have + drained: abort() an armed or failed plan first. A peer's chains and polls may still touch this + rank's buffer and gates, so every rank must call destroy() before any of them is freed. + """ + if self.state is State.DESTROYED: + return + if self.state in (State.ARMED, State.ACTIVE, State.FAILED): + raise RuntimeError(f"plan is {self.state.value}; abort() it before destroying") + self._ledger.release(self) + self.ctx.barrier() + self.state = State.DESTROYED + self.gates = self.view = self.buffer = None diff --git a/iris/triggered/schedule.py b/iris/triggered/schedule.py new file mode 100644 index 000000000..1751c2a12 --- /dev/null +++ b/iris/triggered/schedule.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Schedules for triggered SDMA plans. + +A schedule is the list of batches one rank sends, in the order its producer completes them. Each +batch is one gate: once it has received ``contributors`` publishes, the copy engine copies every +segment of the batch to each peer and signals the batch's arrival there. + +Chains on one copy-engine queue run strictly in arming order, so a batch completed early waits for +every batch armed before it. Order the schedule the way the producer finishes its work. +""" + +from dataclasses import dataclass + +# COPY_LINEAR's count field is 30 bits on SDMA 4.4.x +MAX_LINEAR_BYTES = 1 << 30 + + +@dataclass(frozen=True) +class Segment: + """ + One copy: ``height`` rows of ``width`` bytes, ``pitch`` bytes apart, starting ``offset`` bytes + into the plan's buffer. The same bytes are written at the same offset in each peer's buffer. + """ + + offset: int + width: int + height: int = 1 + pitch: int = 0 + + @property + def row_pitch(self): + return self.pitch or self.width + + @property + def end(self): + return self.offset + (self.height - 1) * self.row_pitch + self.width + + +@dataclass(frozen=True) +class Batch: + """A group of segments behind one gate, complete after ``contributors`` publishes per epoch.""" + + segments: tuple + contributors: int = 1 + + @property + def is_2d(self): + return any(s.height > 1 for s in self.segments) + + +def _byte_offset(view, base): + base = view if base is None else base + offset = view.data_ptr() - base.data_ptr() + if offset < 0 or view.untyped_storage().data_ptr() != base.untyped_storage().data_ptr(): + raise ValueError("view must lie inside base") + return offset + + +def _linear(offset, nbytes): + return tuple( + Segment(offset + start, min(MAX_LINEAR_BYTES, nbytes - start)) for start in range(0, nbytes, MAX_LINEAR_BYTES) + ) + + +def whole(view, base=None, contributors=1): + """One batch covering a contiguous ``view``, published ``contributors`` times per epoch.""" + if not view.is_contiguous(): + raise ValueError("whole() needs a contiguous view") + return [Batch(_linear(_byte_offset(view, base), view.numel() * view.element_size()), contributors)] + + +def chunks(view, num_chunks, base=None, contributors=1): + """ + Split a contiguous ``view`` along dim 0 into ``num_chunks`` batches, in order. Suits producers + that finish row blocks in order: row-block, streaming and elementwise kernels. + """ + if not view.is_contiguous(): + raise ValueError("chunks() needs a contiguous view") + rows = view.shape[0] if view.dim() else 1 + row_bytes = view.numel() * view.element_size() // max(rows, 1) + offset = _byte_offset(view, base) + per_chunk = -(-rows // num_chunks) + return [ + Batch(_linear(offset + r * row_bytes, (min(r + per_chunk, rows) - r) * row_bytes), contributors) + for r in range(0, rows, per_chunk) + ] + + +def waves(view, block_shape, num_programs, order=None, base=None): + """ + One batch per wave of a persistent tiled producer. + + ``view`` is a 2D tensor with unit column stride, split into ``block_shape`` tiles (edge tiles + clipped). Tile ``order[i]`` (row-major tile ids by default) is produced by program + ``i % num_programs`` in wave ``i // num_programs``. Each wave's batch holds one sub-window per + tile and completes after one publish per tile. + + A wave is complete only once every program has finished its tile in it, which each does after + its tile in the previous wave, so waves always complete in order. + """ + if view.dim() != 2 or view.stride(1) != 1: + raise ValueError("waves() needs a 2D view with unit column stride") + rows, cols = view.shape + block_m, block_n = block_shape + tiles_n = -(-cols // block_n) + num_tiles = -(-rows // block_m) * tiles_n + order = range(num_tiles) if order is None else list(order) + if sorted(order) != list(range(num_tiles)): + raise ValueError(f"order must be a permutation of the {num_tiles} tile ids") + + elem = view.element_size() + pitch = view.stride(0) * elem + offset = _byte_offset(view, base) + batches = [] + for start in range(0, num_tiles, num_programs): + segments = [] + for tile in order[start : start + num_programs]: + m0, n0 = (tile // tiles_n) * block_m, (tile % tiles_n) * block_n + height, width = min(block_m, rows - m0), min(block_n, cols - n0) + segments.append(Segment(offset + m0 * pitch + n0 * elem, width * elem, height, pitch)) + batches.append(Batch(tuple(segments), len(segments))) + return batches + + +def validate(schedule, buffer_bytes): + """Check every batch of ``schedule`` fits a buffer of ``buffer_bytes`` and can be encoded.""" + for i, batch in enumerate(schedule): + if not batch.segments: + raise ValueError(f"batch {i} has no segments") + if batch.contributors < 1: + raise ValueError(f"batch {i} needs at least one contributor") + for s in batch.segments: + if s.width <= 0 or s.height <= 0 or s.offset < 0: + raise ValueError(f"batch {i}: empty or negative segment {s}") + if s.pitch and s.pitch < s.width: + raise ValueError(f"batch {i}: pitch smaller than width in {s}") + if s.end > buffer_bytes: + raise ValueError(f"batch {i}: segment {s} runs past the buffer ({buffer_bytes} bytes)") + if not batch.is_2d and s.width > MAX_LINEAR_BYTES: + raise ValueError(f"batch {i}: linear segment over {MAX_LINEAR_BYTES} bytes; split it") + return schedule diff --git a/tests/unittests/test_sdma_packets.py b/tests/unittests/test_sdma_packets.py index cb7d83b91..f23a264f5 100644 --- a/tests/unittests/test_sdma_packets.py +++ b/tests/unittests/test_sdma_packets.py @@ -153,3 +153,41 @@ def test_atomic_add_opcode(dtype, opcode): _atomic_add_kernel[(1,)](flag, heap_bases, ctx) addr = flag.data_ptr() assert _dwords(queue, 9) == [(opcode << 25) | 0xA, addr & 0xFFFFFFFF, addr >> 32, 5, 0, 0, 0, 0, SENTINEL] + + +@triton.jit +def _emit_copy(buf, size, src, dst): + sdma_utils.place_copy_packet(buf, tl.full((), 0, tl.uint64), size, src, dst) + + +@triton.jit +def _emit_poll(buf, flag, expected): + sdma_utils.place_poll_regmem_packet(buf, tl.full((), 0, tl.uint64), flag, expected) + + +@triton.jit +def _emit_nop(buf, padding_bytes): + sdma_utils.place_nop_packet(buf, tl.full((), 0, tl.uint64), padding_bytes) + + +def test_copy_linear_packet(): + buf = _buffer() + dst = ADDR + 0x1000 + _emit_copy[(1,)](buf, 4096, ADDR, dst) + # op 1, sub-op 0; count = bytes - 1; parameters; src; dst + assert _dwords(buf, 8) == [1, 4095, 0, ADDR & 0xFFFFFFFF, ADDR >> 32, dst & 0xFFFFFFFF, dst >> 32, SENTINEL] + + +def test_poll_regmem_packet(): + buf = _buffer() + _emit_poll[(1,)](buf, ADDR, 7) + # op 8, func 5 (>=) at bits 30:28, mem_poll bit 31; addr; reference; mask; retry 0xFFF << 16 | interval 10 + header = (1 << 31) | (5 << 28) | 8 + assert _dwords(buf, 7) == [header, ADDR & 0xFFFFFFFF, ADDR >> 32, 7, 0xFFFFFFFF, (0xFFF << 16) | 10, SENTINEL] + + +def test_nop_packet(): + buf = _buffer() + _emit_nop[(1,)](buf, 5 * 4) + # op 0, count = dwords - 1 at bits 29:16, then zeros + assert _dwords(buf, 6) == [4 << 16, 0, 0, 0, 0, SENTINEL] diff --git a/tests/unittests/test_triggered.py b/tests/unittests/test_triggered.py new file mode 100644 index 000000000..51dd5377f --- /dev/null +++ b/tests/unittests/test_triggered.py @@ -0,0 +1,461 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +# Triggered plans on the copy engines. Each rank owns a slice of a symmetric buffer and sends it to +# every peer; batches are released by kernels with TriggeredView.publish. +# +# Every test that arms a plan aborts it on failure, so no chain is left parked on a queue. + +import time +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch +import triton +import triton.language as tl + +import iris +from iris.host import queue_ledger +from iris.mem.triton import triggered as dv +from iris.mem.utils import read_realtime +from iris.triggered import Batch, Segment, State, TriggeredPlan, TriggeredView, armer, chunks, waves, whole + +EMPTY = -1 + + +def _ctx(): + ctx = iris.iris(1 << 26, allocator_type="vmem_chunked", engines=("sdma",)) + if ctx.get_num_ranks() < 2: + pytest.skip("triggered plans need at least two ranks") + return ctx + + +@contextmanager +def _armed(plan): + try: + yield plan + finally: + if plan.state in (State.ARMED, State.ACTIVE, State.FAILED): + plan.abort(timeout=10) + plan.destroy() + + +def _expected(epoch, n, start=0, device="cuda"): + # Epoch in the top byte, element index below: stale or misplaced data never matches + return (epoch << 24) | torch.arange(start, start + n, dtype=torch.int32, device=device) + + +def _check_before_release(ctx, plan, buffer, slice_elems): + """After begin_device_epoch(), before any producer runs: nothing has arrived from any peer.""" + ctx.barrier() + rank = ctx.get_rank() + others = [p for p in range(ctx.get_num_ranks()) if p != rank] + consume = [plan.batch_id(p, i) for p in others for i in range(len(plan.schedules[p]))] + assert (plan.gates[consume, dv.G_ARRIVAL.value] == plan.epoch - 1).all() + for p in others: + assert (buffer[p * slice_elems : (p + 1) * slice_elems] == EMPTY).all() + ctx.barrier() + + +def _check_and_wipe(ctx, buffer, epoch, slice_elems): + assert torch.equal(buffer, _expected(epoch, buffer.numel()).view_as(buffer)) + rank = ctx.get_rank() + mine = buffer.view(-1)[rank * slice_elems : (rank + 1) * slice_elems].clone() + buffer.fill_(EMPTY) + buffer.view(-1)[rank * slice_elems : (rank + 1) * slice_elems] = mine + + +@triton.jit +def _segments_producer(buf, segs, max_segs, view, first, epoch, BLOCK: tl.constexpr): + """Program b writes every segment (element offset, count) of its batch, then publishes it.""" + b = tl.program_id(0) + for s in range(max_segs): + start = tl.load(segs + (b * max_segs + s) * 2) + count = tl.load(segs + (b * max_segs + s) * 2 + 1) + for off in range(0, count, BLOCK): + idx = start + off + tl.arange(0, BLOCK) + tl.store(buf + idx, (epoch << 24) | idx, mask=off + tl.arange(0, BLOCK) < count) + TriggeredView.initialize(view).publish(first + b) + + +def _segment_table(schedule, elem): + max_segs = max(len(b.segments) for b in schedule) + table = torch.zeros(len(schedule), max_segs, 2, dtype=torch.int64) + for i, batch in enumerate(schedule): + for j, s in enumerate(batch.segments): + table[i, j] = torch.tensor([s.offset // elem, s.width // elem]) + return table.cuda(), max_segs + + +def _strided(slice_elems, rank, num_batches, blocks_per_batch): + """Batch j holds blocks j, j + num_batches, ...: several separate segments per batch.""" + block = slice_elems // (num_batches * blocks_per_batch) * 4 + base = rank * slice_elems * 4 + return [ + Batch(tuple(Segment(base + (j + k * num_batches) * block, block) for k in range(blocks_per_batch))) + for j in range(num_batches) + ] + + +@pytest.mark.parametrize("kind", ["chunks", "strided"]) +def test_linear_batches_across_epochs(kind): + ctx = _ctx() + rank, world = ctx.get_rank(), ctx.get_num_ranks() + slice_elems = 16 * 1024 + buffer = ctx.zeros(world * slice_elems, dtype=torch.int32) + buffer.fill_(EMPTY) + if kind == "chunks": + schedules = [chunks(buffer[r * slice_elems : (r + 1) * slice_elems], 4, base=buffer) for r in range(world)] + else: + schedules = [_strided(slice_elems, r, 4, 4) for r in range(world)] + segs, max_segs = _segment_table(schedules[rank], 4) + + with _armed(TriggeredPlan(ctx, buffer, schedules, timeout=20)) as plan: + for epoch in range(1, 4): + if epoch > 1: + plan.next_epoch() + view = plan.begin_device_epoch() + _check_before_release(ctx, plan, buffer, slice_elems) + _segments_producer[(len(schedules[rank]),)]( + buffer, segs, max_segs, view, plan.batch_id(rank, 0), epoch, BLOCK=1024 + ) + plan.end_device_epoch() + _check_and_wipe(ctx, buffer, epoch, slice_elems) + ctx.barrier() + del ctx + + +def test_heap_growth_while_armed(): + """The chunked heap keeps peer addresses fixed as it grows, so armed chains stay valid.""" + ctx = _ctx() + rank, world = ctx.get_rank(), ctx.get_num_ranks() + alloc = ctx.heap.allocator + slice_elems = 4096 + buffer = ctx.zeros(world * slice_elems, dtype=torch.int32) + buffer.fill_(EMPTY) + schedules = [chunks(buffer[r * slice_elems : (r + 1) * slice_elems], 2, base=buffer) for r in range(world)] + segs, max_segs = _segment_table(schedules[rank], 4) + + with _armed(TriggeredPlan(ctx, buffer, schedules, timeout=20)) as plan: + num_chunks = alloc.get_num_chunks() + ctx.zeros(alloc.chunk_size, dtype=torch.int8) # mapped on every rank while chains are armed + assert alloc.get_num_chunks() > num_chunks + view = plan.begin_device_epoch() + _segments_producer[(2,)](buffer, segs, max_segs, view, plan.batch_id(rank, 0), 1, BLOCK=1024) + plan.end_device_epoch() + _check_and_wipe(ctx, buffer, 1, slice_elems) + + # A buffer larger than a chunk spans a chunk boundary on every rank + slice_elems = -(-(alloc.chunk_size // 4 + 1024) // world) + buffer = ctx.zeros(world * slice_elems, dtype=torch.int32) + buffer.fill_(EMPTY) + mine = buffer[rank * slice_elems : (rank + 1) * slice_elems] + mine.copy_(torch.arange(rank * slice_elems, (rank + 1) * slice_elems, dtype=torch.int32, device=mine.device)) + schedules = [whole(buffer[r * slice_elems : (r + 1) * slice_elems], base=buffer) for r in range(world)] + with _armed(TriggeredPlan(ctx, buffer, schedules, timeout=20)) as plan: + plan.run() + assert torch.equal(buffer, torch.arange(buffer.numel(), dtype=torch.int32, device=buffer.device)) + ctx.barrier() + del ctx + + +@triton.jit +def _waves_producer( + buf, view, first, epoch, row0, M, N, BM: tl.constexpr, BN: tl.constexpr, NUM_PROGRAMS: tl.constexpr +): + """Persistent tiled producer: program p makes tiles p, p + NUM_PROGRAMS, ... and publishes each tile's wave.""" + tv = TriggeredView.initialize(view) + tiles_n = tl.cdiv(N, BN) + for i in range(tl.program_id(0), tl.cdiv(M, BM) * tiles_n, NUM_PROGRAMS): + rows = (i // tiles_n) * BM + tl.arange(0, BM)[:, None] + cols = (i % tiles_n) * BN + tl.arange(0, BN)[None, :] + idx = (row0 + rows) * N + cols + tl.store(buf + idx, (epoch << 24) | idx, mask=(rows < M) & (cols < N)) + tv.publish(first + i // NUM_PROGRAMS) + + +def test_waves_2d_unaligned(): + ctx = _ctx() + rank, world = ctx.get_rank(), ctx.get_num_ranks() + m, n, bm, bn, programs = 40, 72, 16, 32, 4 # 3 x 3 tiles, clipped at the edges; waves of 4, 4, 1 + buffer = ctx.zeros(world * m, n, dtype=torch.int32) + buffer.fill_(EMPTY) + schedules = [waves(buffer[r * m : (r + 1) * m], (bm, bn), programs, base=buffer) for r in range(world)] + + with _armed(TriggeredPlan(ctx, buffer, schedules, timeout=20)) as plan: + for epoch in range(1, 4): + if epoch > 1: + plan.next_epoch() + view = plan.begin_device_epoch() + _check_before_release(ctx, plan, buffer.view(-1), m * n) + _waves_producer[(programs,)]( + buffer, view, plan.batch_id(rank, 0), epoch, rank * m, m, n, BM=bm, BN=bn, NUM_PROGRAMS=programs + ) + plan.end_device_epoch() + _check_and_wipe(ctx, buffer, epoch, m * n) + ctx.barrier() + del ctx + + +@triton.jit +def _staggered_producer(buf, view, batch, epoch, start, part, delay, BLOCK: tl.constexpr): + """Program p waits p * delay ticks, writes its part of the batch, then publishes it.""" + pid = tl.program_id(0) + t0 = read_realtime() + while read_realtime() - t0 < pid * delay: + pass + for off in range(0, part, BLOCK): + idx = start + pid * part + off + tl.arange(0, BLOCK) + tl.store(buf + idx, (epoch << 24) | idx) + TriggeredView.initialize(view).publish(batch) + + +def test_many_contributors_one_gate(): + """The chain must wait for the last contributor, and see every contributor's data.""" + ctx = _ctx() + rank, world = ctx.get_rank(), ctx.get_num_ranks() + programs, slice_elems = 8, 1 << 20 # 4 MiB per rank, one batch + buffer = ctx.zeros(world * slice_elems, dtype=torch.int32) + buffer.fill_(EMPTY) + schedules = [ + whole(buffer[r * slice_elems : (r + 1) * slice_elems], base=buffer, contributors=programs) for r in range(world) + ] + + with _armed(TriggeredPlan(ctx, buffer, schedules, timeout=20)) as plan: + for epoch in range(1, 11): + if epoch > 1: + plan.next_epoch() + view = plan.begin_device_epoch() + _staggered_producer[(programs,)]( + buffer, + view, + plan.batch_id(rank, 0), + epoch, + rank * slice_elems, + slice_elems // programs, + 20_000, # 200 us between contributors + BLOCK=1024, + ) + plan.end_device_epoch() + _check_and_wipe(ctx, buffer, epoch, slice_elems) + ctx.barrier() + del ctx + + +@triton.jit +def _publish_kernel(view, batch): + TriggeredView.initialize(view).publish(batch) + + +def test_head_of_line_order(): + """Chains on one queue run in arming order: a batch released early waits for the one armed before it.""" + ctx = _ctx() + rank, world = ctx.get_rank(), ctx.get_num_ranks() + slice_elems = 4096 + buffer = ctx.zeros(world * slice_elems, dtype=torch.int32) + buffer.fill_(EMPTY) + if rank == 0: + buffer[:slice_elems] = _expected(1, slice_elems) + schedules = [chunks(buffer[:slice_elems], 2)] + [[] for _ in range(world - 1)] + + with _armed(TriggeredPlan(ctx, buffer, schedules, timeout=20)) as plan: + view = plan.begin_device_epoch() + if rank == 0: + _publish_kernel[(1,)](view, 1) + ctx.barrier() + if rank != 0: + time.sleep(0.1) + assert (plan.gates[:, dv.G_ARRIVAL.value] == 0).all() + assert (buffer[:slice_elems] == EMPTY).all() + ctx.barrier() + if rank == 0: + _publish_kernel[(1,)](view, 0) + plan.end_device_epoch() + assert torch.equal(buffer[:slice_elems], _expected(1, slice_elems)) + if rank != 0: + assert (plan.gates[:, dv.G_ARRIVAL.value] == 1).all() + ctx.barrier() + del ctx + + +def test_misuse_raises(): + ctx = _ctx() + rank, world = ctx.get_rank(), ctx.get_num_ranks() + slice_elems = 1024 + buffer = ctx.zeros(world * slice_elems, dtype=torch.int32) + schedules = [whole(buffer[r * slice_elems : (r + 1) * slice_elems], base=buffer) for r in range(world)] + peer = (rank + 1) % world + + with _armed(TriggeredPlan(ctx, buffer, schedules, timeout=20)) as a, _armed( + TriggeredPlan(ctx, buffer, schedules, timeout=20) + ) as b: + with pytest.raises(RuntimeError, match="armed"): + ctx.put(buffer[:slice_elems], to_rank=peer) + with pytest.raises(RuntimeError, match="needs"): + a.end_device_epoch() + with pytest.raises(RuntimeError, match="armed earlier"): + b.begin_device_epoch() + with pytest.raises(RuntimeError, match="abort in the order"): + b.abort() + # Used in arming order, two plans on the same queues are fine + for epoch in range(1, 3): + if epoch > 1: + a.next_epoch() + b.next_epoch() + a.run() + b.run() + with pytest.raises(RuntimeError, match="needs"): + a.begin_device_epoch() + ctx.barrier() + del ctx + + +def test_asymmetric_buffer_raises(): + """A buffer at different heap offsets on different ranks would send every copy to the wrong place.""" + ctx = _ctx() + rank, world = ctx.get_rank(), ctx.get_num_ranks() + slice_elems = 1024 + if rank == 0: + pad = ctx.zeros(slice_elems, dtype=torch.int32) # noqa: F841 shifts rank 0's later allocations + buffer = ctx.zeros(world * slice_elems, dtype=torch.int32) + schedules = [whole(buffer[r * slice_elems : (r + 1) * slice_elems], base=buffer) for r in range(world)] + with pytest.raises(RuntimeError, match="different heap offsets"): + TriggeredPlan(ctx, buffer, schedules, timeout=20) + ctx.barrier() + del ctx + + +def test_capacity_failure_fails_every_rank(): + """A rank that cannot reserve queue space fails all ranks before anything is armed.""" + ctx = _ctx() + rank, world = ctx.get_rank(), ctx.get_num_ranks() + slice_elems = 1024 + buffer = ctx.zeros(world * slice_elems, dtype=torch.int32) + schedules = [whole(buffer[r * slice_elems : (r + 1) * slice_elems], base=buffer) for r in range(world)] + ledger = queue_ledger.ledger() + assert ledger is queue_ledger.ledger() + + hog = SimpleNamespace(state=State.ARMED) + if rank == 0: + ledger.reserve(hog, 1, rank, {p: ledger.capacity for p in range(world) if p != rank}) + try: + with pytest.raises(RuntimeError, match="would exceed" if rank == 0 else "could not reserve"): + TriggeredPlan(ctx, buffer, schedules, timeout=20) + finally: + ledger.release(hog) + assert not any(ledger.entries.values()) + + with _armed(TriggeredPlan(ctx, buffer, schedules, timeout=20)) as plan: + plan.run() + ctx.barrier() + del ctx + + +def test_arm_failure_drains_and_fails_every_rank(monkeypatch): + """A failure while arming raises on every rank and drains what was armed, leaving the queues usable.""" + ctx = _ctx() + rank, world = ctx.get_rank(), ctx.get_num_ranks() + slice_elems = 1024 + buffer = ctx.zeros(world * slice_elems, dtype=torch.int32) + schedules = [chunks(buffer[r * slice_elems : (r + 1) * slice_elems], 4, base=buffer) for r in range(world)] + + arm_chain, armed = armer.arm_chain, [] + + def flaky(*args, **kwargs): + if rank == 0 and len(armed) == 2: + raise RuntimeError("injected arm failure") + armed.append(args) + return arm_chain(*args, **kwargs) + + monkeypatch.setattr(armer, "arm_chain", flaky) + with pytest.raises(RuntimeError, match="injected arm failure" if rank == 0 else "could not arm"): + TriggeredPlan(ctx, buffer, schedules, timeout=20) + monkeypatch.undo() + assert not any(queue_ledger.ledger().entries.values()) + + with _armed(TriggeredPlan(ctx, buffer, schedules, timeout=20)) as plan: + plan.run() + ctx.barrier() + del ctx + + +def test_destroy_waits_for_every_rank(): + """Peers may still touch a rank's buffer and gates, so destroy() returns only once all ranks call it.""" + ctx = _ctx() + rank, world = ctx.get_rank(), ctx.get_num_ranks() + slice_elems = 1024 + buffer = ctx.zeros(world * slice_elems, dtype=torch.int32) + schedules = [whole(buffer[r * slice_elems : (r + 1) * slice_elems], base=buffer) for r in range(world)] + plan = TriggeredPlan(ctx, buffer, schedules, timeout=20) + plan.run() + if rank != 0: + time.sleep(1.0) + start = time.monotonic() + plan.destroy() + if rank == 0: + assert time.monotonic() - start > 0.5 + assert plan.state is State.DESTROYED + plan.destroy() # idempotent + ctx.barrier() + del ctx + + +@triton.jit +def _release_flag(flag, value): + tl.atomic_xchg(flag, value, sem="release", scope="sys") + + +def test_host_put_wait_flag(): + """Host put(wait_flag=, signal_flag=): the copy waits for the flag (>=), never reset between epochs.""" + ctx = _ctx() + rank = ctx.get_rank() + n = 4096 + src = ctx.zeros(n, dtype=torch.int32) + dst = ctx.zeros(n, dtype=torch.int32) + flag = ctx.zeros(1, dtype=torch.int32) + signal = ctx.zeros(1, dtype=torch.int32) + epoch = 0 + try: + for epoch in range(1, 6): + if rank == 0: + src.copy_(_expected(epoch, n)) + ctx.put( + src, to_rank=1, to_tensor=dst, wait_flag=flag, wait_value=epoch, signal_flag=signal, async_op=True + ) + ctx.barrier() + if rank == 1: + time.sleep(0.05) + assert signal.item() == epoch - 1 + ctx.barrier() + if rank == 0: + _release_flag[(1,)](flag, epoch) + ctx.quiet(to_rank=1) + ctx.barrier() + if rank == 1: + assert signal.item() == epoch + assert torch.equal(dst, _expected(epoch, n)) + finally: + if rank == 0: + _release_flag[(1,)](flag, epoch) + ctx.quiet(to_rank=1) + ctx.barrier() + del ctx + + +def test_requires_sdma_engine(): + ctx = iris.iris(1 << 20) + buffer = ctx.zeros(1024, dtype=torch.int32) + with pytest.raises(RuntimeError, match="SDMA is disabled"): + TriggeredPlan(ctx, buffer, [whole(buffer)] * ctx.get_num_ranks()) + ctx.barrier() + del ctx + + +def test_torch_allocator_rejected(): + ctx = iris.iris(1 << 20, engines=("sdma",)) + buffer = ctx.zeros(1024, dtype=torch.int32) + with pytest.raises(ValueError, match="vmem"): + TriggeredPlan(ctx, buffer, [whole(buffer)] * ctx.get_num_ranks()) + ctx.barrier() + del ctx diff --git a/tests/unittests/test_triggered_schedule.py b/tests/unittests/test_triggered_schedule.py new file mode 100644 index 000000000..3e4b11e5f --- /dev/null +++ b/tests/unittests/test_triggered_schedule.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +# Schedule builders and validation for triggered plans (host only). + +import pytest +import torch + +from iris.triggered import Batch, Segment, chunks, waves, whole +from iris.triggered.schedule import MAX_LINEAR_BYTES, validate + + +def test_whole(): + base = torch.zeros(4, 64, dtype=torch.float32) + (batch,) = whole(base[2:], base=base, contributors=3) + assert batch == Batch((Segment(2 * 256, 2 * 256),), 3) + assert not batch.is_2d + + +def test_chunks_uneven(): + base = torch.zeros(10, 16, dtype=torch.int32) + batches = chunks(base, 3) + # 10 rows in chunks of 4: rows 0-3, 4-7, 8-9 + assert [b.segments for b in batches] == [ + (Segment(0, 4 * 64),), + (Segment(4 * 64, 4 * 64),), + (Segment(8 * 64, 2 * 64),), + ] + + +def test_waves_clips_edge_tiles(): + base = torch.zeros(2 * 40, 72, dtype=torch.float32) + view = base[40:] + batches = waves(view, (16, 32), num_programs=4, base=base) + # 3 x 3 tiles in waves of 4, 4, 1 + assert [b.contributors for b in batches] == [4, 4, 1] + pitch = 72 * 4 + first = batches[0].segments + assert first[0] == Segment(40 * pitch, 32 * 4, 16, pitch) + assert first[2] == Segment(40 * pitch + 64 * 4, 8 * 4, 16, pitch) # last column: 8 wide + assert batches[2].segments == (Segment(40 * pitch + 32 * pitch + 64 * 4, 8 * 4, 8, pitch),) # corner: 8 x 8 + assert all(b.is_2d for b in batches) + + +def test_waves_order(): + base = torch.zeros(32, 32, dtype=torch.float16) + batches = waves(base, (16, 16), num_programs=2, order=[3, 2, 1, 0]) + assert batches[0].segments[0].offset == 16 * 64 + 32 + with pytest.raises(ValueError): + waves(base, (16, 16), num_programs=2, order=[0, 1, 2]) + + +@pytest.mark.parametrize( + "batch, buffer_bytes", + [ + (Batch(()), 1024), + (Batch((Segment(0, 64),), contributors=0), 1024), + (Batch((Segment(0, 0),)), 1024), + (Batch((Segment(0, 64, 2, 32),)), 1024), # pitch < width + (Batch((Segment(1000, 64),)), 1024), # past the end + (Batch((Segment(0, MAX_LINEAR_BYTES + 1),)), 1 << 32), + ], +) +def test_validate_rejects(batch, buffer_bytes): + with pytest.raises(ValueError): + validate([batch], buffer_bytes) diff --git a/tests/unittests/test_triggered_view.py b/tests/unittests/test_triggered_view.py new file mode 100644 index 000000000..c0fb585bc --- /dev/null +++ b/tests/unittests/test_triggered_view.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +# TriggeredView against an ordinary gate table: no copy engine involved. + +import torch +import triton +import triton.language as tl + +from iris.mem.triton import triggered as dv +from iris.triggered import Batch, Segment, TriggeredView +from iris.triggered.plan import view_words + + +def _view(rank, schedules, epoch): + num_batches = sum(len(s) for s in schedules) + gates = torch.zeros(num_batches, dv.GATE_WORDS.value, dtype=torch.int64, device="cuda") + words = view_words(rank, len(schedules), gates.data_ptr(), schedules) + words[dv.H_EPOCH.value] = epoch + return torch.tensor(words, dtype=torch.int64, device="cuda"), gates + + +@triton.jit +def _publish_kernel(view, batch): + TriggeredView.initialize(view).publish(batch) + + +@triton.jit +def _arrived_kernel(view, out, n): + tv = TriggeredView.initialize(view) + for b in range(n): + tl.store(out + b, tv.arrived(b).to(tl.int32)) + + +def test_publish_opens_gate_on_last_contributor(): + k = 4 + view, gates = _view(0, [[Batch((Segment(0, 64),), contributors=k)]], epoch=1) + _publish_kernel[(k - 1,)](view, 0) + assert gates[0, dv.G_READY.value].item() == 0 + _publish_kernel[(1,)](view, 0) + assert gates[0, dv.G_READY.value].item() == 1 + + # The counter is never reset: the next epoch needs exactly k more publishes + view[dv.H_EPOCH.value] = 2 + _publish_kernel[(k - 1,)](view, 0) + assert gates[0, dv.G_READY.value].item() == 1 + _publish_kernel[(1,)](view, 0) + assert gates[0, dv.G_READY.value].item() == 2 + assert gates[0, dv.G_COUNTER.value].item() == 2 * k + + +def test_arrived_reads_ready_for_own_batches_and_arrival_for_others(): + one = [Batch((Segment(0, 64),))] + view, gates = _view(1, [one, one], epoch=3) + out = torch.zeros(2, dtype=torch.int32, device="cuda") + + gates[0, dv.G_READY.value] = 3 # another rank's gate: not this rank's to observe + gates[1, dv.G_ARRIVAL.value] = 3 # own batch: arrivals happen at peers, not here + _arrived_kernel[(1,)](view, out, 2) + assert out.tolist() == [0, 0] + + gates[0, dv.G_ARRIVAL.value] = 3 + gates[1, dv.G_READY.value] = 3 + _arrived_kernel[(1,)](view, out, 2) + assert out.tolist() == [1, 1]