diff --git a/iris/host/iris.py b/iris/host/iris.py index a399c7e65..fb56083f1 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -1133,32 +1133,36 @@ def put( wait_val = int(wait_value if wait_value is not None else 0) signal_val = int(signal_value) sdma_ep.wait_flag_then_put( - from_rank, to_rank, channel, wait_ptr, wait_val, from_ptr, to_ptr, size, wait_bits + from_rank, to_rank, channel, wait_ptr, wait_val, to_ptr, from_ptr, size, wait_bits ) sdma_ep.signal(from_rank, to_rank, channel, signal_ptr, signal_val, signal_bits) elif has_wait: # Wait + copy wait_val = int(wait_value if wait_value is not None else 0) sdma_ep.wait_flag_then_put( - from_rank, to_rank, channel, wait_ptr, wait_val, from_ptr, to_ptr, size, wait_bits + from_rank, to_rank, channel, wait_ptr, wait_val, to_ptr, from_ptr, size, wait_bits ) elif has_signal: # Copy + signal signal_val = int(signal_value) - sdma_ep.put_signal(from_rank, to_rank, channel, from_ptr, to_ptr, size, signal_ptr, signal_val, signal_bits) + sdma_ep.put_signal(from_rank, to_rank, channel, to_ptr, from_ptr, size, signal_ptr, signal_val, signal_bits) else: # Simple copy - sdma_ep.put(from_rank, to_rank, channel, from_ptr, to_ptr, size) + sdma_ep.put(from_rank, to_rank, channel, to_ptr, from_ptr, size) if not async_op: sdma_ep.quiet(from_rank, to_rank, channel) def put_tile( self, - tile, + global_addr: int, to_rank: int, - to_ptr: int, - to_stride: int, + dst_addr: int, + data_size: int, + tile_height: int, + tile_width: int, + src_stride: int, + dst_stride: int, wait_flag: int = None, wait_value: int = None, signal_flag: int = None, @@ -1174,10 +1178,14 @@ def put_tile( Low-level API - caller provides pre-translated pointers for performance. Args: - tile: Pre-configured sdma_ep.Tile object with data pointer and dimensions set + global_addr: Source address at the first element of the tile to_rank: Destination rank - to_ptr: Destination pointer (already translated to remote address space) - to_stride: Destination row stride in bytes + dst_addr: Destination address (already translated to remote address space) + data_size: Element size in bytes + tile_height: Number of rows + tile_width: Number of elements per row + src_stride: Source row stride in elements + dst_stride: Destination row stride in elements wait_flag: Optional LOCAL flag pointer to poll before transfer wait_value: Expected value for wait_flag signal_flag: Optional REMOTE flag pointer to atomic-add after transfer (already translated) @@ -1185,22 +1193,6 @@ def put_tile( async_op: If True, don't wait for completion channel: 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) """ sdma_ep = self._require_copy_engine() from_rank = self.get_rank() @@ -1216,34 +1208,85 @@ def put_tile( wait_val = int(wait_value if wait_value is not None else 0) signal_val = int(signal_value) sdma_ep.wait_flag_then_put_tile( - from_rank, to_rank, channel, wait_ptr, wait_val, tile, int(to_ptr), int(to_stride), wait_bits + from_rank, + to_rank, + channel, + wait_ptr, + wait_val, + int(dst_addr), + int(global_addr), + int(data_size), + int(tile_height), + int(tile_width), + int(dst_stride), + int(src_stride), + wait_bits, ) sdma_ep.signal(from_rank, to_rank, channel, signal_ptr, signal_val, signal_bits) elif has_wait: # Wait + tile copy wait_val = int(wait_value if wait_value is not None else 0) sdma_ep.wait_flag_then_put_tile( - from_rank, to_rank, channel, wait_ptr, wait_val, tile, int(to_ptr), int(to_stride), wait_bits + from_rank, + to_rank, + channel, + wait_ptr, + wait_val, + int(dst_addr), + int(global_addr), + int(data_size), + int(tile_height), + int(tile_width), + int(dst_stride), + int(src_stride), + wait_bits, ) elif has_signal: # Tile copy + signal signal_val = int(signal_value) sdma_ep.put_tile_signal( - from_rank, to_rank, channel, tile, int(to_ptr), int(to_stride), signal_ptr, signal_val, signal_bits + from_rank, + to_rank, + channel, + int(dst_addr), + int(global_addr), + int(data_size), + int(tile_height), + int(tile_width), + int(dst_stride), + int(src_stride), + signal_ptr, + signal_val, + signal_bits, ) else: # Simple tile copy - sdma_ep.put_tile(from_rank, to_rank, channel, tile, int(to_ptr), int(to_stride)) + sdma_ep.put_tile( + from_rank, + to_rank, + channel, + int(dst_addr), + int(global_addr), + int(data_size), + int(tile_height), + int(tile_width), + int(dst_stride), + int(src_stride), + ) if not async_op: sdma_ep.quiet(from_rank, to_rank, channel) def put_tiles( self, - tiles, + global_addrs, to_rank: int, - to_ptrs, - to_strides, + dst_addrs, + data_size: int, + tile_heights, + tile_widths, + src_strides, + dst_strides, wait_flag: int = None, wait_value: int = None, signal_flag: int = None, @@ -1257,10 +1300,14 @@ def put_tiles( Requires initialization with ``engines=("sdma",)``. Args: - tiles: Sequence of pre-configured sdma_ep.Tile objects + global_addrs: Source addresses at the first element of each tile to_rank: Destination rank - to_ptrs: Sequence of translated destination pointers - to_strides: Sequence of destination row strides in bytes + dst_addrs: Sequence of translated destination pointers + data_size: Element size in bytes + tile_heights: Number of rows for each tile + tile_widths: Number of elements per row for each tile + src_strides: Source row strides in elements + dst_strides: Destination row strides in elements wait_flag: Optional LOCAL flag pointer to poll before all transfers wait_value: Expected value for wait_flag signal_flag: Optional REMOTE flag pointer to atomic-add after all transfers @@ -1271,8 +1318,15 @@ def put_tiles( sdma_ep = self._require_copy_engine() from_rank = self.get_rank() - if len(tiles) != len(to_ptrs) or len(tiles) != len(to_strides): - raise ValueError("tiles, to_ptrs, and to_strides must have the same length") + if not ( + len(global_addrs) + == len(dst_addrs) + == len(tile_heights) + == len(tile_widths) + == len(src_strides) + == len(dst_strides) + ): + raise ValueError("tile transfer arrays must have the same length") wait_ptr, wait_bits = self._flag_pointer_and_bits(wait_flag, default_bits=32) signal_ptr, signal_bits = self._flag_pointer_and_bits(signal_flag, default_bits=32) @@ -1280,31 +1334,81 @@ def put_tiles( has_wait = wait_ptr != 0 has_signal = signal_ptr != 0 - to_ptr_list = [int(p) for p in to_ptrs] - to_stride_list = [int(s) for s in to_strides] + src_ptr_list = [int(p) for p in global_addrs] + to_ptr_list = [int(p) for p in dst_addrs] + tile_height_list = [int(v) for v in tile_heights] + tile_width_list = [int(v) for v in tile_widths] + src_stride_list = [int(v) for v in src_strides] + dst_stride_list = [int(v) for v in dst_strides] if has_wait and has_signal: # Wait + tiles copy + signal (two calls) wait_val = int(wait_value if wait_value is not None else 0) signal_val = int(signal_value) sdma_ep.wait_flag_then_put_tiles( - from_rank, to_rank, channel, wait_ptr, wait_val, list(tiles), to_ptr_list, to_stride_list, wait_bits + from_rank, + to_rank, + channel, + wait_ptr, + wait_val, + to_ptr_list, + src_ptr_list, + int(data_size), + tile_height_list, + tile_width_list, + dst_stride_list, + src_stride_list, + wait_bits, ) sdma_ep.signal(from_rank, to_rank, channel, signal_ptr, signal_val, signal_bits) elif has_wait: # Wait + tiles copy wait_val = int(wait_value if wait_value is not None else 0) sdma_ep.wait_flag_then_put_tiles( - from_rank, to_rank, channel, wait_ptr, wait_val, list(tiles), to_ptr_list, to_stride_list, wait_bits + from_rank, + to_rank, + channel, + wait_ptr, + wait_val, + to_ptr_list, + src_ptr_list, + int(data_size), + tile_height_list, + tile_width_list, + dst_stride_list, + src_stride_list, + wait_bits, ) elif has_signal: # Tiles copy + signal (loop + signal) signal_val = int(signal_value) - sdma_ep.put_tiles(from_rank, to_rank, channel, list(tiles), to_ptr_list, to_stride_list) + sdma_ep.put_tiles( + from_rank, + to_rank, + channel, + to_ptr_list, + src_ptr_list, + int(data_size), + tile_height_list, + tile_width_list, + dst_stride_list, + src_stride_list, + ) sdma_ep.signal(from_rank, to_rank, channel, signal_ptr, signal_val, signal_bits) else: # Simple tiles copy - sdma_ep.put_tiles(from_rank, to_rank, channel, list(tiles), to_ptr_list, to_stride_list) + sdma_ep.put_tiles( + from_rank, + to_rank, + channel, + to_ptr_list, + src_ptr_list, + int(data_size), + tile_height_list, + tile_width_list, + dst_stride_list, + src_stride_list, + ) if not async_op: sdma_ep.quiet(from_rank, to_rank, channel) diff --git a/pyproject.toml b/pyproject.toml index cc8c4e31d..6ec825c0b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ amd = [ ] sdma = [ # SDMA copy engine (use_copy_engine=True, host-initiated put/quiet). Install via `pip install iris[sdma]`. - "rocm-xio @ git+https://github.com/ROCm/rocm-xio.git@cbe97e6392066bef7901121965ffadad19404da4", + "rocm-xio @ git+https://github.com/ROCm/rocm-xio.git@5c3bc937bbe9d02026d7431c8c2bbd72cef4f5a7", ] dev = [ "pytest", diff --git a/tests/unittests/test_sdma_host_convention.py b/tests/unittests/test_sdma_host_convention.py new file mode 100644 index 000000000..c2aa1eddf --- /dev/null +++ b/tests/unittests/test_sdma_host_convention.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +# Call convention of the host put methods against a recording stand-in for ``xio.sdma_ep`` (host only). +# rocm-xio 5c3bc93 takes ``(dst, src)`` and raw tile arguments; a swapped pair is silent on hardware, so +# these check that the destination is passed first and the source second. + +from types import SimpleNamespace + +import pytest + +from iris import Iris + +SRC, DST = 0x1000, 0x9000 +REMOTE = 0x4000_0000 # what the fake heap adds when translating to a peer + + +class FakeSdmaEp: + def __init__(self): + self.calls = [] + + def __getattr__(self, name): + def record(*args): + self.calls.append((name, args)) + + return record + + +class FakeTensor: + def __init__(self, ptr, nbytes): + self._ptr, self._nbytes = ptr, nbytes + + def data_ptr(self): + return self._ptr + + def numel(self): + return self._nbytes + + def element_size(self): + return 1 + + +def make_ctx(): + ep = FakeSdmaEp() + ctx = SimpleNamespace( + get_rank=lambda: 0, + _require_copy_engine=lambda: ep, + heap=SimpleNamespace(translate=lambda ptr, src_rank, dst_rank: ptr + REMOTE), + ) + ctx._flag_pointer_and_bits = lambda flag, **kw: Iris._flag_pointer_and_bits(ctx, flag, **kw) + return ctx, ep + + +def test_put_passes_dst_before_src(): + ctx, ep = make_ctx() + Iris.put(ctx, FakeTensor(SRC, 256), to_rank=1) + assert ep.calls == [("put", (0, 1, 0, SRC + REMOTE, SRC, 256)), ("quiet", (0, 1, 0))] + + +def test_put_with_wait_and_signal_passes_dst_before_src(): + ctx, ep = make_ctx() + Iris.put(ctx, FakeTensor(SRC, 256), to_rank=1, wait_flag=0xA000, wait_value=3, signal_flag=0xB000, async_op=True) + assert ep.calls == [ + ("wait_flag_then_put", (0, 1, 0, 0xA000, 3, SRC + REMOTE, SRC, 256, 32)), + ("signal", (0, 1, 0, 0xB000 + REMOTE, 1, 32)), + ] + + +def test_put_signal_passes_dst_before_src(): + ctx, ep = make_ctx() + Iris.put(ctx, FakeTensor(SRC, 256), to_rank=1, signal_flag=0xB000, async_op=True) + assert ep.calls == [("put_signal", (0, 1, 0, SRC + REMOTE, SRC, 256, 0xB000 + REMOTE, 1, 32))] + + +def test_put_tile_takes_raw_arguments(): + ctx, ep = make_ctx() + Iris.put_tile(ctx, SRC, 1, DST, 4, 8, 16, 64, 128, async_op=True) + assert ep.calls == [("put_tile", (0, 1, 0, DST, SRC, 4, 8, 16, 128, 64))] # ..., dst_stride, src_stride + + +def test_put_tile_with_wait_takes_raw_arguments(): + ctx, ep = make_ctx() + Iris.put_tile(ctx, SRC, 1, DST, 4, 8, 16, 64, 128, wait_flag=0xA000, wait_value=2, async_op=True) + assert ep.calls == [("wait_flag_then_put_tile", (0, 1, 0, 0xA000, 2, DST, SRC, 4, 8, 16, 128, 64, 32))] + + +def test_put_tiles_takes_raw_arguments(): + ctx, ep = make_ctx() + Iris.put_tiles(ctx, [SRC, SRC + 64], 1, [DST, DST + 64], 4, [8, 4], [16, 8], [64, 64], [128, 128], async_op=True) + assert ep.calls == [ + ( + "put_tiles", + (0, 1, 0, [DST, DST + 64], [SRC, SRC + 64], 4, [8, 4], [16, 8], [128, 128], [64, 64]), + ) + ] + + +def test_put_tiles_rejects_mismatched_lengths(): + ctx, ep = make_ctx() + with pytest.raises(ValueError, match="same length"): + Iris.put_tiles(ctx, [SRC], 1, [DST, DST + 64], 4, [8], [16], [64], [128]) + assert ep.calls == []