Copy Engine (SDMA)

Copy Engine (SDMA)#

Host-initiated transfers that run on the GPU’s SDMA copy engines rather than on compute units. The host enqueues a descriptor and the DMA engine moves the bytes, so no kernel launch is involved and no CUs are consumed.

Note

These operations require the copy engine to be initialised. See iris.iris() for how to enable it, and the Installation guide for the optional rocm-xio dependency they rely on.

Unlike the device-side operations, there is no engine choice to make here: a CU store is an instruction a wave executes, so the host cannot issue one. Anything host-initiated goes through the copy engine.

put#

Iris.put(from_tensor, to_rank, to_tensor=None, wait_flag=None, wait_value=None, signal_flag=None, signal_value=1, async_op=False, channel=0)[source]#

One-sided put operation with optional wait (POLL) and signal (ATOMIC).

Requires initialization with engines=("sdma",).

Supports:

  • Simple copy: put(src, to_rank)

  • Copy + signal: put(src, to_rank, signal_flag=flag)

  • Wait + copy: put(src, to_rank, wait_flag=flag, wait_value=N)

  • Wait + copy + signal: put(src, to_rank, wait_flag=…, signal_flag=…)

Parameters:

  • from_tensor (torch.Tensor) – Source tensor (local, must be symmetric)

  • to_rank (int) – Destination rank

  • to_tensor (torch.Tensor) – Destination tensor (symmetric). If None, uses from_tensor.

  • wait_flag (torch.Tensor) – Optional LOCAL flag tensor to poll before transfer (POLL packet)

  • wait_value (int) – Expected value for wait_flag

  • signal_flag (torch.Tensor) – Optional flag tensor to atomic-add on REMOTE rank after transfer (will be translated)

  • signal_value (int) – Value to add to signal_flag (default 1)

  • async_op (bool) – If True, don’t wait for completion

  • channel (int) – SDMA channel to use

Examples

>>> # Simple copy
>>> shmem.put(data, to_rank=1)

>>> # Copy with completion signal
>>> shmem.put(data, to_rank=1, signal_flag=completion_flag)

>>> # Wait for ready signal, then copy
>>> shmem.put(data, to_rank=1, wait_flag=ready_flag, wait_value=1)

>>> # Full pipeline: wait, copy, signal
>>> shmem.put(data, to_rank=1,
...          wait_flag=batch_ready, wait_value=256,
...          signal_flag=transfer_done, signal_value=1)

put_tile#

Iris.put_tile(tile, to_rank, to_ptr, to_stride, wait_flag=None, wait_value=None, signal_flag=None, signal_value=1, async_op=False, channel=0)[source]#

2D tile transfer with optional wait/signal (sub-window copy).

Requires initialization with engines=("sdma",).

Low-level API - caller provides pre-translated pointers for performance.

Parameters:

  • tile – Pre-configured sdma_ep.Tile object with data pointer and dimensions set

  • to_rank (int) – Destination rank

  • to_ptr (int) – Destination pointer (already translated to remote address space)

  • to_stride (int) – Destination row stride in bytes

  • wait_flag (int) – Optional LOCAL flag pointer to poll before transfer

  • wait_value (int) – Expected value for wait_flag

  • signal_flag (int) – Optional REMOTE flag pointer to atomic-add after transfer (already translated)

  • signal_value (int) – Value to add to signal_flag

  • async_op (bool) – If True, don’t wait for completion

  • channel (int) – SDMA channel to use

Examples

>>> from xio import sdma_ep
>>> tile = sdma_ep.Tile()
>>> tile.pid_m = 0
>>> tile.pid_n = 0
>>> tile.block_m = 256
>>> tile.block_n = 256
>>> tile.elem_size = A.element_size()
>>> tile.src_stride = A.stride(0) * tile.elem_size
>>> tile.data = A.data_ptr()
>>> dst_ptr = shmem.translate(A.data_ptr(), src_rank, dst_rank)
>>> dst_stride = A.stride(0) * tile.elem_size
>>> wait_ptr = flag.data_ptr()
>>> signal_ptr = shmem.translate(flag.data_ptr(), src_rank, dst_rank)
>>> 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)

put_tiles#

Iris.put_tiles(tiles, to_rank, to_ptrs, to_strides, wait_flag=None, wait_value=None, signal_flag=None, signal_value=1, async_op=False, channel=0)[source]#

Batched 2D tile transfer with optional shared wait/signal.

Requires initialization with engines=("sdma",).

Parameters:

  • tiles – Sequence of pre-configured sdma_ep.Tile objects

  • to_rank (int) – Destination rank

  • to_ptrs – Sequence of translated destination pointers

  • to_strides – Sequence of destination row strides in bytes

  • wait_flag (int) – Optional LOCAL flag pointer to poll before all transfers

  • wait_value (int) – Expected value for wait_flag

  • signal_flag (int) – Optional REMOTE flag pointer to atomic-add after all transfers

  • signal_value (int) – Value to add to signal_flag

  • async_op (bool) – If True, don’t wait for completion

  • channel (int) – SDMA channel to use

quiet#

Iris.quiet(to_rank=None, channel=0)[source]#

Wait for all outstanding SDMA operations to complete.

Requires initialization with engines=("sdma",).

Parameters:

  • to_rank (int) – If specified, wait only for ops to this rank. If None, wait for ops to all ranks.

  • channel (int) – SDMA channel

Example

>>> shmem.put(tensor, dst_rank=1, async_op=True)
>>> shmem.quiet(to_rank=1)  # Wait for completion
>>> shmem.quiet()  # Wait for all ranks

get_copy_engine_ctx#

Iris.get_copy_engine_ctx()[source]

Get the device-side copy engine context for SDMA operations.

The returned tensor holds the per-rank SDMA queue, read and write pointer, and doorbell addresses that device-side kernels need in order to enqueue copy packets. Pass it to iris.put as copy_engine_ctx when driving the copy engine from inside a kernel.

Returns:

The copy engine device context, or None when Iris was not initialized with engines=("sdma",).

Return type:

torch.Tensor

Contents