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(global_addr, to_rank, dst_addr, data_size, tile_height, tile_width, src_stride, dst_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:

  • global_addr (int) – Source address of the tile’s first element

  • to_rank (int) – Destination rank

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

  • data_size (int) – Element size in bytes

  • tile_height (int) – Number of rows

  • tile_width (int) – Number of elements per row

  • src_stride (int) – Source row stride in elements

  • dst_stride (int) – Destination row stride in elements

  • 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

>>> # Copy the top-left 256x256 tile of A to the same place on rank 1
>>> dst_addr = shmem.heap.translate(A.data_ptr(), src_rank, dst_rank)
>>> signal_ptr = shmem.heap.translate(flag.data_ptr(), src_rank, dst_rank)
>>> shmem.put_tile(A.data_ptr(), to_rank=1, dst_addr=dst_addr, data_size=A.element_size(),
...                tile_height=256, tile_width=256, src_stride=A.stride(0), dst_stride=A.stride(0),
...                wait_flag=flag.data_ptr(), wait_value=256, signal_flag=signal_ptr)

put_tiles#

Iris.put_tiles(global_addrs, to_rank, dst_addrs, data_size, tile_heights, tile_widths, src_strides, dst_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:

  • global_addrs – Source address of each tile’s first element

  • to_rank (int) – Destination rank

  • dst_addrs – Translated destination address of each tile

  • data_size (int) – Element size in bytes, shared by all tiles

  • tile_heights – Number of rows of each tile

  • tile_widths – Number of elements per row of each tile

  • src_strides – Source row stride of each tile, in elements

  • dst_strides – Destination row stride of each tile, in elements

  • 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