From 490a11cf779fb262dc138f1155ea5c79f1db4740 Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Wed, 7 Oct 2026 14:10:57 -0500 Subject: [PATCH 1/8] Add golden-byte tests for the SDMA copy, poll and NOP emitters Only the atomic emitters had byte-level tests. Check COPY_LINEAR, POLL_REGMEM and NOP against the layouts in the Linux amdgpu SDMA 4.4 driver, writing into an ordinary buffer rather than a queue. Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/unittests/test_sdma_packets.py | 38 ++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) 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] From cbfbd29e40cbc9174d107642afc690cbbd9ed3e8 Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Wed, 7 Oct 2026 14:10:57 -0500 Subject: [PATCH 2/8] Add triggered SDMA plans (iris.triggered) The host pre-arms copy-engine chains that wait on gates; a kernel opens a gate once its data is written, and the copy engine copies it to every peer and signals its arrival there, without spending compute units. - TriggeredPlan arms epoch 1 when created, then runs begin_device_epoch / end_device_epoch / next_epoch per epoch, plus run, abort and destroy. Gates are 64-byte rows polled with >= on their low 32 bits and never reset. - Schedules list each rank's batches (ordered groups of 1D or 2D segments behind one gate) in the order its producer completes them. Builders: whole, chunks, waves. - Chains are armed through rocm-xio's host-initiated queue to each peer: POLL_REGMEM, one copy per segment, ATOMIC ADD_RTN_64 on the peer's arrival gate. All of these packets have run on gfx942/gfx950. - TriggeredView is the device side: publish (the last of a batch's contributors opens its gate), arrived, wait, reusable. - end_device_epoch checks completion with a kernel reading peers' arrival gates, since quiet() waits for everything on the queue. - Chains on one queue run in arming order, so plans sharing queues must begin and abort in arming order, and Iris.put/put_tile/ put_tiles/quiet raise while chains are armed on that peer's queue. - Plans need stable peer mappings: the torch allocator re-maps peer heaps on every allocation, so it is rejected; use vmem_chunked or vmem. Co-Authored-By: Claude Opus 5.5 (1M context) --- iris/__init__.py | 2 + iris/host/iris.py | 12 + iris/mem/triton/triggered.py | 134 ++++++++ iris/triggered/__init__.py | 15 + iris/triggered/armer.py | 72 +++++ iris/triggered/kernels.py | 59 ++++ iris/triggered/plan.py | 301 +++++++++++++++++ iris/triggered/schedule.py | 144 +++++++++ tests/unittests/test_triggered.py | 360 +++++++++++++++++++++ tests/unittests/test_triggered_schedule.py | 66 ++++ tests/unittests/test_triggered_view.py | 65 ++++ 11 files changed, 1230 insertions(+) create mode 100644 iris/mem/triton/triggered.py create mode 100644 iris/triggered/__init__.py create mode 100644 iris/triggered/armer.py create mode 100644 iris/triggered/kernels.py create mode 100644 iris/triggered/plan.py create mode 100644 iris/triggered/schedule.py create mode 100644 tests/unittests/test_triggered.py create mode 100644 tests/unittests/test_triggered_schedule.py create mode 100644 tests/unittests/test_triggered_view.py 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..7bd9c1467 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -160,6 +160,8 @@ def __init__(self, heap_size=1 << 30, allocator_type="torch", engines=()): self.copy_engines_device_ctx = None if self._sdma_ep is not None: self._init_copy_engines(num_gpus, num_ranks, cur_rank) + # Chains armed by iris.triggered plans on the host queues, per peer + self._triggered_ledger = None # Initialize CCL interface self.ccl = self.CCL(self) @@ -1113,6 +1115,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 +1206,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 +1273,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 +1333,19 @@ 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 + if self._triggered_ledger is not None: + self._triggered_ledger.check_idle(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/mem/triton/triggered.py b/iris/mem/triton/triggered.py new file mode 100644 index 000000000..257112671 --- /dev/null +++ b/iris/mem/triton/triggered.py @@ -0,0 +1,134 @@ +# 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).""" + 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..e37190adc --- /dev/null +++ b/iris/triggered/armer.py @@ -0,0 +1,72 @@ +# 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 xio import 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): + 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``.""" + 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..0b551b999 --- /dev/null +++ b/iris/triggered/kernels.py @@ -0,0 +1,59 @@ +# 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, 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 (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): + 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..ef295e89d --- /dev/null +++ b/iris/triggered/plan.py @@ -0,0 +1,301 @@ +# 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() + +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 +from collections import defaultdict + +import numpy as np +import torch + +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" + + +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) # peer -> [(plan, epoch, nbytes)] + + def reserve(self, plan, epoch, nbytes): + for peer, n in nbytes.items(): + used = sum(e[2] for e in self.entries[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[peer].append((plan, epoch, n)) + + def release(self, plan): + for peer in self.entries: + self.entries[peer] = [e for e in self.entries[peer] 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, peer): + if self.entries.get(peer): + raise RuntimeError( + f"triggered chains are armed on the copy-engine queue to rank {peer}; " + "complete, abort or destroy those plans before using it" + ) + + +def ledger(ctx): + if getattr(ctx, "_triggered_ledger", None) is None: + from xio import sdma_ep + + # Leave half the ring for wrap padding and work queued before the plan + ctx._triggered_ledger = _QueueLedger(sdma_ep.SDMA_QUEUE_SIZE // 2) + return ctx._triggered_ledger + + +def _digest(buffer_bytes, schedules): + text = repr((buffer_bytes, schedules)).encode() + return np.frombuffer(hashlib.sha1(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 + + # 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) + 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._ledger = ledger(ctx) + self.epoch = 0 + self._arm(1) + ctx.barrier() + + def batch_id(self, rank, index): + """Global id of batch ``index`` in ``schedules[rank]``.""" + return self._first[rank] + index + + def _arm(self, epoch): + if epoch > UINT32_MAX: + raise RuntimeError("epoch exceeds 32 bits; destroy the plan and create a new one") + peers = [p for p in range(self.world_size) if p != self.rank] + own = self.schedules[self.rank] + nbytes = sum(armer.chain_bytes(b) for b in own) + self._ledger.reserve(self, epoch, {p: nbytes for p in peers}) + self.epoch = epoch + self.state = State.ARMED + + 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 + try: + for i, batch in enumerate(own): + gate = self.gates.data_ptr() + self.batch_id(self.rank, i) * 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) + except Exception: + # Chains armed so far stay parked: abort() releases them + self.state = State.FAILED + raise + + 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._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.""" + 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): + """Forget the plan. Its chains must have drained: abort() an armed or failed plan first.""" + 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.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_triggered.py b/tests/unittests/test_triggered.py new file mode 100644 index 000000000..ed17e5ef2 --- /dev/null +++ b/tests/unittests/test_triggered.py @@ -0,0 +1,360 @@ +# 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 + +import pytest +import torch +import triton +import triton.language as tl + +import iris +from iris.mem.triton import triggered as dv +from iris.mem.utils import read_realtime +from iris.triggered import Batch, Segment, State, TriggeredPlan, TriggeredView, chunks, waves, whole + +EMPTY = -1 + + +def _ctx(): + ctx = iris.iris(1 << 26, allocator_type="vmem_chunked") + 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 + + +@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_torch_allocator_rejected(): + ctx = iris.iris(1 << 20) + 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] From eb35c4ee35ed130f33cefff8f46477de97937269 Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Wed, 7 Oct 2026 14:10:57 -0500 Subject: [PATCH 3/8] Document triggered SDMA A conceptual page on the model, the kernel author's ordering contract and the measured GPU-to-SDMA ordering, and an API reference page. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/conceptual/triggered-sdma.md | 114 ++++++++++++++++++++++++++++++ docs/index.md | 1 + docs/reference/host/triggered.md | 25 +++++++ docs/sphinx/_toc.yml | 2 + docs/sphinx/_toc.yml.in | 2 + 5 files changed, 144 insertions(+) create mode 100644 docs/conceptual/triggered-sdma.md create mode 100644 docs/reference/host/triggered.md diff --git a/docs/conceptual/triggered-sdma.md b/docs/conceptual/triggered-sdma.md new file mode 100644 index 000000000..938bad5a2 --- /dev/null +++ b/docs/conceptual/triggered-sdma.md @@ -0,0 +1,114 @@ +# 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 **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") +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() +``` + +## 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.** + +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. + +## 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. +- **One queue per peer.** Plans use the host-initiated copy-engine queue to each peer, shared with + host `Iris.put()`. 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 From 57b5c50616a3ed786c21b7b996ab262d0f1a2c70 Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Wed, 7 Oct 2026 15:49:23 -0500 Subject: [PATCH 4/8] Use SHA-256 for the plan's schedule digest Bandit flags SHA1 as a weak hash (B324, high severity), which fails the security scan. The digest only checks that every rank passed the same schedules, so any hash works. Co-Authored-By: Claude Opus 5.5 (1M context) --- iris/triggered/plan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/iris/triggered/plan.py b/iris/triggered/plan.py index ef295e89d..51196ee9c 100644 --- a/iris/triggered/plan.py +++ b/iris/triggered/plan.py @@ -97,7 +97,7 @@ def ledger(ctx): def _digest(buffer_bytes, schedules): text = repr((buffer_bytes, schedules)).encode() - return np.frombuffer(hashlib.sha1(text).digest()[:8], dtype=np.int64).copy() + return np.frombuffer(hashlib.sha256(text).digest()[:8], dtype=np.int64).copy() def view_words(rank, world_size, gates_ptr, schedules): From ddd2b20bbc54a2b3bb08ff0132967e25ae143fee Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Fri, 9 Oct 2026 11:49:45 -0500 Subject: [PATCH 5/8] Document that arrived/wait/reusable must be called uniformly per CTA Each is a scalar atomic: it runs on one lane, and the compiler hands the result to the other waves through LDS behind a CTA barrier. A polling loop that exits on a per-wave condition, such as a clock read, makes the waves leave after different numbers of iterations. The barriers then no longer pair up, so waves read the result before it is written (and read data before it arrives) or deadlock. Seen on MI300X with a consumer bounded by s_memrealtime; bounding the loop by a poll count fixes it. Stated in the kernel author's contract and in TriggeredView.arrived. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/conceptual/triggered-sdma.md | 4 ++++ iris/mem/triton/triggered.py | 7 ++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/docs/conceptual/triggered-sdma.md b/docs/conceptual/triggered-sdma.md index 938bad5a2..2978fcc75 100644 --- a/docs/conceptual/triggered-sdma.md +++ b/docs/conceptual/triggered-sdma.md @@ -59,6 +59,10 @@ plan.destroy() 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 diff --git a/iris/mem/triton/triggered.py b/iris/mem/triton/triggered.py index 257112671..b6acf6fff 100644 --- a/iris/mem/triton/triggered.py +++ b/iris/mem/triton/triggered.py @@ -116,7 +116,12 @@ def publish(self, batch): @triton.jit def arrived(self, batch): - """True once batch ``batch`` is visible here this epoch (published, for this rank's own).""" + """ + 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 From aaa31d2c15bc09f46d8a68b2d0991c0604002ac7 Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Fri, 9 Oct 2026 13:37:02 -0500 Subject: [PATCH 6/8] Refuse plans whose buffer or gates differ in heap offset across ranks The vmem_chunked heap frees an allocation when its tensor is garbage collected, independently on each rank. Frees processed in a different order (its free lists are LIFO), or a reference cycle collected on one rank before another, leave the ranks with different free lists, and the same allocation then lands at different offsets. Every peer address a plan arms is then wrong, silently: copies land elsewhere in the peer's heap while arrivals still look correct if the gates happen to match. Check both offsets across ranks when the plan is created. Marked FIXME: it works around the allocator and goes once allocation is symmetric. Co-Authored-By: Claude Opus 5.5 (1M context) --- iris/triggered/plan.py | 6 ++++++ tests/unittests/test_triggered.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/iris/triggered/plan.py b/iris/triggered/plan.py index 51196ee9c..ff754811b 100644 --- a/iris/triggered/plan.py +++ b/iris/triggered/plan.py @@ -166,6 +166,12 @@ def __init__(self, ctx, buffer, schedules, timeout=30.0): 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( diff --git a/tests/unittests/test_triggered.py b/tests/unittests/test_triggered.py index ed17e5ef2..f5af871ee 100644 --- a/tests/unittests/test_triggered.py +++ b/tests/unittests/test_triggered.py @@ -309,6 +309,21 @@ def test_misuse_raises(): 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 + + @triton.jit def _release_flag(flag, value): tl.atomic_xchg(flag, value, sem="release", scope="sys") From 41e47af6c0cba7521c8c59b25d4f62beae435c1d Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Fri, 16 Oct 2026 16:16:08 -0500 Subject: [PATCH 7/8] Arm triggered plans on every rank or none, and make destroy collective Address review on #573: - Plans reserve queue space and arm in two agreed steps, so a rank that cannot arm fails every rank instead of leaving its peers at a barrier. next_epoch() gets the same treatment. - A failed construction drains the chains it armed (and waits for peers to finish draining into this rank) before raising, since the caller has no plan to abort(). - The progress kernel waits only on chains that were actually armed, so aborting a partially armed plan no longer times out. - destroy() is collective: peers' chains and polls can still touch a rank's buffer and gates after its own work is done. - The queue ledger is process-wide rather than per Iris context, because every context's plans share the same host queue (channel 0). - Clarify in the docs how segments relate to producer tiles. Co-Authored-By: Claude Sonnet 5.5 --- docs/conceptual/triggered-sdma.md | 16 +++- iris/host/iris.py | 4 +- iris/host/queue_ledger.py | 64 +++++++++++++ iris/triggered/kernels.py | 10 ++- iris/triggered/plan.py | 143 ++++++++++++++++-------------- tests/unittests/test_triggered.py | 79 ++++++++++++++++- 6 files changed, 238 insertions(+), 78 deletions(-) create mode 100644 iris/host/queue_ledger.py diff --git a/docs/conceptual/triggered-sdma.md b/docs/conceptual/triggered-sdma.md index 2978fcc75..ccd44d586 100644 --- a/docs/conceptual/triggered-sdma.md +++ b/docs/conceptual/triggered-sdma.md @@ -18,7 +18,9 @@ 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. + 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 @@ -44,7 +46,7 @@ for step in range(steps): plan.end_device_epoch() # every batch arrived and drained ... plan.next_epoch() # barrier, re-arm, barrier -plan.destroy() +plan.destroy() # collective: peers may still touch this rank's memory ``` ## The kernel author's contract @@ -68,6 +70,12 @@ Plans that share queues must begin, and be aborted, in the order they were armed 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). @@ -110,8 +118,8 @@ after it in the same kernel is fresh (no stale elements over 9,600 tile-epochs). reserved range as the heap grows. The default torch allocator re-maps peer heaps on every allocation and is rejected. - **One queue per peer.** Plans use the host-initiated copy-engine queue to each peer, shared with - host `Iris.put()`. Kernel-initiated copies (`put(..., use_copy_engine=True)`) use a separate - queue and are unaffected. + 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 diff --git a/iris/host/iris.py b/iris/host/iris.py index 7bd9c1467..1194bfd98 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 @@ -1343,8 +1344,7 @@ def quiet(self, to_rank: int = None, channel: int = 0): def _check_host_queue_idle(self, to_rank): # Chains parked on a host queue block everything queued behind them, quiet() included - if self._triggered_ledger is not None: - self._triggered_ledger.check_idle(to_rank) + queue_ledger.ledger().check_idle(self.get_rank(), to_rank) def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Tensor]: """ diff --git a/iris/host/queue_ledger.py b/iris/host/queue_ledger.py new file mode 100644 index 000000000..c161a5823 --- /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 xio import 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(sdma_ep.SDMA_QUEUE_SIZE // 2) + return _ledger diff --git a/iris/triggered/kernels.py b/iris/triggered/kernels.py index 0b551b999..e06f32709 100644 --- a/iris/triggered/kernels.py +++ b/iris/triggered/kernels.py @@ -36,11 +36,13 @@ def _reached(ptr, epoch, budget: tl.constexpr): @triton.jit -def progress_kernel(view, heap_bases, status, BUDGET: tl.constexpr, WORLD_SIZE: tl.constexpr, OWN_ONLY: tl.constexpr): +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 (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. + 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) @@ -50,7 +52,7 @@ def progress_kernel(view, heap_bases, status, BUDGET: tl.constexpr, WORLD_SIZE: if own: done = _reached(gate + G_READY, tv.epoch, BUDGET) for peer in tl.static_range(WORLD_SIZE): - if done & (peer != tv.rank): + 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) diff --git a/iris/triggered/plan.py b/iris/triggered/plan.py index ff754811b..48d8565d7 100644 --- a/iris/triggered/plan.py +++ b/iris/triggered/plan.py @@ -11,7 +11,7 @@ producer[grid](..., view) # TriggeredView.publish(batch) plan.end_device_epoch() # every batch complete, or TimeoutError plan.next_epoch() # barrier -> re-arm -> barrier - plan.destroy() + 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 @@ -21,11 +21,11 @@ import enum import hashlib import time -from collections import defaultdict 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 @@ -45,56 +45,6 @@ class State(enum.Enum): DESTROYED = "destroyed" -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) # peer -> [(plan, epoch, nbytes)] - - def reserve(self, plan, epoch, nbytes): - for peer, n in nbytes.items(): - used = sum(e[2] for e in self.entries[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[peer].append((plan, epoch, n)) - - def release(self, plan): - for peer in self.entries: - self.entries[peer] = [e for e in self.entries[peer] 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, peer): - if self.entries.get(peer): - raise RuntimeError( - f"triggered chains are armed on the copy-engine queue to rank {peer}; " - "complete, abort or destroy those plans before using it" - ) - - -def ledger(ctx): - if getattr(ctx, "_triggered_ledger", None) is None: - from xio import sdma_ep - - # Leave half the ring for wrap padding and work queued before the plan - ctx._triggered_ledger = _QueueLedger(sdma_ep.SDMA_QUEUE_SIZE // 2) - return ctx._triggered_ledger - - def _digest(buffer_bytes, schedules): text = repr((buffer_bytes, schedules)).encode() return np.frombuffer(hashlib.sha256(text).digest()[:8], dtype=np.int64).copy() @@ -178,32 +128,84 @@ def __init__(self, ctx, buffer, schedules, timeout=30.0): 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._ledger = ledger(ctx) + self._armed = torch.zeros(self.num_batches * self.world_size, dtype=torch.int32, device=buffer.device) + self._ledger = queue_ledger.ledger() self.epoch = 0 - self._arm(1) + 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): - if epoch > UINT32_MAX: - raise RuntimeError("epoch exceeds 32 bits; destroy the plan and create a new one") + """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] - own = self.schedules[self.rank] - nbytes = sum(armer.chain_bytes(b) for b in own) - self._ledger.reserve(self, epoch, {p: nbytes for p in peers}) + 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(own): - gate = self.gates.data_ptr() + self.batch_id(self.rank, i) * gate_row + 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( @@ -217,10 +219,9 @@ def _arm(self, epoch): ] 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) - except Exception: - # Chains armed so far stay parked: abort() releases them - self.state = State.FAILED - raise + 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: @@ -243,6 +244,7 @@ def _wait(self, own_only, timeout): 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, @@ -269,7 +271,7 @@ def end_device_epoch(self, timeout=None): self.state = State.COMPLETE def next_epoch(self): - """Collective: re-arm every chain for the next epoch.""" + """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) @@ -299,9 +301,16 @@ def abort(self, timeout=None): self._ledger.release(self) def destroy(self): - """Forget the plan. Its chains must have drained: abort() an armed or failed plan first.""" + """ + 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/tests/unittests/test_triggered.py b/tests/unittests/test_triggered.py index f5af871ee..d1859b83d 100644 --- a/tests/unittests/test_triggered.py +++ b/tests/unittests/test_triggered.py @@ -8,6 +8,7 @@ import time from contextlib import contextmanager +from types import SimpleNamespace import pytest import torch @@ -15,9 +16,10 @@ 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, chunks, waves, whole +from iris.triggered import Batch, Segment, State, TriggeredPlan, TriggeredView, armer, chunks, waves, whole EMPTY = -1 @@ -324,6 +326,81 @@ def test_asymmetric_buffer_raises(): 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") From 1e8a64bbec533eef463d8780889cdf53cb30db5d Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Fri, 16 Oct 2026 17:02:48 -0500 Subject: [PATCH 8/8] Require the SDMA engine for triggered plans Upstream made SDMA opt-in (engines=("sdma",)) and rocm-xio an optional extra, so iris.triggered must not import xio at module load. Bind sdma_ep lazily through bind_sdma_ep(), and have TriggeredPlan raise the usual "SDMA is disabled" error for a context created without the engine. Co-Authored-By: Claude Sonnet 5.5 --- docs/conceptual/triggered-sdma.md | 3 ++- iris/host/iris.py | 2 -- iris/host/queue_ledger.py | 4 ++-- iris/triggered/armer.py | 4 +++- iris/triggered/plan.py | 1 + tests/unittests/test_triggered.py | 13 +++++++++++-- 6 files changed, 19 insertions(+), 8 deletions(-) diff --git a/docs/conceptual/triggered-sdma.md b/docs/conceptual/triggered-sdma.md index ccd44d586..6e89772da 100644 --- a/docs/conceptual/triggered-sdma.md +++ b/docs/conceptual/triggered-sdma.md @@ -35,7 +35,7 @@ per wave of a persistent tiled kernel, such as a GEMM). Any list of batches work ```python from iris.triggered import TriggeredPlan, TriggeredView, waves -ctx = iris.iris(heap_size, allocator_type="vmem_chunked") +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 @@ -117,6 +117,7 @@ after it in the same kernel is fresh (no stale elements over 9,600 tile-epochs). 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. diff --git a/iris/host/iris.py b/iris/host/iris.py index 1194bfd98..24d7d55b3 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -161,8 +161,6 @@ def __init__(self, heap_size=1 << 30, allocator_type="torch", engines=()): self.copy_engines_device_ctx = None if self._sdma_ep is not None: self._init_copy_engines(num_gpus, num_ranks, cur_rank) - # Chains armed by iris.triggered plans on the host queues, per peer - self._triggered_ledger = None # Initialize CCL interface self.ccl = self.CCL(self) diff --git a/iris/host/queue_ledger.py b/iris/host/queue_ledger.py index c161a5823..b765000f7 100644 --- a/iris/host/queue_ledger.py +++ b/iris/host/queue_ledger.py @@ -9,7 +9,7 @@ from collections import defaultdict -from xio import sdma_ep +from iris.device.sdma_utils import bind_sdma_ep class QueueLedger: @@ -60,5 +60,5 @@ def ledger(): global _ledger if _ledger is None: # Leave half the ring for wrap padding and work queued before the plan - _ledger = QueueLedger(sdma_ep.SDMA_QUEUE_SIZE // 2) + _ledger = QueueLedger(bind_sdma_ep().SDMA_QUEUE_SIZE // 2) return _ledger diff --git a/iris/triggered/armer.py b/iris/triggered/armer.py index e37190adc..5afbc0bd7 100644 --- a/iris/triggered/armer.py +++ b/iris/triggered/armer.py @@ -10,7 +10,7 @@ import ctypes from dataclasses import dataclass -from xio import sdma_ep +from iris.device.sdma_utils import bind_sdma_ep CHANNEL = 0 POLL_REGMEM_BYTES = 24 @@ -36,12 +36,14 @@ class Copy: 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: diff --git a/iris/triggered/plan.py b/iris/triggered/plan.py index 48d8565d7..764f1ea5a 100644 --- a/iris/triggered/plan.py +++ b/iris/triggered/plan.py @@ -93,6 +93,7 @@ def __init__(self, ctx, buffer, schedules, timeout=30.0): 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): diff --git a/tests/unittests/test_triggered.py b/tests/unittests/test_triggered.py index d1859b83d..51dd5377f 100644 --- a/tests/unittests/test_triggered.py +++ b/tests/unittests/test_triggered.py @@ -25,7 +25,7 @@ def _ctx(): - ctx = iris.iris(1 << 26, allocator_type="vmem_chunked") + 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 @@ -443,9 +443,18 @@ def test_host_put_wait_flag(): del ctx -def test_torch_allocator_rejected(): +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()