diff --git a/.github/scripts/run_new_examples.sh b/.github/scripts/run_new_examples.sh index 9767f26c4..b79736833 100755 --- a/.github/scripts/run_new_examples.sh +++ b/.github/scripts/run_new_examples.sh @@ -24,11 +24,11 @@ fi if [ "$INSTALL_METHOD" = "git" ]; then REPO=${GITHUB_REPOSITORY:-"ROCm/iris"} SHA=${GITHUB_SHA:-"HEAD"} - INSTALL_CMD="pip install git+https://github.com/${REPO}.git@${SHA}" + INSTALL_CMD="pip install 'iris[sdma] @ git+https://github.com/${REPO}.git@${SHA}'" elif [ "$INSTALL_METHOD" = "editable" ]; then - INSTALL_CMD="pip install -e ." + INSTALL_CMD="pip install -e '.[sdma]'" elif [ "$INSTALL_METHOD" = "install" ]; then - INSTALL_CMD="pip install ." + INSTALL_CMD="pip install '.[sdma]'" else echo "[ERROR] Invalid install_method: $INSTALL_METHOD" exit 1 diff --git a/.github/scripts/run_perf_benchmark.sh b/.github/scripts/run_perf_benchmark.sh index 26abd6254..5db9ba575 100755 --- a/.github/scripts/run_perf_benchmark.sh +++ b/.github/scripts/run_perf_benchmark.sh @@ -29,7 +29,7 @@ echo "[PERF-BENCHMARK] Using GPUs: $GPU_DEVICES" set -e cd /iris_workspace - pip install -e . + pip install -e '.[sdma]' torchrun --rdzv-backend=c10d --rdzv-endpoint=localhost:0 --nnodes=1 --nproc_per_node=8 examples/${EXAMPLE_PATH}/benchmark.py \ --benchmark \ --validate \ diff --git a/.github/scripts/run_tests.sh b/.github/scripts/run_tests.sh index f5ad694a8..dcd84c268 100755 --- a/.github/scripts/run_tests.sh +++ b/.github/scripts/run_tests.sh @@ -8,9 +8,9 @@ # num_ranks: number of GPU ranks (1, 2, 4, or 8) # gpu_devices: comma-separated GPU device IDs (optional, if not provided will use allocator) # install_method: pip install method - "git", "editable", or "install" (optional, default: "editable") -# - "git": pip install git+https://github.com/${{ github.repository }}.git@${{ github.sha }} -# - "editable": pip install -e . -# - "install": pip install . +# - "git": pip install 'iris[sdma] @ git+https://github.com/${{ github.repository }}.git@${{ github.sha }}' +# - "editable": pip install -e '.[sdma]' +# - "install": pip install '.[sdma]' # # CONTAINER_IMAGE, if set, selects the container image (e.g. the nightly Triton # image) instead of the one container_build.sh produced. @@ -68,11 +68,11 @@ if [ "$INSTALL_METHOD" = "git" ]; then # For git install, we need the repository and SHA from environment or use defaults REPO=${GITHUB_REPOSITORY:-"ROCm/iris"} SHA=${GITHUB_SHA:-"HEAD"} - INSTALL_CMD="pip install git+https://github.com/${REPO}.git@${SHA}" + INSTALL_CMD="pip install 'iris[sdma] @ git+https://github.com/${REPO}.git@${SHA}'" elif [ "$INSTALL_METHOD" = "editable" ]; then - INSTALL_CMD="pip install -e ." + INSTALL_CMD="pip install -e '.[sdma]'" elif [ "$INSTALL_METHOD" = "install" ]; then - INSTALL_CMD="pip install ." + INSTALL_CMD="pip install '.[sdma]'" fi # Run tests in container diff --git a/.github/workflows/copilot-setup-steps.yml b/.github/workflows/copilot-setup-steps.yml index 4a2582ae3..41c67d21e 100644 --- a/.github/workflows/copilot-setup-steps.yml +++ b/.github/workflows/copilot-setup-steps.yml @@ -30,7 +30,7 @@ jobs: /opt/venv/bin/python -m venv --system-site-packages .venv source .venv/bin/activate python -m pip install --upgrade pip - python -m pip install -e . --no-cache-dir + python -m pip install -e ".[sdma]" --no-cache-dir - name: Install native build dependencies run: | diff --git a/README.md b/README.md index ee2ba71d0..3dc623343 100644 --- a/README.md +++ b/README.md @@ -203,6 +203,12 @@ For a quick installation directly from the repository: pip install git+https://github.com/ROCm/iris.git ``` +The SDMA copy engine needs the optional `sdma` extra ([rocm-xio](https://github.com/ROCm/rocm-xio)): + +```shell +pip install "iris[sdma] @ git+https://github.com/ROCm/iris.git" +``` + ### Docker Compose (Recommended for Development) The recommended way to get started is using Docker Compose, which provides a development environment with the Iris directory mounted inside the container. This allows you to make changes to the code outside the container and see them reflected inside. diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index dcc9588f2..99e851b02 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -94,6 +94,26 @@ pip install -e . ``` +## Optional: SDMA copy engine + +Shader load/store, atomics, and collectives work with the base install. The SDMA +copy engine (`use_copy_engine=True` in device code, and host-initiated `put`, +`put_tile`, `put_tiles`, and `quiet`) needs [rocm-xio](https://github.com/ROCm/rocm-xio), +which is an optional extra: + +```shell +pip install "iris[sdma] @ git+https://github.com/ROCm/iris.git" +# or, from a checkout +pip install -e ".[sdma]" +``` + +To add SDMA to an existing Iris install, rerun the install with the extra. If an +older rocm-xio is already installed, run `pip uninstall rocm-xio` first: every +rocm-xio commit has the same version number, so pip keeps the installed one. + +Installing rocm-xio does not turn SDMA on; see [Iris Class](../reference/host/class.md) +for enabling it. + ## Selecting the AMD runtime for fabric communication When multiple ROCm versions are installed, set `LD_LIBRARY_PATH` before starting diff --git a/docs/reference/host/class.md b/docs/reference/host/class.md index 87a0187bb..46845d8b8 100644 --- a/docs/reference/host/class.md +++ b/docs/reference/host/class.md @@ -8,6 +8,29 @@ Prefer using the convenience factory over calling the constructor directly: .. autofunction:: iris.host.iris.iris ``` +## Optional SDMA Copy Engine + +The `engines` parameter defaults to `()`. In this mode, Iris skips SDMA +initialization and queue creation, and `get_copy_engine_ctx()` returns `None`. +Regular Triton and Gluon memory operations and collectives work without SDMA. + +Pass `engines=("sdma",)` to initialize both host- and device-initiated SDMA queues. +This needs the `sdma` extra ([Installation](../../getting-started/installation.md)): + +```python +ctx = iris.iris(heap_size=1 << 30, engines=("sdma",)) +copy_engine_ctx = ctx.get_copy_engine_ctx() +``` + +This is required for the host methods `put`, `put_tile`, `put_tiles`, and `quiet`, +and for device operations using `use_copy_engine=True`. Host SDMA methods raise +`RuntimeError` if the copy engine is disabled. Pass the returned context to +device operations through their `copy_engine_ctx` argument. + +```{eval-rst} +.. automethod:: iris.host.iris.Iris.get_copy_engine_ctx +``` + ## Core Methods ```{eval-rst} diff --git a/examples/06_message_passing/message_passing_host_initiated.py b/examples/06_message_passing/message_passing_host_initiated.py index 51b59fd5b..dfebff220 100644 --- a/examples/06_message_passing/message_passing_host_initiated.py +++ b/examples/06_message_passing/message_passing_host_initiated.py @@ -174,7 +174,7 @@ def _worker(local_rank: int, world_size: int, init_url: str, args: dict): ) # Main benchmark logic - shmem = iris.iris(args["heap_size"]) + shmem = iris.iris(args["heap_size"], engines=("sdma",)) dtype = torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() diff --git a/examples/06_message_passing/message_passing_put.py b/examples/06_message_passing/message_passing_put.py index 6cad6928f..f27ad7028 100755 --- a/examples/06_message_passing/message_passing_put.py +++ b/examples/06_message_passing/message_passing_put.py @@ -157,7 +157,7 @@ def _worker(local_rank: int, world_size: int, init_url: str, args: dict): ) # Main benchmark logic - shmem = iris.iris(args["heap_size"]) + shmem = iris.iris(args["heap_size"], engines=("sdma",) if args["use_copy_engine"] else ()) dtype = torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() diff --git a/iris/device/sdma_utils.py b/iris/device/sdma_utils.py index b50335234..baf7b170d 100644 --- a/iris/device/sdma_utils.py +++ b/iris/device/sdma_utils.py @@ -11,9 +11,29 @@ import triton import triton.language as tl -from xio import sdma_ep from iris.mem.utils import wait_cnt +sdma_ep = None + + +def bind_sdma_ep(): + """Import rocm-xio and bind ``sdma_ep`` for the SDMA device functions.""" + global sdma_ep + if sdma_ep is None: + try: + from xio import sdma_ep as imported + except ModuleNotFoundError as exc: + if exc.name != "xio": + raise + raise ModuleNotFoundError( + "rocm-xio is required for the SDMA copy engine. " + 'Install it with `pip install iris[sdma]`, or drop "sdma" from engines ' + "to use shader load/store only.", + name="xio", + ) from exc + sdma_ep = imported + return sdma_ep + @triton.jit def wrap_into_ring(index: tl.uint64): diff --git a/iris/host/iris.py b/iris/host/iris.py index aa1674334..49d4a9913 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -53,7 +53,6 @@ count_devices, ) -from xio import sdma_ep from iris.host.memory.symmetric_heap import SymmetricHeap import numpy as np from typing import Any @@ -72,6 +71,8 @@ from iris.host.memory import tensors as tensor_creation from iris.host.platform.utils import is_simulation_env +_SUPPORTED_ENGINES = frozenset({"sdma"}) + class Iris: """ @@ -83,6 +84,9 @@ class Iris: Args: heap_size (int): Size of the symmetric heap in bytes. Default: 1GB (2^30) allocator_type (str): Type of allocator to use. Options: "torch" (default), "vmem" + engines (str or tuple of str, optional): Engine name or names to set up. Use + ``("sdma",)`` for host- and device-initiated SDMA operations (needs + ``pip install iris[sdma]``). Defaults to ``()`` (none). Example: >>> ctx = iris.iris(heap_size=2**31) # 2GB heap with torch allocator @@ -91,9 +95,25 @@ class Iris: >>> # Use VMem allocator for memory oversubscription >>> ctx = iris.iris(heap_size=2**31, allocator_type="vmem") + """ - def __init__(self, heap_size=1 << 30, allocator_type="torch"): + def __init__(self, heap_size=1 << 30, allocator_type="torch", engines=()): + if engines is None: + engines = () + elif isinstance(engines, str): + engines = (engines,) + else: + engines = tuple(engines) + unsupported = set(engines) - _SUPPORTED_ENGINES + if unsupported: + raise ValueError(f"Unsupported engines: {sorted(unsupported)}. Supported: {sorted(_SUPPORTED_ENGINES)}.") + self._sdma_ep = None + if "sdma" in engines: + from iris.device.sdma_utils import bind_sdma_ep + + self._sdma_ep = bind_sdma_ep() + # Initialize distributed environment comm, cur_rank, num_ranks = init_distributed() num_gpus = count_devices() @@ -137,7 +157,27 @@ def __init__(self, heap_size=1 << 30, allocator_type="torch"): distributed_barrier() - # initialize copy engines + self.copy_engines_device_ctx = None + if self._sdma_ep is not None: + self._init_copy_engines(num_gpus, num_ranks, cur_rank) + + # Initialize CCL interface + self.ccl = self.CCL(self) + + # Lazy initialization for ops interface + self._ops = None + + # Device-side barrier state, keyed by process group (None = all ranks). + self._device_barrier_state: dict[Any, torch.Tensor] = {} + + # Initialize tracing + self.tracing = Tracing(self) + + # Pre-build the device context tensor (rebuilt when tracing is enabled) + self._build_device_context() + + def _init_copy_engines(self, num_gpus, num_ranks, cur_rank): + sdma_ep = self._sdma_ep sdma_ep.init() context_size = sdma_ep.QUEUE_DEVICE_CTX_SIZE @@ -167,20 +207,13 @@ 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 - # Initialize CCL interface - self.ccl = self.CCL(self) - - # Lazy initialization for ops interface - self._ops = None - # Device-side barrier state, keyed by process group (None = all ranks). - self._device_barrier_state: dict[Any, torch.Tensor] = {} - - # Initialize tracing - self.tracing = Tracing(self) - - # Pre-build the device context tensor (rebuilt when tracing is enabled) - self._build_device_context() + def _require_copy_engine(self): + if self._sdma_ep is None: + raise RuntimeError( + 'SDMA is disabled; initialize Iris with engines=("sdma",) (needs `pip install iris[sdma]`).' + ) + return self._sdma_ep def __del__(self): """Cleanup resources on deletion.""" @@ -950,7 +983,8 @@ def get_heap_bases(self): """ return self.heap_bases - def get_copy_engine_ctx(self): + def get_copy_engine_ctx(self) -> torch.Tensor | None: + """Return the device copy-engine context, or None when copy engines are disabled.""" return self.copy_engines_device_ctx @staticmethod @@ -1003,6 +1037,8 @@ def put( """ 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) @@ -1035,6 +1071,7 @@ def put( ... wait_flag=batch_ready, wait_value=256, ... signal_flag=transfer_done, signal_value=1) """ + sdma_ep = self._require_copy_engine() if to_tensor is None: to_tensor = from_tensor @@ -1094,6 +1131,8 @@ def put_tile( """ 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. Args: @@ -1125,6 +1164,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) """ + sdma_ep = self._require_copy_engine() from_rank = self.get_rank() wait_ptr, wait_bits = self._flag_pointer_and_bits(wait_flag, default_bits=32) @@ -1176,6 +1216,8 @@ def put_tiles( """ Batched 2D tile transfer with optional shared wait/signal. + Requires initialization with ``engines=("sdma",)``. + Args: tiles: Sequence of pre-configured sdma_ep.Tile objects to_rank: Destination rank @@ -1188,6 +1230,7 @@ def put_tiles( async_op: If True, don't wait for completion channel: SDMA channel to use """ + sdma_ep = self._require_copy_engine() from_rank = self.get_rank() if len(tiles) != len(to_ptrs) or len(tiles) != len(to_strides): @@ -1232,6 +1275,8 @@ def quiet(self, to_rank: int = None, channel: int = 0): """ Wait for all outstanding SDMA operations to complete. + Requires initialization with ``engines=("sdma",)``. + Args: to_rank: If specified, wait only for ops to this rank. If None, wait for ops to all ranks. @@ -1242,6 +1287,7 @@ def quiet(self, to_rank: int = None, channel: int = 0): >>> shmem.quiet(to_rank=1) # Wait for completion >>> shmem.quiet() # Wait for all ranks """ + sdma_ep = self._require_copy_engine() src_rank = self.get_rank() if to_rank is not None: sdma_ep.quiet(src_rank, to_rank, channel) @@ -1698,7 +1744,7 @@ def reduce_scatter(self, output_tensor, input_tensor, op=None, group=None, async ) -def iris(heap_size=1 << 30, allocator_type="torch"): +def iris(heap_size=1 << 30, allocator_type="torch", engines=()): """ Create and return an Iris instance with the specified heap size. @@ -1706,6 +1752,7 @@ def iris(heap_size=1 << 30, allocator_type="torch"): heap_size (int): Size of the heap in bytes. Defaults to 1GB. allocator_type (str): Type of allocator to use. Options: "torch" (default), "vmem". Can be overridden with IRIS_ALLOCATOR environment variable. + engines (str or tuple of str, optional): See :class:`Iris`. Returns: Iris: An initialized Iris instance. @@ -1719,4 +1766,4 @@ def iris(heap_size=1 << 30, allocator_type="torch"): >>> iris_ctx = iris.iris(2**30, allocator_type="vmem") >>> tensor = iris_ctx.zeros(1024, 1024) """ - return Iris(heap_size, allocator_type) + return Iris(heap_size, allocator_type, engines=engines) diff --git a/iris/mem/triton/ops.py b/iris/mem/triton/ops.py index 1d022efac..f56544758 100644 --- a/iris/mem/triton/ops.py +++ b/iris/mem/triton/ops.py @@ -10,7 +10,6 @@ import triton import triton.language as tl from iris.mem.triton.context import __translate -from xio import sdma_ep from iris.device import sdma_utils @@ -377,7 +376,7 @@ def put( tl.store(translated_to_ptr, data, mask=mask, cache_modifier=store_cache_modifier) else: - ctx = copy_engine_ctx + (sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) + ctx = copy_engine_ctx + (sdma_utils.sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) queue_ptr_u32 = tl.load(ctx + 0).to(tl.pointer_type(tl.uint32)) read_ptr = tl.load(ctx + 1).to(tl.pointer_type(tl.uint64)) write_ptr = tl.load(ctx + 2).to(tl.pointer_type(tl.uint64)) @@ -422,7 +421,9 @@ def put( # Determine packet size based on copy type # Linear copy packet: 32 bytes for 1D, Sub-window copy packet: 80 bytes for 2D command_in_bytes = ( - sdma_ep.COPY_LINEAR_SUB_WINDOW_COMMAND_BYTES if is_2d_copy else sdma_ep.COPY_LINEAR_COMMAND_BYTES + sdma_utils.sdma_ep.COPY_LINEAR_SUB_WINDOW_COMMAND_BYTES + if is_2d_copy + else sdma_utils.sdma_ep.COPY_LINEAR_COMMAND_BYTES ) # Acquire space in the queue @@ -557,7 +558,7 @@ def atomic_add( if not use_copy_engine: return tl.atomic_add(translated_ptr, val, mask=mask, sem=sem, scope=scope) else: - handle = copy_engine_ctx + (sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) + handle = copy_engine_ctx + (sdma_utils.sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) queue_ptr_u32 = tl.load(handle + 0).to(tl.pointer_type(tl.uint32)) read_ptr = tl.load(handle + 1).to(tl.pointer_type(tl.uint64)) write_ptr = tl.load(handle + 2).to(tl.pointer_type(tl.uint64)) @@ -567,7 +568,7 @@ def atomic_add( dst_ptr_val = translated_ptr.to(tl.uint64) - command_in_bytes = sdma_ep.ATOMIC_COMMAND_BYTES + command_in_bytes = sdma_utils.sdma_ep.ATOMIC_COMMAND_BYTES # Acquire space (returns base index and wraparound offset) base, offset = sdma_utils.acquire_fadd( # base = sdma_utils.acquire( @@ -690,7 +691,7 @@ def atomic_cas( if not use_copy_engine: return tl.atomic_cas(translated_ptr, cmp, val, sem=sem, scope=scope) else: - handle = copy_engine_ctx + (sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) + handle = copy_engine_ctx + (sdma_utils.sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) queue_ptr_u32 = tl.load(handle + 0).to(tl.pointer_type(tl.uint32)) read_ptr = tl.load(handle + 1).to(tl.pointer_type(tl.uint64)) write_ptr = tl.load(handle + 2).to(tl.pointer_type(tl.uint64)) @@ -700,7 +701,7 @@ def atomic_cas( dst_ptr_val = translated_ptr.to(tl.uint64) - command_in_bytes = sdma_ep.ATOMIC_COMMAND_BYTES + command_in_bytes = sdma_utils.sdma_ep.ATOMIC_COMMAND_BYTES # Acquire space (returns base index and wraparound offset) base, offset = sdma_utils.acquire_fadd( queue_ptr_u32, @@ -984,7 +985,7 @@ def quiet(copy_engine_ctx: tl.tensor, to_rank): """ # Extract queue pointers from context # Context layout: [queue_buf, rptr, wptr, doorbell, cached_wptr, committed_wptr] - handle = copy_engine_ctx + (sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) + handle = copy_engine_ctx + (sdma_utils.sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) read_ptr = tl.load(handle + 1).to(tl.pointer_type(tl.uint64)) committed_wptr = tl.load(handle + 5).to(tl.pointer_type(tl.uint64)) diff --git a/pyproject.toml b/pyproject.toml index d85bab970..7828223d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,6 @@ dependencies = [ "requests", "ruff<0.16", # 0.16 formats code blocks in Markdown; see lint.yml "tritonblas @ git+https://github.com/ROCm/tritonBLAS.git@muhaawad/iris", - "rocm-xio @ git+https://github.com/ROCm/rocm-xio.git@cbe97e6392066bef7901121965ffadad19404da4", ] [project.urls] @@ -38,6 +37,10 @@ amd = [ # amdsmi is provided by the ROCm system package # Listed here for completeness ] +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", +] dev = [ "pytest", "black", diff --git a/scripts/run_core_tests_slurm.sh b/scripts/run_core_tests_slurm.sh index 55814d3ed..35ab99236 100755 --- a/scripts/run_core_tests_slurm.sh +++ b/scripts/run_core_tests_slurm.sh @@ -103,6 +103,6 @@ docker run --rm \ -w "$WORKSPACE_DIR" \ --entrypoint bash \ "$IMAGE_NAME" \ - -lc 'set -euo pipefail; git config --global --add safe.directory "$PWD"; python3 -m pip install -e ".[dev]"; bash scripts/run_core_tests.sh' + -lc 'set -euo pipefail; git config --global --add safe.directory "$PWD"; python3 -m pip install -e ".[dev,sdma]"; bash scripts/run_core_tests.sh' echo "Logs copied to $PERSIST_LOG_ROOT" diff --git a/scripts/run_example_slurm.sh b/scripts/run_example_slurm.sh index 6e461f4b0..30cfd988d 100644 --- a/scripts/run_example_slurm.sh +++ b/scripts/run_example_slurm.sh @@ -73,15 +73,15 @@ fi case "$INSTALL_METHOD" in editable) - INSTALL_CMD='python3 -m pip install -e ".[dev]"' + INSTALL_CMD='python3 -m pip install -e ".[dev,sdma]"' ;; install) - INSTALL_CMD='python3 -m pip install .' + INSTALL_CMD='python3 -m pip install ".[sdma]"' ;; git) REPO=${GITHUB_REPOSITORY:-ROCm/iris} SHA=${GITHUB_SHA:-HEAD} - INSTALL_CMD="python3 -m pip install git+https://github.com/${REPO}.git@${SHA}" + INSTALL_CMD="python3 -m pip install 'iris[sdma] @ git+https://github.com/${REPO}.git@${SHA}'" ;; *) echo "Unsupported INSTALL_METHOD: $INSTALL_METHOD" >&2 diff --git a/tests/examples/test_message_passing.py b/tests/examples/test_message_passing.py index 6df7d6886..a5e8f1712 100644 --- a/tests/examples/test_message_passing.py +++ b/tests/examples/test_message_passing.py @@ -12,6 +12,11 @@ current_dir = Path(__file__).parent +requires_copy_engine = pytest.mark.skipif( + importlib.util.find_spec("xio") is None, + reason="rocm-xio not installed (pip install iris[sdma])", +) + def load_example_module(relative_path: str, module_name: str): file_path = (current_dir / relative_path).resolve() @@ -40,7 +45,7 @@ def run_message_passing_kernels(module, args, *, use_copy_engine: bool = False): """Run the core message passing logic without command line argument parsing.""" shmem = None try: - shmem = iris.iris(args["heap_size"]) + shmem = iris.iris(args["heap_size"], engines=("sdma",) if use_copy_engine else ()) dtype = module.torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() @@ -189,6 +194,7 @@ def test_message_passing_put(dtype_str, buffer_size, heap_size, block_size): assert success, "Message passing put validation failed" +@requires_copy_engine @pytest.mark.parametrize("dtype_str", ["fp16", "fp32"]) @pytest.mark.parametrize("buffer_size, heap_size", [(4096, 1 << 20)]) @pytest.mark.parametrize("block_size", [512]) @@ -203,7 +209,7 @@ def run_host_initiated_copy_engine(module, args): """Execute the host-initiated message passing example logic.""" shmem = None try: - shmem = iris.iris(args["heap_size"]) + shmem = iris.iris(args["heap_size"], engines=("sdma",)) dtype = module.torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() @@ -263,6 +269,7 @@ def run_host_initiated_copy_engine(module, args): gc.collect() +@requires_copy_engine @pytest.mark.parametrize("dtype_str", ["fp16", "fp32"]) @pytest.mark.parametrize("buffer_size, heap_size", [(4096, 1 << 20)]) @pytest.mark.parametrize("block_size", [512]) diff --git a/tests/unittests/test_copy_engine_ops.py b/tests/unittests/test_copy_engine_ops.py index c2c4220e7..50f9ddaa8 100644 --- a/tests/unittests/test_copy_engine_ops.py +++ b/tests/unittests/test_copy_engine_ops.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: MIT +import importlib.util + import pytest import torch import triton @@ -7,6 +9,11 @@ import iris +pytestmark = pytest.mark.skipif( + importlib.util.find_spec("xio") is None, + reason="rocm-xio not installed (pip install iris[sdma])", +) + @triton.jit def _copy_engine_linear_kernel( @@ -90,7 +97,7 @@ def _make_grid(n, block): @pytest.mark.parametrize("num_elements", [256, 1024]) def test_copy_engine_device_linear_put(num_elements): - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -130,7 +137,7 @@ def test_copy_engine_device_linear_put(num_elements): def test_copy_engine_device_linear_put_no_mask(): num_elements = 512 block_size = 128 - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -166,7 +173,7 @@ def test_copy_engine_device_linear_put_no_mask(): @pytest.mark.parametrize("num_elements", [512, 2048]) def test_copy_engine_host_put(num_elements): - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -192,7 +199,7 @@ def test_copy_engine_host_put(num_elements): def test_copy_engine_host_put_loopback(): """Test a host-initiated copy-engine put when source and destination are the same rank.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) num_elements = 512 src = _allocate_symmetric_range(shmem, num_elements, torch.float32) @@ -218,7 +225,7 @@ def test_copy_engine_host_put_loopback(): def test_copy_engine_device_linear_put_loopback(): """Test a device-initiated copy-engine put when source and destination share one rank.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) rank = shmem.get_rank() num_elements = 256 @@ -271,7 +278,7 @@ def _copy_engine_atomic_kernel( # For int64, adding to 0xFFFFFFFF carries into the upper dword only if the SDMA add is 64-bit. @pytest.mark.parametrize("dtype, initial", [(torch.int32, 0), (torch.int64, 0xFFFFFFFF)]) def test_copy_engine_atomic_add(dtype, initial): - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -327,7 +334,7 @@ def _copy_engine_atomic_cas_kernel( def test_copy_engine_atomic_cas(): - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -449,7 +456,7 @@ def _copy_engine_2d_no_mask_kernel( @pytest.mark.parametrize("M,N", [(16, 16), (32, 64)]) def test_copy_engine_2d_tiled(M, N): """Test 2D tiled copy with strides.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -492,7 +499,7 @@ def test_copy_engine_2d_tiled(M, N): def test_copy_engine_2d_tiled_no_mask(): """Test full-tile 2D copy using shape-inferred SDMA packet size.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -579,7 +586,7 @@ def _copy_engine_put_signal_kernel( def test_copy_engine_put_with_signal(): """Test copy followed by atomic signal.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -667,7 +674,7 @@ def _copy_engine_multi_block_kernel( @pytest.mark.parametrize("num_blocks", [4, 8]) def test_copy_engine_multi_block_concurrent(num_blocks): """Test multiple workgroups using copy engine concurrently.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -712,7 +719,7 @@ def test_copy_engine_multi_block_concurrent(num_blocks): def test_copy_engine_zero_size(): """Test copy engine with zero-size transfer (should be no-op).""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -738,7 +745,7 @@ def test_copy_engine_zero_size(): def test_copy_engine_device_zero_size(): """Test device-side copy engine puts with fully-masked tiles (should be no-op).""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -792,7 +799,7 @@ def test_copy_engine_device_zero_size(): @pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.int32]) def test_copy_engine_different_dtypes(dtype): """Test copy engine with different data types.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -831,7 +838,7 @@ def test_copy_engine_different_dtypes(dtype): def test_copy_engine_bidirectional(): """Test both ranks doing copy engine operations simultaneously.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() diff --git a/tests/unittests/test_optional_sdma.py b/tests/unittests/test_optional_sdma.py new file mode 100644 index 000000000..b6c91659d --- /dev/null +++ b/tests/unittests/test_optional_sdma.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +import iris + + +def test_import_iris_does_not_import_xio(): + code = "import sys, iris; assert sys.modules.get('xio') is None, 'import iris loaded rocm-xio'" + root = Path(__file__).resolve().parents[2] + env = os.environ.copy() + env["PYTHONPATH"] = str(root) + os.pathsep + env.get("PYTHONPATH", "") + subprocess.check_call([sys.executable, "-c", code], env=env) + + +def test_default_context_rejects_host_sdma(): + shmem = iris.iris(1 << 20) + assert shmem.get_copy_engine_ctx() is None + + rank = shmem.get_rank() + peer = (rank + 1) % shmem.get_num_ranks() + buffer = shmem.zeros(16, dtype=torch.float32) + + with pytest.raises(RuntimeError, match=r"engines=\(\"sdma\",\)"): + shmem.put(buffer, to_rank=peer) + with pytest.raises(RuntimeError, match=r"engines=\(\"sdma\",\)"): + shmem.quiet() diff --git a/tests/unittests/test_sdma_packets.py b/tests/unittests/test_sdma_packets.py index 61372ebcc..cb7d83b91 100644 --- a/tests/unittests/test_sdma_packets.py +++ b/tests/unittests/test_sdma_packets.py @@ -7,11 +7,18 @@ import torch import triton import triton.language as tl -from xio import sdma_ep import iris from iris.device import sdma_utils +# The packet emitters read constants from rocm-xio, which only iris.iris() binds otherwise. +try: + sdma_ep = sdma_utils.bind_sdma_ep() +except ImportError: + sdma_ep = None + +pytestmark = pytest.mark.skipif(sdma_ep is None, reason="rocm-xio not installed (pip install iris[sdma])") + SENTINEL = 0x5A5A5A5A ADDR = 0x00007F12_34567890