Triggered SDMA Plans

docs/conceptual/triggered-sdma.md

@@ -0,0 +1,118 @@
+# 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.**
+5. **`arrived`, `wait` and `reusable` are called uniformly across the CTA:** every wave makes the
+   same number of calls. Each is one atomic whose result reaches the other waves through a CTA
+   barrier, so a polling loop must not exit on a per-wave condition (such as a clock read):
+   waves would read the result before it is written, or deadlock. Count polls instead.
+
+Plans that share queues must begin, and be aborted, in the order they were armed. While chains
+are armed on a queue, host `Iris.put()` and `quiet()` to that peer raise rather than queue behind
+them.
+
+## 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.

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

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
+```

docs/sphinx/_toc.yml

@@ -10,6 +10,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
@@ -18,6 +19,7 @@ subtrees:
       - file: reference/host/class.md
       - file: reference/host/tensor-creation.md
       - file: reference/host/ccl.md
+      - file: reference/host/triggered.md
       - file: reference/triton/overview.md
         entries:
         - file: reference/triton/device-functions.md

docs/sphinx/_toc.yml.in

@@ -10,6 +10,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
@@ -18,6 +19,7 @@ subtrees:
       - file: reference/host/class.md
       - file: reference/host/tensor-creation.md
       - file: reference/host/ccl.md
+      - file: reference/host/triggered.md
       - file: reference/triton/overview.md
         entries:
         - file: reference/triton/device-functions.md

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",

iris/host/iris.py

@@ -167,6 +167,9 @@ def __init__(self, heap_size=1 << 30, allocator_type="torch"):
             self.copy_engines_device_ctx[local_rank][3] = handle.doorbell
             self.copy_engines_device_ctx[local_rank][4] = handle.cached_wptr
             self.copy_engines_device_ctx[local_rank][5] = handle.committed_wptr
+        # Chains armed by iris.triggered plans on the host queues, per peer
+        self._triggered_ledger = None
+
         # Initialize CCL interface
         self.ccl = self.CCL(self)
 
@@ -1038,6 +1041,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)
@@ -1125,6 +1129,7 @@ def put_tile(
             >>> shmem.put_tile(tile, to_rank=1, to_ptr=dst_ptr, to_stride=dst_stride,
             ...               wait_flag=wait_ptr, wait_value=256, signal_flag=signal_ptr)
         """
+        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)
@@ -1188,6 +1193,7 @@ def put_tiles(
             async_op: If True, don't wait for completion
             channel: SDMA channel to use
         """
+        self._check_host_queue_idle(to_rank)
         from_rank = self.get_rank()
 
         if len(tiles) != len(to_ptrs) or len(tiles) != len(to_strides):
@@ -1244,12 +1250,19 @@ def quiet(self, to_rank: int = None, channel: int = 0):
         """
         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.

iris/mem/triton/triggered.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+
+"""
+Device side of triggered SDMA plans (see ``iris.triggered``).
+
+The host arms one chain per (batch, peer) on the copy engine: poll the batch's ``ready`` gate,
+copy the batch's segments to the peer, then add 1 to the batch's ``arrival`` gate on the peer.
+A kernel opens the gate with ``publish`` and observes arrivals with ``arrived`` / ``wait``.
+
+View layout (int64 words): an 8-word header, then one 4-word record per batch. Gate table: one
+64-byte row per batch.
+"""
+
+import triton
+import triton.language as tl
+from triton.language.core import _aggregate as aggregate
+
+from iris.mem.utils import wait_cnt
+
+ABI_VERSION = 1
+
+# Header words
+HEADER_WORDS = tl.constexpr(8)
+H_VERSION = tl.constexpr(0)
+H_NUM_BATCHES = tl.constexpr(1)
+H_RANK = tl.constexpr(2)
+H_EPOCH = tl.constexpr(3)
+H_GATES = tl.constexpr(4)
+H_WORLD_SIZE = tl.constexpr(5)
+
+# Batch record words
+RECORD_WORDS = tl.constexpr(4)
+R_ROLE = tl.constexpr(0)
+R_CONTRIBUTORS = tl.constexpr(1)
+R_OWNER = tl.constexpr(2)
+
+ROLE_PUBLISH = tl.constexpr(1)
+ROLE_CONSUME = tl.constexpr(2)
+
+# Gate row words
+GATE_WORDS = tl.constexpr(8)
+G_READY = tl.constexpr(0)
+G_ARRIVAL = tl.constexpr(1)
+G_COMPLETION = tl.constexpr(2)
+G_ERROR = tl.constexpr(3)
+G_COUNTER = tl.constexpr(4)
+
+
+@aggregate
+class TriggeredView:
+    """
+    Device handle for one epoch of a triggered plan.
+
+    Usage::
+
+        @triton.jit
+        def producer(out, view, ...):
+            tv = TriggeredView.initialize(view)
+            ...                          # write this program's part of batch b
+            tv.publish(b)                # whole CTA, after its writes
+
+        @triton.jit
+        def consumer(out, view, ...):
+            tv = TriggeredView.initialize(view)
+            tv.wait(b)                   # batch b's data is now visible
+            x = tl.load(out + ...)
+    """
+
+    view: tl.tensor
+    gates: tl.tensor
+    epoch: tl.tensor
+    rank: tl.tensor
+
+    @triton.constexpr_function
+    def __init__(self, view, gates, epoch, rank):
+        self.view = view
+        self.gates = gates
+        self.epoch = epoch
+        self.rank = rank
+
+    @staticmethod
+    @triton.jit
+    def initialize(view):
+        """Read the plan's view tensor (from ``TriggeredPlan.begin_device_epoch()``)."""
+        gates = tl.cast(tl.load(view + H_GATES), tl.pointer_type(tl.int64))
+        return TriggeredView(view, gates, tl.load(view + H_EPOCH), tl.load(view + H_RANK))
+
+    @triton.jit
+    def _record(self, batch, word):
+        return tl.load(self.view + HEADER_WORDS + batch * RECORD_WORDS + word)
+
+    @triton.jit
+    def can_publish(self, batch):
+        return (self._record(batch, R_ROLE) & ROLE_PUBLISH) != 0
+
+    @triton.jit
+    def can_consume(self, batch):
+        return (self._record(batch, R_ROLE) & ROLE_CONSUME) != 0
+
+    @triton.jit
+    def publish(self, batch):
+        """
+        Contribute to batch ``batch``; the last of its contributors opens the gate.
+
+        Call from the whole CTA, outside divergent branches, after the CTA's writes to the batch.
+        """
+        # Copy engines read behind L2: every wave's stores must land before the release writes back
+        wait_cnt()
+        tl.debug_barrier()
+        gate = self.gates + batch * GATE_WORDS
+        contributors = self._record(batch, R_CONTRIBUTORS)
+        ticket = tl.atomic_add(gate + G_COUNTER, 1, sem="acq_rel", scope="sys")
+        if ticket % contributors == contributors - 1:
+            tl.atomic_xchg(gate + G_READY, self.epoch, sem="release", scope="sys")
+
+    @triton.jit
+    def arrived(self, batch):
+        """
+        True once batch ``batch`` is visible here this epoch (published, for this rank's own).
+
+        Call uniformly across the CTA: the result reaches the other waves through a CTA barrier, so
+        every wave must make the same number of calls (true of ``wait`` and ``reusable`` too).
+        """
+        gate = self.gates + batch * GATE_WORDS
+        word = tl.where(self._record(batch, R_OWNER) == self.rank, G_READY, G_ARRIVAL)
+        return tl.atomic_add(gate + word, 0, sem="acquire", scope="sys") >= self.epoch
+
+    @triton.jit
+    def wait(self, batch):
+        """Spin until ``arrived(batch)``."""
+        while self.arrived(batch) == 0:
+            pass
+
+    @triton.jit
+    def reusable(self, batch):
+        """True once batch ``batch`` is complete this epoch (set by ``end_device_epoch()``)."""
+        gate = self.gates + batch * GATE_WORDS
+        return tl.atomic_add(gate + G_COMPLETION, 0, sem="acquire", scope="sys") >= self.epoch