Triggered Plans

Experimental (iris.experimental.plan): may change between releases without deprecation.

Copy-engine transfers prepared once and released by kernels, every round. See Triggered Plans for the model and its ordering rules.

TriggeredPlan

class TriggeredPlan(ctx, buffer, schedule, timeout=30.0)

Bases: object

Sends regions of buffer from the rank that writes them to every other rank, using copy engines that a kernel releases chunk by chunk.

Each rank passes its own schedule: the chunks it writes, in the order its producer completes them. Every chunk is copied to the same offsets in each peer’s buffer. Chunk ids are global: rank r’s chunks come after those of ranks 0..r-1 (see chunk_id), and schedules[r] holds rank r’s schedule. A rank that writes nothing passes []. No byte may be in two ranks’ schedules.

Collective: every rank creates the plan for the same symmetric buffer. Use it as a context manager, or call close() on every rank when done.

Parameters:

• ctx – Iris instance, created with engines=("sdma",).

• buffer – Contiguous tensor on the symmetric heap.

• schedule – This rank’s list of Chunk, in the order it completes them.

• timeout – Seconds wait() and close() wait before raising.

__init__(ctx, buffer, schedule, timeout=30.0)

chunk_id(rank, index)

Global id of chunk index in rank rank’s schedule.

start()

Start the prepared round and return the handle tensor its kernels take.

wait(timeout=None)

Wait until every chunk is delivered this round: own ones at every peer, others here.

next_round()

Collective: prepare every chain for the next round. If any rank cannot, every rank raises.

Calling it declares that this rank is done reading the current round’s data: once it returns, a peer may start the next round, and its copies can overwrite buffer at any time. Work queued before the call is covered (the barrier synchronizes the device first); anything issued after it must not read buffer until the next wait().

run(timeout=None)

Run one round with no producer kernel: send every own chunk as it stands, then wait.

close(timeout=None)

Collective: wait until every rank is done with the plan, then free it.

Prepared copies cannot be withdrawn from the queue, so closing a plan whose round has not finished sends this rank’s chunks as they stand. A peer’s copies and polls may still touch this rank’s buffer and gates, so every rank must call close() before any of them is freed.

Plans may be closed in any order. If a plan prepared earlier on the same queues still has copies waiting, this plan’s copies are queued behind them and cannot run yet: close() then returns with the plan CLOSING, and closing the earlier plan finishes this one too.

Schedules

class Region(offset_bytes, row_bytes, rows=1, stride_bytes=None)

Bases: object

One copy: rows rows of row_bytes bytes, stride_bytes apart (packed if None), starting offset_bytes into the plan’s buffer. The same bytes are written at the same offset in each peer’s buffer.

Parameters:

• offset_bytes (int)

• row_bytes (int)

• rows (int)

• stride_bytes (int | None)

offset_bytes: int

row_bytes: int

rows: int = 1

stride_bytes: int | None = None

property row_stride

Bytes between the starts of consecutive rows (row_bytes if packed).

property end

Byte offset one past the last byte this region copies.

__init__(offset_bytes, row_bytes, rows=1, stride_bytes=None)

Parameters:

• offset_bytes (int)

• row_bytes (int)

• rows (int)

• stride_bytes (int | None)

Return type:

None

class Chunk(regions, writers=1)

Bases: object

Regions sent together once all writers have marked the chunk ready, every round.

Parameters:

• regions (tuple)

• writers (int)

regions: tuple

writers: int = 1

property is_2d

True if any region has more than one row.

__init__(regions, writers=1)

Parameters:

• regions (tuple)

• writers (int)

Return type:

None

whole(view, base=None, writers=1)

One chunk covering a contiguous view, marked ready by writers programs per round.

chunks(view, n, base=None, writers=1)

Split a contiguous view along dim 0 into n chunks, in order. Suits producers that finish row blocks in order: row-block, streaming and elementwise kernels.

tile_chunks(view, block_shape, tiles_per_chunk, order=None, base=None)

Group the tiles of a tiled producer into chunks of tiles_per_chunk, in completion order.

view is a 2D tensor with unit column stride, split into block_shape tiles (edge tiles clipped; tile ids are row-major). order gives the tile the producer makes at each step i: a list of tile ids, or a function of i. Chunk k holds the tiles of steps k * tiles_per_chunk up to the next chunk, one region per tile, and is marked ready once per tile. The producer must mark the tile of step i ready as part of chunk i // tiles_per_chunk, so pass the mapping the kernel itself uses: a @triton.jit mapping written with plain integer arithmetic runs on the host as lambda i: mapping.fn(i, ...).

For a persistent kernel of P programs where program p takes steps p, p + P, …, use tiles_per_chunk=P: chunk k is every program’s k-th tile, so chunks complete in order.

TriggeredView (device)

class TriggeredView(view, gates, round, rank)

Bases: object

Device handle for one round of a triggered plan.

Usage:

@triton.jit
def producer(out, handle, ...):
    cv = TriggeredView.initialize(handle)
    ...                          # write this program's part of chunk c
    cv.mark_ready(c)             # whole CTA, after its writes

@triton.jit
def consumer(out, handle, ...):
    cv = TriggeredView.initialize(handle)
    cv.wait(c)                   # chunk c's data is now visible
    x = tl.load(out + ...)

Parameters:

• view (tensor)

• gates (tensor)

• round (tensor)

• rank (tensor)

static initialize(handle)

Read the plan’s handle tensor (from TriggeredPlan.start()).

owns(chunk)

True if this rank produces chunk chunk.

mark_ready(chunk)

Declare this CTA’s writes to chunk chunk done; the last of its writers sends it.

Call from the whole CTA, outside divergent branches, after the CTA’s writes to the chunk.

arrived(chunk)

True once chunk chunk is visible here this round (marked ready, 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 too).

wait(chunk)

Spin until arrived(chunk).

TriggeredView (Gluon device)

class TriggeredView(view, gates, round, rank)

Bases: object

Gluon device handle for one round of a triggered plan.

Usage:

@gluon.jit
def producer(out, handle, ...):
    cv = TriggeredView.initialize(handle)
    ...                          # write this program's part of chunk c
    cv.mark_ready(c)             # whole CTA, after its writes

@gluon.jit
def consumer(out, handle, ...):
    cv = TriggeredView.initialize(handle)
    cv.wait(c)                   # chunk c's data is now visible
    x = gl.load(out + ...)

Parameters:

• view (tensor)

• gates (tensor)

• round (tensor)

• rank (tensor)

static initialize(handle)

Read the plan’s handle tensor (from TriggeredPlan.start()).

owns(chunk)

True if this rank produces chunk chunk.

mark_ready(chunk)

Declare this CTA’s writes to chunk chunk done; the last of its writers sends it.

Call from the whole CTA, after the CTA’s writes to the chunk: outside divergent branches and outside gl.warp_specialize partitions.

arrived(chunk)

True once chunk chunk is visible here this round (marked ready, 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 too).

wait(chunk)

Spin until arrived(chunk).