From 52e1b625f3879135715d161e3915c97fe2742129 Mon Sep 17 00:00:00 2001 From: "Prajapati, Dimple" Date: Fri, 2 Oct 2026 22:32:35 +0000 Subject: [PATCH 1/8] Make rocm-xio an optional sdma extra; add enable_copy_engine flag Shader load/store users no longer need rocm-xio. It moves from the required dependencies to the sdma extra, and import iris no longer imports xio. iris.iris() / Iris() take enable_copy_engine (True/False/None). None reads IRIS_ENABLE_COPY_ENGINE (1/0/auto, strictly parsed); in auto mode SDMA is enabled only if rocm-xio is installed. With SDMA off, get_copy_engine_ctx() returns None and the host SDMA APIs raise RuntimeError. CI installs .[sdma]; external validation stays on the base install. SDMA tests skip when the copy engine is disabled. Signed-off-by: Prajapati, Dimple --- .github/scripts/run_new_examples.sh | 6 +- .github/scripts/run_perf_benchmark.sh | 2 +- .github/scripts/run_tests.sh | 12 +- .github/workflows/copilot-setup-steps.yml | 2 +- README.md | 6 + docs/getting-started/installation.md | 36 ++++++ iris/device/sdma_utils.py | 22 +++- iris/host/iris.py | 123 +++++++++++++++++--- iris/mem/triton/ops.py | 3 +- pyproject.toml | 5 +- scripts/run_example_slurm.sh | 6 +- tests/examples/test_message_passing.py | 8 ++ tests/unittests/test_copy_engine_flag.py | 135 ++++++++++++++++++++++ tests/unittests/test_copy_engine_ops.py | 6 + 14 files changed, 337 insertions(+), 35 deletions(-) create mode 100644 tests/unittests/test_copy_engine_flag.py 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..63a45389b 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): + +```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..2a402e9b2 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -94,6 +94,42 @@ 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, install rocm-xio on its own. Use the +commit pinned by the `sdma` extra in `pyproject.toml`, since Iris builds SDMA +packets against that rocm-xio version: + +```shell +pip install "rocm-xio @ git+https://github.com/ROCm/rocm-xio.git@cbe97e6392066bef7901121965ffadad19404da4" +``` + +No Iris reinstall is needed; the next `iris.iris()` picks it up in auto mode. + +`iris.iris()` decides whether to initialize SDMA queues as follows: + +| `enable_copy_engine` argument | `IRIS_ENABLE_COPY_ENGINE` | Result | +|---|---|---| +| `True` | any | SDMA on; error if rocm-xio is not installed | +| `False` | any | SDMA off; rocm-xio is never imported | +| `None` (default) | `1`, `true`, `yes`, `on` | SDMA on; error if rocm-xio is not installed | +| `None` (default) | `0`, `false`, `no`, `off` | SDMA off | +| `None` (default) | unset, empty, or `auto` | SDMA on only if rocm-xio is installed | + +Any other value of `IRIS_ENABLE_COPY_ENGINE` raises `ValueError`. With SDMA off, +`get_copy_engine_ctx()` returns `None` and the host SDMA APIs raise `RuntimeError`. + ## Selecting the AMD runtime for fabric communication When multiple ROCm versions are installed, set `LD_LIBRARY_PATH` before starting diff --git a/iris/device/sdma_utils.py b/iris/device/sdma_utils.py index b50335234..edc328c05 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 ImportError as exc: + raise ImportError( + "rocm-xio is required for the SDMA copy engine. " + "Install it with `pip install iris[sdma]`, or pass enable_copy_engine=False " + "to iris.iris() to use shader load/store only." + ) from exc + sdma_ep = imported + import iris.mem.triton.ops as triton_ops + + triton_ops.sdma_ep = sdma_ep + 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..5784d42fd 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -38,6 +38,7 @@ >>> data = device_ctx.load(buffer, from_rank=remote_rank) """ +import importlib.util import os from iris.host.distributed.helpers import ( @@ -53,7 +54,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 +72,54 @@ from iris.host.memory import tensors as tensor_creation from iris.host.platform.utils import is_simulation_env +_COPY_ENGINE_DISABLED = ( + "SDMA copy engine is disabled. Install rocm-xio with `pip install iris[sdma]` and " + "pass enable_copy_engine=True to iris.iris() (or set IRIS_ENABLE_COPY_ENGINE=1) " + "to use host-initiated SDMA puts and use_copy_engine=True." +) + +_ENV_TRUE = ("1", "true", "yes", "on") +_ENV_FALSE = ("0", "false", "no", "off") +_ENV_AUTO = ("", "auto") + + +def _xio_installed(): + return importlib.util.find_spec("xio") is not None + + +def _parse_copy_engine_env(value): + """Map ``IRIS_ENABLE_COPY_ENGINE`` to True, False, or None (auto).""" + normalized = value.strip().lower() + if normalized in _ENV_TRUE: + return True + if normalized in _ENV_FALSE: + return False + if normalized in _ENV_AUTO: + return None + raise ValueError( + f"Invalid IRIS_ENABLE_COPY_ENGINE={value!r}. " + f"Expected one of {', '.join(_ENV_TRUE + _ENV_FALSE)}, or auto/empty." + ) + + +def _resolve_enable_copy_engine(enable_copy_engine): + """ + Decide whether to initialize the SDMA copy engine. + + An explicit argument wins over ``IRIS_ENABLE_COPY_ENGINE``. In auto mode + (neither set), SDMA is enabled only if rocm-xio is installed. + """ + if enable_copy_engine is not None and not isinstance(enable_copy_engine, bool): + raise TypeError(f"enable_copy_engine must be True, False, or None, got {enable_copy_engine!r}") + requested = enable_copy_engine + if requested is None: + env = os.environ.get("IRIS_ENABLE_COPY_ENGINE") + if env is not None: + requested = _parse_copy_engine_env(env) + if requested is None: + return _xio_installed() + return requested + class Iris: """ @@ -83,6 +131,10 @@ 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" + enable_copy_engine (bool, optional): Initialize SDMA queues via rocm-xio + (``pip install iris[sdma]``). None (default) reads ``IRIS_ENABLE_COPY_ENGINE`` + (``1``/``0``/``auto``); in auto mode SDMA is enabled only if rocm-xio is + installed. False uses shader load/store only and never imports rocm-xio. Example: >>> ctx = iris.iris(heap_size=2**31) # 2GB heap with torch allocator @@ -91,9 +143,19 @@ class Iris: >>> # Use VMem allocator for memory oversubscription >>> ctx = iris.iris(heap_size=2**31, allocator_type="vmem") + + >>> # Shader load/store only + >>> ctx = iris.iris(heap_size=2**31, enable_copy_engine=False) """ - def __init__(self, heap_size=1 << 30, allocator_type="torch"): + def __init__(self, heap_size=1 << 30, allocator_type="torch", enable_copy_engine=None): + self.enable_copy_engine = _resolve_enable_copy_engine(enable_copy_engine) + self._sdma_ep = None + if self.enable_copy_engine: + 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 +199,29 @@ def __init__(self, heap_size=1 << 30, allocator_type="torch"): distributed_barrier() - # initialize copy engines + self.copy_engines_device_ctx = None + if self.enable_copy_engine: + self._init_copy_engines(num_gpus, num_ranks, cur_rank) + else: + self.debug("SDMA copy engine disabled; using shader load/store only") + + # 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 +251,11 @@ 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(_COPY_ENGINE_DISABLED) + return self._sdma_ep def __del__(self): """Cleanup resources on deletion.""" @@ -951,6 +1026,7 @@ def get_heap_bases(self): return self.heap_bases def get_copy_engine_ctx(self): + """Return the SDMA queue context, or None if the copy engine is disabled.""" return self.copy_engines_device_ctx @staticmethod @@ -1035,6 +1111,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 @@ -1125,6 +1202,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) @@ -1188,6 +1266,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): @@ -1242,6 +1321,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 +1778,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", enable_copy_engine=None): """ Create and return an Iris instance with the specified heap size. @@ -1706,6 +1786,10 @@ 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. + enable_copy_engine (bool, optional): Initialize SDMA queues via rocm-xio + (``pip install iris[sdma]``). None (default) reads ``IRIS_ENABLE_COPY_ENGINE`` + (``1``/``0``/``auto``); in auto mode SDMA is enabled only if rocm-xio is + installed. False uses shader load/store only and never imports rocm-xio. Returns: Iris: An initialized Iris instance. @@ -1718,5 +1802,8 @@ def iris(heap_size=1 << 30, allocator_type="torch"): >>> # Use VMem allocator >>> iris_ctx = iris.iris(2**30, allocator_type="vmem") >>> tensor = iris_ctx.zeros(1024, 1024) + + >>> # Shader load/store only + >>> iris_ctx = iris.iris(2**30, enable_copy_engine=False) """ - return Iris(heap_size, allocator_type) + return Iris(heap_size, allocator_type, enable_copy_engine=enable_copy_engine) diff --git a/iris/mem/triton/ops.py b/iris/mem/triton/ops.py index 1d022efac..1805f4a60 100644 --- a/iris/mem/triton/ops.py +++ b/iris/mem/triton/ops.py @@ -10,9 +10,10 @@ 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 +sdma_ep = None + @triton.jit def load( 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_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..2c3527d9c 100644 --- a/tests/examples/test_message_passing.py +++ b/tests/examples/test_message_passing.py @@ -6,12 +6,18 @@ import torch import triton import iris +from iris.host.iris import _resolve_enable_copy_engine import importlib.util from pathlib import Path current_dir = Path(__file__).parent +requires_copy_engine = pytest.mark.skipif( + not _resolve_enable_copy_engine(None), + reason="SDMA copy engine disabled (rocm-xio not installed or IRIS_ENABLE_COPY_ENGINE=0)", +) + def load_example_module(relative_path: str, module_name: str): file_path = (current_dir / relative_path).resolve() @@ -189,6 +195,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]) @@ -263,6 +270,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_flag.py b/tests/unittests/test_copy_engine_flag.py new file mode 100644 index 000000000..42fb73251 --- /dev/null +++ b/tests/unittests/test_copy_engine_flag.py @@ -0,0 +1,135 @@ +# 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 triton +import triton.language as tl + +import iris +import iris.device.sdma_utils as sdma_utils +from iris.host.iris import _parse_copy_engine_env, _resolve_enable_copy_engine + + +@pytest.fixture +def no_xio(monkeypatch): + """Make rocm-xio look uninstalled, as with a base `pip install iris`.""" + monkeypatch.setitem(sys.modules, "xio", None) + monkeypatch.setattr(sdma_utils, "sdma_ep", None) + + +@pytest.mark.parametrize( + "value, expected", + [ + ("1", True), + ("TRUE", True), + (" yes ", True), + ("on", True), + ("0", False), + ("False", False), + ("no", False), + ("off", False), + ("", None), + ("auto", None), + ], +) +def test_parse_copy_engine_env(value, expected): + assert _parse_copy_engine_env(value) is expected + + +@pytest.mark.parametrize("value", ["2", "enable", "ture", "disabled"]) +def test_parse_copy_engine_env_rejects_unknown(value): + with pytest.raises(ValueError, match="IRIS_ENABLE_COPY_ENGINE"): + _parse_copy_engine_env(value) + + +def test_explicit_argument_overrides_env(monkeypatch): + monkeypatch.setenv("IRIS_ENABLE_COPY_ENGINE", "1") + assert _resolve_enable_copy_engine(False) is False + monkeypatch.setenv("IRIS_ENABLE_COPY_ENGINE", "0") + assert _resolve_enable_copy_engine(True) is True + + +@pytest.mark.parametrize("value", [1, 0, "yes"]) +def test_explicit_argument_must_be_bool(value): + with pytest.raises(TypeError, match="enable_copy_engine"): + _resolve_enable_copy_engine(value) + + +def test_auto_follows_xio_install(monkeypatch, no_xio): + monkeypatch.delenv("IRIS_ENABLE_COPY_ENGINE", raising=False) + assert _resolve_enable_copy_engine(None) is False + monkeypatch.setenv("IRIS_ENABLE_COPY_ENGINE", "auto") + assert _resolve_enable_copy_engine(None) is False + + +def test_import_iris_does_not_import_xio(): + code = "import sys, iris; assert 'xio' not in sys.modules, '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_invalid_env_fails_before_init(monkeypatch): + monkeypatch.setenv("IRIS_ENABLE_COPY_ENGINE", "maybe") + with pytest.raises(ValueError, match="IRIS_ENABLE_COPY_ENGINE"): + iris.iris(1 << 20) + + +def test_enable_without_xio_fails_before_init(no_xio): + with pytest.raises(ImportError, match=r"iris\[sdma\]"): + iris.iris(1 << 20, enable_copy_engine=True) + + +@triton.jit +def _store_to_peer(buffer, n_elements, cur_rank, peer_rank, heap_bases, BLOCK_SIZE: tl.constexpr): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + values = tl.full([BLOCK_SIZE], 1, tl.float32) * (cur_rank + 1) + iris.store(buffer + offsets, values, cur_rank, peer_rank, heap_bases, mask=mask) + + +def _check_shader_store_without_sdma(shmem): + assert shmem.enable_copy_engine is False + assert shmem.get_copy_engine_ctx() is None + + rank = shmem.get_rank() + world_size = shmem.get_num_ranks() + peer = (rank + 1) % world_size + n_elements = 4096 + buffer = shmem.zeros(n_elements, dtype=torch.float32) + shmem.barrier() + + _store_to_peer[(triton.cdiv(n_elements, 1024),)]( + buffer, n_elements, rank, peer, shmem.get_heap_bases(), BLOCK_SIZE=1024 + ) + shmem.barrier() + + writer = (rank - 1) % world_size + torch.testing.assert_close(buffer, torch.full_like(buffer, writer + 1)) + + with pytest.raises(RuntimeError, match="copy engine is disabled"): + shmem.put(buffer, to_rank=peer) + with pytest.raises(RuntimeError, match="copy engine is disabled"): + shmem.quiet() + + +def test_disabled_shader_store(monkeypatch): + monkeypatch.delenv("IRIS_ENABLE_COPY_ENGINE", raising=False) + _check_shader_store_without_sdma(iris.iris(1 << 20, enable_copy_engine=False)) + + +def test_env_disabled_shader_store(monkeypatch): + monkeypatch.setenv("IRIS_ENABLE_COPY_ENGINE", "0") + _check_shader_store_without_sdma(iris.iris(1 << 20)) + + +def test_auto_without_xio_shader_store(monkeypatch, no_xio): + monkeypatch.delenv("IRIS_ENABLE_COPY_ENGINE", raising=False) + _check_shader_store_without_sdma(iris.iris(1 << 20)) diff --git a/tests/unittests/test_copy_engine_ops.py b/tests/unittests/test_copy_engine_ops.py index c2c4220e7..3832b9cb5 100644 --- a/tests/unittests/test_copy_engine_ops.py +++ b/tests/unittests/test_copy_engine_ops.py @@ -6,6 +6,12 @@ import triton.language as tl import iris +from iris.host.iris import _resolve_enable_copy_engine + +pytestmark = pytest.mark.skipif( + not _resolve_enable_copy_engine(None), + reason="SDMA copy engine disabled (rocm-xio not installed or IRIS_ENABLE_COPY_ENGINE=0)", +) @triton.jit From b7d8862553d00646711ea7b6d6d0fc1fc1d68cd2 Mon Sep 17 00:00:00 2001 From: "Prajapati, Dimple" Date: Sat, 10 Oct 2026 00:03:35 +0000 Subject: [PATCH 2/8] Drop IRIS_ENABLE_COPY_ENGINE; enable_copy_engine argument and sdma extra cover it - Remove the env var and its parser; None (default) enables SDMA only if rocm-xio is installed. - Bind rocm-xio in test_sdma_packets.py and skip it without rocm-xio. - Trim the flag tests to the no-rocm-xio cases and rename to test_optional_sdma.py. - Docs: use --force-reinstall --no-deps for a standalone rocm-xio install. Signed-off-by: Prajapati, Dimple --- docs/getting-started/installation.md | 28 ++++---- iris/host/iris.py | 54 ++------------- tests/examples/test_message_passing.py | 2 +- tests/unittests/test_copy_engine_ops.py | 2 +- ...y_engine_flag.py => test_optional_sdma.py} | 65 +------------------ tests/unittests/test_sdma_packets.py | 9 ++- 6 files changed, 36 insertions(+), 124 deletions(-) rename tests/unittests/{test_copy_engine_flag.py => test_optional_sdma.py} (51%) diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index 2a402e9b2..2fb324824 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -112,23 +112,27 @@ commit pinned by the `sdma` extra in `pyproject.toml`, since Iris builds SDMA packets against that rocm-xio version: ```shell -pip install "rocm-xio @ git+https://github.com/ROCm/rocm-xio.git@cbe97e6392066bef7901121965ffadad19404da4" +pip install --force-reinstall --no-deps "rocm-xio @ git+https://github.com/ROCm/rocm-xio.git@cbe97e6392066bef7901121965ffadad19404da4" ``` -No Iris reinstall is needed; the next `iris.iris()` picks it up in auto mode. +`--force-reinstall` makes pip rebuild rocm-xio even if another commit is already +installed, since every commit reports the same package version. `--no-deps` keeps +pip from reinstalling unrelated packages such as PyTorch. Iris itself does not need +to be reinstalled; the next `iris.iris()` picks rocm-xio up. -`iris.iris()` decides whether to initialize SDMA queues as follows: +The `enable_copy_engine` argument of `iris.iris()` controls whether SDMA queues are +initialized: -| `enable_copy_engine` argument | `IRIS_ENABLE_COPY_ENGINE` | Result | -|---|---|---| -| `True` | any | SDMA on; error if rocm-xio is not installed | -| `False` | any | SDMA off; rocm-xio is never imported | -| `None` (default) | `1`, `true`, `yes`, `on` | SDMA on; error if rocm-xio is not installed | -| `None` (default) | `0`, `false`, `no`, `off` | SDMA off | -| `None` (default) | unset, empty, or `auto` | SDMA on only if rocm-xio is installed | +| `enable_copy_engine` | Result | +|---|---| +| `None` (default) | SDMA on only if rocm-xio is installed | +| `True` | SDMA on; `ImportError` if rocm-xio is not installed | +| `False` | SDMA off; rocm-xio is never imported | -Any other value of `IRIS_ENABLE_COPY_ENGINE` raises `ValueError`. With SDMA off, -`get_copy_engine_ctx()` returns `None` and the host SDMA APIs raise `RuntimeError`. +Code that only uses shader load/store should pass `enable_copy_engine=False`, so it +does not initialize SDMA queues in environments that happen to have rocm-xio. With +SDMA off, `get_copy_engine_ctx()` returns `None` and the host SDMA APIs raise +`RuntimeError`. ## Selecting the AMD runtime for fabric communication diff --git a/iris/host/iris.py b/iris/host/iris.py index 5784d42fd..3d4db4b87 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -74,51 +74,15 @@ _COPY_ENGINE_DISABLED = ( "SDMA copy engine is disabled. Install rocm-xio with `pip install iris[sdma]` and " - "pass enable_copy_engine=True to iris.iris() (or set IRIS_ENABLE_COPY_ENGINE=1) " - "to use host-initiated SDMA puts and use_copy_engine=True." + "pass enable_copy_engine=True to iris.iris() to use host-initiated SDMA puts and use_copy_engine=True." ) -_ENV_TRUE = ("1", "true", "yes", "on") -_ENV_FALSE = ("0", "false", "no", "off") -_ENV_AUTO = ("", "auto") - - -def _xio_installed(): - return importlib.util.find_spec("xio") is not None - - -def _parse_copy_engine_env(value): - """Map ``IRIS_ENABLE_COPY_ENGINE`` to True, False, or None (auto).""" - normalized = value.strip().lower() - if normalized in _ENV_TRUE: - return True - if normalized in _ENV_FALSE: - return False - if normalized in _ENV_AUTO: - return None - raise ValueError( - f"Invalid IRIS_ENABLE_COPY_ENGINE={value!r}. " - f"Expected one of {', '.join(_ENV_TRUE + _ENV_FALSE)}, or auto/empty." - ) - def _resolve_enable_copy_engine(enable_copy_engine): - """ - Decide whether to initialize the SDMA copy engine. - - An explicit argument wins over ``IRIS_ENABLE_COPY_ENGINE``. In auto mode - (neither set), SDMA is enabled only if rocm-xio is installed. - """ - if enable_copy_engine is not None and not isinstance(enable_copy_engine, bool): - raise TypeError(f"enable_copy_engine must be True, False, or None, got {enable_copy_engine!r}") - requested = enable_copy_engine - if requested is None: - env = os.environ.get("IRIS_ENABLE_COPY_ENGINE") - if env is not None: - requested = _parse_copy_engine_env(env) - if requested is None: - return _xio_installed() - return requested + """Return ``enable_copy_engine``, or whether rocm-xio is installed if it is None.""" + if enable_copy_engine is None: + return importlib.util.find_spec("xio") is not None + return bool(enable_copy_engine) class Iris: @@ -132,8 +96,7 @@ class Iris: 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" enable_copy_engine (bool, optional): Initialize SDMA queues via rocm-xio - (``pip install iris[sdma]``). None (default) reads ``IRIS_ENABLE_COPY_ENGINE`` - (``1``/``0``/``auto``); in auto mode SDMA is enabled only if rocm-xio is + (``pip install iris[sdma]``). None (default) enables it only if rocm-xio is installed. False uses shader load/store only and never imports rocm-xio. Example: @@ -1786,10 +1749,7 @@ def iris(heap_size=1 << 30, allocator_type="torch", enable_copy_engine=None): 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. - enable_copy_engine (bool, optional): Initialize SDMA queues via rocm-xio - (``pip install iris[sdma]``). None (default) reads ``IRIS_ENABLE_COPY_ENGINE`` - (``1``/``0``/``auto``); in auto mode SDMA is enabled only if rocm-xio is - installed. False uses shader load/store only and never imports rocm-xio. + enable_copy_engine (bool, optional): See :class:`Iris`. Returns: Iris: An initialized Iris instance. diff --git a/tests/examples/test_message_passing.py b/tests/examples/test_message_passing.py index 2c3527d9c..ab24c0afc 100644 --- a/tests/examples/test_message_passing.py +++ b/tests/examples/test_message_passing.py @@ -15,7 +15,7 @@ requires_copy_engine = pytest.mark.skipif( not _resolve_enable_copy_engine(None), - reason="SDMA copy engine disabled (rocm-xio not installed or IRIS_ENABLE_COPY_ENGINE=0)", + reason="rocm-xio not installed (pip install iris[sdma])", ) diff --git a/tests/unittests/test_copy_engine_ops.py b/tests/unittests/test_copy_engine_ops.py index 3832b9cb5..c4459b035 100644 --- a/tests/unittests/test_copy_engine_ops.py +++ b/tests/unittests/test_copy_engine_ops.py @@ -10,7 +10,7 @@ pytestmark = pytest.mark.skipif( not _resolve_enable_copy_engine(None), - reason="SDMA copy engine disabled (rocm-xio not installed or IRIS_ENABLE_COPY_ENGINE=0)", + reason="rocm-xio not installed (pip install iris[sdma])", ) diff --git a/tests/unittests/test_copy_engine_flag.py b/tests/unittests/test_optional_sdma.py similarity index 51% rename from tests/unittests/test_copy_engine_flag.py rename to tests/unittests/test_optional_sdma.py index 42fb73251..d4b025fea 100644 --- a/tests/unittests/test_copy_engine_flag.py +++ b/tests/unittests/test_optional_sdma.py @@ -13,7 +13,6 @@ import iris import iris.device.sdma_utils as sdma_utils -from iris.host.iris import _parse_copy_engine_env, _resolve_enable_copy_engine @pytest.fixture @@ -23,65 +22,14 @@ def no_xio(monkeypatch): monkeypatch.setattr(sdma_utils, "sdma_ep", None) -@pytest.mark.parametrize( - "value, expected", - [ - ("1", True), - ("TRUE", True), - (" yes ", True), - ("on", True), - ("0", False), - ("False", False), - ("no", False), - ("off", False), - ("", None), - ("auto", None), - ], -) -def test_parse_copy_engine_env(value, expected): - assert _parse_copy_engine_env(value) is expected - - -@pytest.mark.parametrize("value", ["2", "enable", "ture", "disabled"]) -def test_parse_copy_engine_env_rejects_unknown(value): - with pytest.raises(ValueError, match="IRIS_ENABLE_COPY_ENGINE"): - _parse_copy_engine_env(value) - - -def test_explicit_argument_overrides_env(monkeypatch): - monkeypatch.setenv("IRIS_ENABLE_COPY_ENGINE", "1") - assert _resolve_enable_copy_engine(False) is False - monkeypatch.setenv("IRIS_ENABLE_COPY_ENGINE", "0") - assert _resolve_enable_copy_engine(True) is True - - -@pytest.mark.parametrize("value", [1, 0, "yes"]) -def test_explicit_argument_must_be_bool(value): - with pytest.raises(TypeError, match="enable_copy_engine"): - _resolve_enable_copy_engine(value) - - -def test_auto_follows_xio_install(monkeypatch, no_xio): - monkeypatch.delenv("IRIS_ENABLE_COPY_ENGINE", raising=False) - assert _resolve_enable_copy_engine(None) is False - monkeypatch.setenv("IRIS_ENABLE_COPY_ENGINE", "auto") - assert _resolve_enable_copy_engine(None) is False - - def test_import_iris_does_not_import_xio(): - code = "import sys, iris; assert 'xio' not in sys.modules, 'import iris loaded rocm-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_invalid_env_fails_before_init(monkeypatch): - monkeypatch.setenv("IRIS_ENABLE_COPY_ENGINE", "maybe") - with pytest.raises(ValueError, match="IRIS_ENABLE_COPY_ENGINE"): - iris.iris(1 << 20) - - def test_enable_without_xio_fails_before_init(no_xio): with pytest.raises(ImportError, match=r"iris\[sdma\]"): iris.iris(1 << 20, enable_copy_engine=True) @@ -120,16 +68,9 @@ def _check_shader_store_without_sdma(shmem): shmem.quiet() -def test_disabled_shader_store(monkeypatch): - monkeypatch.delenv("IRIS_ENABLE_COPY_ENGINE", raising=False) +def test_disabled_shader_store(): _check_shader_store_without_sdma(iris.iris(1 << 20, enable_copy_engine=False)) -def test_env_disabled_shader_store(monkeypatch): - monkeypatch.setenv("IRIS_ENABLE_COPY_ENGINE", "0") - _check_shader_store_without_sdma(iris.iris(1 << 20)) - - -def test_auto_without_xio_shader_store(monkeypatch, no_xio): - monkeypatch.delenv("IRIS_ENABLE_COPY_ENGINE", raising=False) +def test_auto_without_xio_shader_store(no_xio): _check_shader_store_without_sdma(iris.iris(1 << 20)) 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 From dc139d6f6b20c452385950ba22d28993d6566f95 Mon Sep 17 00:00:00 2001 From: Eric Feng Date: Fri, 9 Oct 2026 18:04:03 +0000 Subject: [PATCH 3/8] Make SDMA initialization opt-in --- docs/reference/host/class.md | 22 +++++++ .../message_passing_host_initiated.py | 2 +- .../06_message_passing/message_passing_put.py | 2 +- iris/host/iris.py | 61 +++++++++---------- tests/examples/test_message_passing.py | 4 +- tests/unittests/test_copy_engine_ops.py | 30 ++++----- tests/unittests/test_device_context.py | 1 + 7 files changed, 70 insertions(+), 52 deletions(-) diff --git a/docs/reference/host/class.md b/docs/reference/host/class.md index 87a0187bb..a0106ca9c 100644 --- a/docs/reference/host/class.md +++ b/docs/reference/host/class.md @@ -8,6 +8,28 @@ Prefer using the convenience factory over calling the constructor directly: .. autofunction:: iris.host.iris.iris ``` +## Optional SDMA Copy Engine + +The `copy_engine` parameter defaults to `None`. 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 `copy_engine="sdma"` to initialize both host- and device-initiated SDMA queues: + +```python +ctx = iris.iris(heap_size=1 << 30, copy_engine="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..29be2afd9 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"], copy_engine="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..8c4855db7 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"], copy_engine="sdma" if args["use_copy_engine"] else None) dtype = torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() diff --git a/iris/host/iris.py b/iris/host/iris.py index 3d4db4b87..e9d3671b0 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -38,7 +38,6 @@ >>> data = device_ctx.load(buffer, from_rank=remote_rank) """ -import importlib.util import os from iris.host.distributed.helpers import ( @@ -72,18 +71,6 @@ from iris.host.memory import tensors as tensor_creation from iris.host.platform.utils import is_simulation_env -_COPY_ENGINE_DISABLED = ( - "SDMA copy engine is disabled. Install rocm-xio with `pip install iris[sdma]` and " - "pass enable_copy_engine=True to iris.iris() to use host-initiated SDMA puts and use_copy_engine=True." -) - - -def _resolve_enable_copy_engine(enable_copy_engine): - """Return ``enable_copy_engine``, or whether rocm-xio is installed if it is None.""" - if enable_copy_engine is None: - return importlib.util.find_spec("xio") is not None - return bool(enable_copy_engine) - class Iris: """ @@ -95,9 +82,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" - enable_copy_engine (bool, optional): Initialize SDMA queues via rocm-xio - (``pip install iris[sdma]``). None (default) enables it only if rocm-xio is - installed. False uses shader load/store only and never imports rocm-xio. + copy_engine (str, optional): Copy engine to initialize. Use "sdma" for host- and + device-initiated SDMA operations; requires rocm-xio (``pip install iris[sdma]``). + Defaults to None (disabled; rocm-xio is never imported). Example: >>> ctx = iris.iris(heap_size=2**31) # 2GB heap with torch allocator @@ -107,14 +94,13 @@ class Iris: >>> # Use VMem allocator for memory oversubscription >>> ctx = iris.iris(heap_size=2**31, allocator_type="vmem") - >>> # Shader load/store only - >>> ctx = iris.iris(heap_size=2**31, enable_copy_engine=False) """ - def __init__(self, heap_size=1 << 30, allocator_type="torch", enable_copy_engine=None): - self.enable_copy_engine = _resolve_enable_copy_engine(enable_copy_engine) + def __init__(self, heap_size=1 << 30, allocator_type="torch", copy_engine=None): + if copy_engine not in (None, "sdma"): + raise ValueError(f"Unsupported copy engine: {copy_engine!r}. Expected None or 'sdma'.") self._sdma_ep = None - if self.enable_copy_engine: + if copy_engine == "sdma": from iris.device.sdma_utils import bind_sdma_ep self._sdma_ep = bind_sdma_ep() @@ -163,10 +149,8 @@ def __init__(self, heap_size=1 << 30, allocator_type="torch", enable_copy_engine distributed_barrier() self.copy_engines_device_ctx = None - if self.enable_copy_engine: + if self._sdma_ep is not None: self._init_copy_engines(num_gpus, num_ranks, cur_rank) - else: - self.debug("SDMA copy engine disabled; using shader load/store only") # Initialize CCL interface self.ccl = self.CCL(self) @@ -217,7 +201,9 @@ def _init_copy_engines(self, num_gpus, num_ranks, cur_rank): def _require_copy_engine(self): if self._sdma_ep is None: - raise RuntimeError(_COPY_ENGINE_DISABLED) + raise RuntimeError( + 'SDMA is disabled; initialize Iris with copy_engine="sdma" (needs `pip install iris[sdma]`).' + ) return self._sdma_ep def __del__(self): @@ -988,10 +974,14 @@ def get_heap_bases(self): """ return self.heap_bases - def get_copy_engine_ctx(self): - """Return the SDMA queue context, or None if the copy engine is disabled.""" + 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 + def _require_copy_engine(self): + if self.copy_engines_device_ctx is None: + raise RuntimeError('SDMA is disabled; initialize Iris with copy_engine="sdma".') + @staticmethod def _dtype_to_flag_bits(dtype: torch.dtype) -> int: if dtype in (torch.int32, torch.int): @@ -1042,6 +1032,8 @@ def put( """ One-sided put operation with optional wait (POLL) and signal (ATOMIC). + Requires initialization with ``copy_engine="sdma"``. + Supports: - Simple copy: put(src, to_rank) - Copy + signal: put(src, to_rank, signal_flag=flag) @@ -1134,6 +1126,8 @@ def put_tile( """ 2D tile transfer with optional wait/signal (sub-window copy). + Requires initialization with ``copy_engine="sdma"``. + Low-level API - caller provides pre-translated pointers for performance. Args: @@ -1217,6 +1211,8 @@ def put_tiles( """ Batched 2D tile transfer with optional shared wait/signal. + Requires initialization with ``copy_engine="sdma"``. + Args: tiles: Sequence of pre-configured sdma_ep.Tile objects to_rank: Destination rank @@ -1274,6 +1270,8 @@ def quiet(self, to_rank: int = None, channel: int = 0): """ Wait for all outstanding SDMA operations to complete. + Requires initialization with ``copy_engine="sdma"``. + Args: to_rank: If specified, wait only for ops to this rank. If None, wait for ops to all ranks. @@ -1741,7 +1739,7 @@ def reduce_scatter(self, output_tensor, input_tensor, op=None, group=None, async ) -def iris(heap_size=1 << 30, allocator_type="torch", enable_copy_engine=None): +def iris(heap_size=1 << 30, allocator_type="torch", copy_engine=None): """ Create and return an Iris instance with the specified heap size. @@ -1749,7 +1747,7 @@ def iris(heap_size=1 << 30, allocator_type="torch", enable_copy_engine=None): 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. - enable_copy_engine (bool, optional): See :class:`Iris`. + copy_engine (str, optional): See :class:`Iris`. Returns: Iris: An initialized Iris instance. @@ -1762,8 +1760,5 @@ def iris(heap_size=1 << 30, allocator_type="torch", enable_copy_engine=None): >>> # Use VMem allocator >>> iris_ctx = iris.iris(2**30, allocator_type="vmem") >>> tensor = iris_ctx.zeros(1024, 1024) - - >>> # Shader load/store only - >>> iris_ctx = iris.iris(2**30, enable_copy_engine=False) """ - return Iris(heap_size, allocator_type, enable_copy_engine=enable_copy_engine) + return Iris(heap_size, allocator_type, copy_engine=copy_engine) diff --git a/tests/examples/test_message_passing.py b/tests/examples/test_message_passing.py index ab24c0afc..d0646c1e6 100644 --- a/tests/examples/test_message_passing.py +++ b/tests/examples/test_message_passing.py @@ -46,7 +46,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"], copy_engine="sdma" if use_copy_engine else None) dtype = module.torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() @@ -210,7 +210,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"], copy_engine="sdma") dtype = module.torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() diff --git a/tests/unittests/test_copy_engine_ops.py b/tests/unittests/test_copy_engine_ops.py index c4459b035..c1ba701d5 100644 --- a/tests/unittests/test_copy_engine_ops.py +++ b/tests/unittests/test_copy_engine_ops.py @@ -96,7 +96,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -136,7 +136,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -172,7 +172,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -198,7 +198,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, copy_engine="sdma") num_elements = 512 src = _allocate_symmetric_range(shmem, num_elements, torch.float32) @@ -224,7 +224,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, copy_engine="sdma") rank = shmem.get_rank() num_elements = 256 @@ -277,7 +277,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -333,7 +333,7 @@ def _copy_engine_atomic_cas_kernel( def test_copy_engine_atomic_cas(): - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -455,7 +455,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -498,7 +498,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -585,7 +585,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -673,7 +673,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -718,7 +718,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -744,7 +744,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -798,7 +798,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -837,7 +837,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, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() diff --git a/tests/unittests/test_device_context.py b/tests/unittests/test_device_context.py index 874afa102..5c628da31 100644 --- a/tests/unittests/test_device_context.py +++ b/tests/unittests/test_device_context.py @@ -478,6 +478,7 @@ def test_device_context_tracing_1d_address(): def test_device_context_initialize(): """Test DeviceContext.initialize() creates valid context.""" ctx = iris.iris(1 << 20) + assert ctx.get_copy_engine_ctx() is None cur_rank = ctx.get_rank() num_ranks = ctx.get_num_ranks() From 010b5959bcfa4a1d440690c5edae17dabed84ae6 Mon Sep 17 00:00:00 2001 From: "Prajapati, Dimple" Date: Sat, 10 Oct 2026 00:21:23 +0000 Subject: [PATCH 4/8] Adapt optional rocm-xio packaging to copy_engine="sdma" opt-in - Bind rocm-xio lazily only when copy_engine="sdma"; ImportError points to iris[sdma]. - Skip SDMA tests when rocm-xio is not installed. - test_optional_sdma.py: no-xio import, fail-fast, and default shader-store checks. - Docs: installation covers the extra, class reference covers copy_engine. Signed-off-by: Prajapati, Dimple --- README.md | 2 +- docs/getting-started/installation.md | 18 ++++-------------- docs/reference/host/class.md | 4 +++- iris/device/sdma_utils.py | 4 ++-- tests/examples/test_message_passing.py | 3 +-- tests/unittests/test_copy_engine_ops.py | 5 +++-- tests/unittests/test_optional_sdma.py | 20 ++++++-------------- 7 files changed, 20 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 63a45389b..7360a12f7 100644 --- a/README.md +++ b/README.md @@ -203,7 +203,7 @@ 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): +The SDMA copy engine is opt-in: install the optional `sdma` extra (rocm-xio) and pass `copy_engine="sdma"` to `iris.iris()`: ```shell pip install "iris[sdma] @ git+https://github.com/ROCm/iris.git" diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index 2fb324824..4472ae5a0 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -118,21 +118,11 @@ pip install --force-reinstall --no-deps "rocm-xio @ git+https://github.com/ROCm/ `--force-reinstall` makes pip rebuild rocm-xio even if another commit is already installed, since every commit reports the same package version. `--no-deps` keeps pip from reinstalling unrelated packages such as PyTorch. Iris itself does not need -to be reinstalled; the next `iris.iris()` picks rocm-xio up. +to be reinstalled. -The `enable_copy_engine` argument of `iris.iris()` controls whether SDMA queues are -initialized: - -| `enable_copy_engine` | Result | -|---|---| -| `None` (default) | SDMA on only if rocm-xio is installed | -| `True` | SDMA on; `ImportError` if rocm-xio is not installed | -| `False` | SDMA off; rocm-xio is never imported | - -Code that only uses shader load/store should pass `enable_copy_engine=False`, so it -does not initialize SDMA queues in environments that happen to have rocm-xio. With -SDMA off, `get_copy_engine_ctx()` returns `None` and the host SDMA APIs raise -`RuntimeError`. +Installing rocm-xio does not turn SDMA on. Pass `copy_engine="sdma"` to `iris.iris()` +to initialize the SDMA queues; see [Iris Class](../reference/host/class.md). Without +it, Iris neither imports rocm-xio nor creates SDMA queues. ## Selecting the AMD runtime for fabric communication diff --git a/docs/reference/host/class.md b/docs/reference/host/class.md index a0106ca9c..ffcd0a5d6 100644 --- a/docs/reference/host/class.md +++ b/docs/reference/host/class.md @@ -14,7 +14,9 @@ The `copy_engine` parameter defaults to `None`. 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 `copy_engine="sdma"` to initialize both host- and device-initiated SDMA queues: +Pass `copy_engine="sdma"` to initialize both host- and device-initiated SDMA queues. +This needs rocm-xio (`pip install iris[sdma]`, see [Installation](../../getting-started/installation.md)); +without it, `iris.iris()` raises `ImportError`: ```python ctx = iris.iris(heap_size=1 << 30, copy_engine="sdma") diff --git a/iris/device/sdma_utils.py b/iris/device/sdma_utils.py index edc328c05..12086b94d 100644 --- a/iris/device/sdma_utils.py +++ b/iris/device/sdma_utils.py @@ -25,8 +25,8 @@ def bind_sdma_ep(): except ImportError as exc: raise ImportError( "rocm-xio is required for the SDMA copy engine. " - "Install it with `pip install iris[sdma]`, or pass enable_copy_engine=False " - "to iris.iris() to use shader load/store only." + 'Install it with `pip install iris[sdma]`, or drop copy_engine="sdma" ' + "to use shader load/store only." ) from exc sdma_ep = imported import iris.mem.triton.ops as triton_ops diff --git a/tests/examples/test_message_passing.py b/tests/examples/test_message_passing.py index d0646c1e6..9c39d9e42 100644 --- a/tests/examples/test_message_passing.py +++ b/tests/examples/test_message_passing.py @@ -6,7 +6,6 @@ import torch import triton import iris -from iris.host.iris import _resolve_enable_copy_engine import importlib.util from pathlib import Path @@ -14,7 +13,7 @@ current_dir = Path(__file__).parent requires_copy_engine = pytest.mark.skipif( - not _resolve_enable_copy_engine(None), + importlib.util.find_spec("xio") is None, reason="rocm-xio not installed (pip install iris[sdma])", ) diff --git a/tests/unittests/test_copy_engine_ops.py b/tests/unittests/test_copy_engine_ops.py index c1ba701d5..d94c3db23 100644 --- a/tests/unittests/test_copy_engine_ops.py +++ b/tests/unittests/test_copy_engine_ops.py @@ -1,15 +1,16 @@ # SPDX-License-Identifier: MIT +import importlib.util + import pytest import torch import triton import triton.language as tl import iris -from iris.host.iris import _resolve_enable_copy_engine pytestmark = pytest.mark.skipif( - not _resolve_enable_copy_engine(None), + importlib.util.find_spec("xio") is None, reason="rocm-xio not installed (pip install iris[sdma])", ) diff --git a/tests/unittests/test_optional_sdma.py b/tests/unittests/test_optional_sdma.py index d4b025fea..bb805d7d3 100644 --- a/tests/unittests/test_optional_sdma.py +++ b/tests/unittests/test_optional_sdma.py @@ -30,9 +30,9 @@ def test_import_iris_does_not_import_xio(): subprocess.check_call([sys.executable, "-c", code], env=env) -def test_enable_without_xio_fails_before_init(no_xio): +def test_sdma_without_xio_fails_before_init(no_xio): with pytest.raises(ImportError, match=r"iris\[sdma\]"): - iris.iris(1 << 20, enable_copy_engine=True) + iris.iris(1 << 20, copy_engine="sdma") @triton.jit @@ -43,8 +43,8 @@ def _store_to_peer(buffer, n_elements, cur_rank, peer_rank, heap_bases, BLOCK_SI iris.store(buffer + offsets, values, cur_rank, peer_rank, heap_bases, mask=mask) -def _check_shader_store_without_sdma(shmem): - assert shmem.enable_copy_engine is False +def test_default_shader_store_without_xio(no_xio): + shmem = iris.iris(1 << 20) assert shmem.get_copy_engine_ctx() is None rank = shmem.get_rank() @@ -62,15 +62,7 @@ def _check_shader_store_without_sdma(shmem): writer = (rank - 1) % world_size torch.testing.assert_close(buffer, torch.full_like(buffer, writer + 1)) - with pytest.raises(RuntimeError, match="copy engine is disabled"): + with pytest.raises(RuntimeError, match='copy_engine="sdma"'): shmem.put(buffer, to_rank=peer) - with pytest.raises(RuntimeError, match="copy engine is disabled"): + with pytest.raises(RuntimeError, match='copy_engine="sdma"'): shmem.quiet() - - -def test_disabled_shader_store(): - _check_shader_store_without_sdma(iris.iris(1 << 20, enable_copy_engine=False)) - - -def test_auto_without_xio_shader_store(no_xio): - _check_shader_store_without_sdma(iris.iris(1 << 20)) From 517bf00c3003df473bdf468008234e5946dc434a Mon Sep 17 00:00:00 2001 From: "Prajapati, Dimple" Date: Sat, 10 Oct 2026 00:34:31 +0000 Subject: [PATCH 5/8] Remove duplicate _require_copy_engine; trim repeated install docs Signed-off-by: Prajapati, Dimple --- README.md | 2 +- docs/getting-started/installation.md | 5 ++--- docs/reference/host/class.md | 3 +-- iris/host/iris.py | 8 ++------ 4 files changed, 6 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 7360a12f7..63a45389b 100644 --- a/README.md +++ b/README.md @@ -203,7 +203,7 @@ For a quick installation directly from the repository: pip install git+https://github.com/ROCm/iris.git ``` -The SDMA copy engine is opt-in: install the optional `sdma` extra (rocm-xio) and pass `copy_engine="sdma"` to `iris.iris()`: +The SDMA copy engine needs the optional `sdma` extra (rocm-xio): ```shell pip install "iris[sdma] @ git+https://github.com/ROCm/iris.git" diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index 4472ae5a0..3a0f19b9d 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -120,9 +120,8 @@ installed, since every commit reports the same package version. `--no-deps` keep pip from reinstalling unrelated packages such as PyTorch. Iris itself does not need to be reinstalled. -Installing rocm-xio does not turn SDMA on. Pass `copy_engine="sdma"` to `iris.iris()` -to initialize the SDMA queues; see [Iris Class](../reference/host/class.md). Without -it, Iris neither imports rocm-xio nor creates SDMA queues. +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 diff --git a/docs/reference/host/class.md b/docs/reference/host/class.md index ffcd0a5d6..f34668e03 100644 --- a/docs/reference/host/class.md +++ b/docs/reference/host/class.md @@ -15,8 +15,7 @@ initialization and queue creation, and `get_copy_engine_ctx()` returns `None`. Regular Triton and Gluon memory operations and collectives work without SDMA. Pass `copy_engine="sdma"` to initialize both host- and device-initiated SDMA queues. -This needs rocm-xio (`pip install iris[sdma]`, see [Installation](../../getting-started/installation.md)); -without it, `iris.iris()` raises `ImportError`: +This needs the `sdma` extra ([Installation](../../getting-started/installation.md)): ```python ctx = iris.iris(heap_size=1 << 30, copy_engine="sdma") diff --git a/iris/host/iris.py b/iris/host/iris.py index e9d3671b0..6b093518c 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -83,8 +83,8 @@ class Iris: 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" copy_engine (str, optional): Copy engine to initialize. Use "sdma" for host- and - device-initiated SDMA operations; requires rocm-xio (``pip install iris[sdma]``). - Defaults to None (disabled; rocm-xio is never imported). + device-initiated SDMA operations (needs ``pip install iris[sdma]``). Defaults to + None (disabled). Example: >>> ctx = iris.iris(heap_size=2**31) # 2GB heap with torch allocator @@ -978,10 +978,6 @@ 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 - def _require_copy_engine(self): - if self.copy_engines_device_ctx is None: - raise RuntimeError('SDMA is disabled; initialize Iris with copy_engine="sdma".') - @staticmethod def _dtype_to_flag_bits(dtype: torch.dtype) -> int: if dtype in (torch.int32, torch.int): From 40a251d5d19de68d4f7dfb03a5cbd0cbdfa34bcb Mon Sep 17 00:00:00 2001 From: "Prajapati, Dimple" Date: Tue, 13 Oct 2026 01:11:24 +0000 Subject: [PATCH 6/8] Address review feedback engines argument, sdma_utils owns sdma_ep, simpler install docs Signed-off-by: Prajapati, Dimple --- README.md | 2 +- docs/getting-started/installation.md | 15 ++----- docs/reference/host/class.md | 6 +-- .../message_passing_host_initiated.py | 2 +- .../06_message_passing/message_passing_put.py | 2 +- iris/device/sdma_utils.py | 14 +++---- iris/host/iris.py | 39 ++++++++++++------- iris/mem/triton/ops.py | 18 ++++----- tests/examples/test_message_passing.py | 4 +- tests/unittests/test_copy_engine_ops.py | 30 +++++++------- tests/unittests/test_optional_sdma.py | 19 ++------- 11 files changed, 69 insertions(+), 82 deletions(-) diff --git a/README.md b/README.md index 63a45389b..3dc623343 100644 --- a/README.md +++ b/README.md @@ -203,7 +203,7 @@ 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): +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" diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index 3a0f19b9d..99e851b02 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -107,18 +107,9 @@ pip install "iris[sdma] @ git+https://github.com/ROCm/iris.git" pip install -e ".[sdma]" ``` -To add SDMA to an existing Iris install, install rocm-xio on its own. Use the -commit pinned by the `sdma` extra in `pyproject.toml`, since Iris builds SDMA -packets against that rocm-xio version: - -```shell -pip install --force-reinstall --no-deps "rocm-xio @ git+https://github.com/ROCm/rocm-xio.git@cbe97e6392066bef7901121965ffadad19404da4" -``` - -`--force-reinstall` makes pip rebuild rocm-xio even if another commit is already -installed, since every commit reports the same package version. `--no-deps` keeps -pip from reinstalling unrelated packages such as PyTorch. Iris itself does not need -to be reinstalled. +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. diff --git a/docs/reference/host/class.md b/docs/reference/host/class.md index f34668e03..46845d8b8 100644 --- a/docs/reference/host/class.md +++ b/docs/reference/host/class.md @@ -10,15 +10,15 @@ Prefer using the convenience factory over calling the constructor directly: ## Optional SDMA Copy Engine -The `copy_engine` parameter defaults to `None`. In this mode, Iris skips SDMA +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 `copy_engine="sdma"` to initialize both host- and device-initiated SDMA queues. +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, copy_engine="sdma") +ctx = iris.iris(heap_size=1 << 30, engines=("sdma",)) copy_engine_ctx = ctx.get_copy_engine_ctx() ``` diff --git a/examples/06_message_passing/message_passing_host_initiated.py b/examples/06_message_passing/message_passing_host_initiated.py index 29be2afd9..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"], copy_engine="sdma") + 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 8c4855db7..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"], copy_engine="sdma" if args["use_copy_engine"] else None) + 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 12086b94d..baf7b170d 100644 --- a/iris/device/sdma_utils.py +++ b/iris/device/sdma_utils.py @@ -22,16 +22,16 @@ def bind_sdma_ep(): if sdma_ep is None: try: from xio import sdma_ep as imported - except ImportError as exc: - raise ImportError( + 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 copy_engine="sdma" ' - "to use shader load/store only." + '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 - import iris.mem.triton.ops as triton_ops - - triton_ops.sdma_ep = sdma_ep return sdma_ep diff --git a/iris/host/iris.py b/iris/host/iris.py index 6b093518c..49d4a9913 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -71,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: """ @@ -82,9 +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" - copy_engine (str, optional): Copy engine to initialize. Use "sdma" for host- and - device-initiated SDMA operations (needs ``pip install iris[sdma]``). Defaults to - None (disabled). + 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 @@ -96,11 +98,18 @@ class Iris: """ - def __init__(self, heap_size=1 << 30, allocator_type="torch", copy_engine=None): - if copy_engine not in (None, "sdma"): - raise ValueError(f"Unsupported copy engine: {copy_engine!r}. Expected None or 'sdma'.") + 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 copy_engine == "sdma": + if "sdma" in engines: from iris.device.sdma_utils import bind_sdma_ep self._sdma_ep = bind_sdma_ep() @@ -202,7 +211,7 @@ def _init_copy_engines(self, num_gpus, num_ranks, cur_rank): def _require_copy_engine(self): if self._sdma_ep is None: raise RuntimeError( - 'SDMA is disabled; initialize Iris with copy_engine="sdma" (needs `pip install iris[sdma]`).' + 'SDMA is disabled; initialize Iris with engines=("sdma",) (needs `pip install iris[sdma]`).' ) return self._sdma_ep @@ -1028,7 +1037,7 @@ def put( """ One-sided put operation with optional wait (POLL) and signal (ATOMIC). - Requires initialization with ``copy_engine="sdma"``. + Requires initialization with ``engines=("sdma",)``. Supports: - Simple copy: put(src, to_rank) @@ -1122,7 +1131,7 @@ def put_tile( """ 2D tile transfer with optional wait/signal (sub-window copy). - Requires initialization with ``copy_engine="sdma"``. + Requires initialization with ``engines=("sdma",)``. Low-level API - caller provides pre-translated pointers for performance. @@ -1207,7 +1216,7 @@ def put_tiles( """ Batched 2D tile transfer with optional shared wait/signal. - Requires initialization with ``copy_engine="sdma"``. + Requires initialization with ``engines=("sdma",)``. Args: tiles: Sequence of pre-configured sdma_ep.Tile objects @@ -1266,7 +1275,7 @@ def quiet(self, to_rank: int = None, channel: int = 0): """ Wait for all outstanding SDMA operations to complete. - Requires initialization with ``copy_engine="sdma"``. + Requires initialization with ``engines=("sdma",)``. Args: to_rank: If specified, wait only for ops to this rank. @@ -1735,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", copy_engine=None): +def iris(heap_size=1 << 30, allocator_type="torch", engines=()): """ Create and return an Iris instance with the specified heap size. @@ -1743,7 +1752,7 @@ def iris(heap_size=1 << 30, allocator_type="torch", copy_engine=None): 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. - copy_engine (str, optional): See :class:`Iris`. + engines (str or tuple of str, optional): See :class:`Iris`. Returns: Iris: An initialized Iris instance. @@ -1757,4 +1766,4 @@ def iris(heap_size=1 << 30, allocator_type="torch", copy_engine=None): >>> iris_ctx = iris.iris(2**30, allocator_type="vmem") >>> tensor = iris_ctx.zeros(1024, 1024) """ - return Iris(heap_size, allocator_type, copy_engine=copy_engine) + return Iris(heap_size, allocator_type, engines=engines) diff --git a/iris/mem/triton/ops.py b/iris/mem/triton/ops.py index 1805f4a60..f56544758 100644 --- a/iris/mem/triton/ops.py +++ b/iris/mem/triton/ops.py @@ -12,8 +12,6 @@ from iris.mem.triton.context import __translate from iris.device import sdma_utils -sdma_ep = None - @triton.jit def load( @@ -378,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)) @@ -423,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 @@ -558,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)) @@ -568,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( @@ -691,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)) @@ -701,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, @@ -985,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/tests/examples/test_message_passing.py b/tests/examples/test_message_passing.py index 9c39d9e42..a5e8f1712 100644 --- a/tests/examples/test_message_passing.py +++ b/tests/examples/test_message_passing.py @@ -45,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"], copy_engine="sdma" if use_copy_engine else None) + 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() @@ -209,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"], copy_engine="sdma") + 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() diff --git a/tests/unittests/test_copy_engine_ops.py b/tests/unittests/test_copy_engine_ops.py index d94c3db23..50f9ddaa8 100644 --- a/tests/unittests/test_copy_engine_ops.py +++ b/tests/unittests/test_copy_engine_ops.py @@ -97,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -137,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -173,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -199,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) num_elements = 512 src = _allocate_symmetric_range(shmem, num_elements, torch.float32) @@ -225,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) rank = shmem.get_rank() num_elements = 256 @@ -278,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -334,7 +334,7 @@ def _copy_engine_atomic_cas_kernel( def test_copy_engine_atomic_cas(): - shmem = iris.iris(1 << 20, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -456,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -499,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -586,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -674,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -719,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -745,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -799,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, copy_engine="sdma") + shmem = iris.iris(1 << 20, engines=("sdma",)) _require_two_ranks(shmem) rank = shmem.get_rank() @@ -838,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, copy_engine="sdma") + 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 index bb805d7d3..fae409477 100644 --- a/tests/unittests/test_optional_sdma.py +++ b/tests/unittests/test_optional_sdma.py @@ -12,14 +12,6 @@ import triton.language as tl import iris -import iris.device.sdma_utils as sdma_utils - - -@pytest.fixture -def no_xio(monkeypatch): - """Make rocm-xio look uninstalled, as with a base `pip install iris`.""" - monkeypatch.setitem(sys.modules, "xio", None) - monkeypatch.setattr(sdma_utils, "sdma_ep", None) def test_import_iris_does_not_import_xio(): @@ -30,11 +22,6 @@ def test_import_iris_does_not_import_xio(): subprocess.check_call([sys.executable, "-c", code], env=env) -def test_sdma_without_xio_fails_before_init(no_xio): - with pytest.raises(ImportError, match=r"iris\[sdma\]"): - iris.iris(1 << 20, copy_engine="sdma") - - @triton.jit def _store_to_peer(buffer, n_elements, cur_rank, peer_rank, heap_bases, BLOCK_SIZE: tl.constexpr): offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) @@ -43,7 +30,7 @@ def _store_to_peer(buffer, n_elements, cur_rank, peer_rank, heap_bases, BLOCK_SI iris.store(buffer + offsets, values, cur_rank, peer_rank, heap_bases, mask=mask) -def test_default_shader_store_without_xio(no_xio): +def test_default_shader_store_without_sdma(): shmem = iris.iris(1 << 20) assert shmem.get_copy_engine_ctx() is None @@ -62,7 +49,7 @@ def test_default_shader_store_without_xio(no_xio): writer = (rank - 1) % world_size torch.testing.assert_close(buffer, torch.full_like(buffer, writer + 1)) - with pytest.raises(RuntimeError, match='copy_engine="sdma"'): + with pytest.raises(RuntimeError, match=r"engines=\(\"sdma\",\)"): shmem.put(buffer, to_rank=peer) - with pytest.raises(RuntimeError, match='copy_engine="sdma"'): + with pytest.raises(RuntimeError, match=r"engines=\(\"sdma\",\)"): shmem.quiet() From edd4d1d940750568cc60968e2c74e43dd74d608c Mon Sep 17 00:00:00 2001 From: "Prajapati, Dimple" Date: Tue, 13 Oct 2026 02:02:49 +0000 Subject: [PATCH 7/8] Install the sdma extra in the SLURM core-test script run_core_tests.sh runs every tests//test_*.py, including the SDMA tests, which skip without rocm-xio. Match run_example_slurm.sh. Signed-off-by: Prajapati, Dimple --- scripts/run_core_tests_slurm.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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" From 638f3d097a3bae9c401624f5b810e68905d0166a Mon Sep 17 00:00:00 2001 From: "Prajapati, Dimple" Date: Tue, 13 Oct 2026 06:01:43 +0000 Subject: [PATCH 8/8] Trim the default-SDMA checks to what the suite does not already cover. Shader stores already run on a default Iris context. Keep the import check and the host put/quiet rejection, and drop the extra copy-engine assert from the device-context test. Signed-off-by: Prajapati, Dimple --- tests/unittests/test_device_context.py | 1 - tests/unittests/test_optional_sdma.py | 27 +++----------------------- 2 files changed, 3 insertions(+), 25 deletions(-) diff --git a/tests/unittests/test_device_context.py b/tests/unittests/test_device_context.py index 5c628da31..874afa102 100644 --- a/tests/unittests/test_device_context.py +++ b/tests/unittests/test_device_context.py @@ -478,7 +478,6 @@ def test_device_context_tracing_1d_address(): def test_device_context_initialize(): """Test DeviceContext.initialize() creates valid context.""" ctx = iris.iris(1 << 20) - assert ctx.get_copy_engine_ctx() is None cur_rank = ctx.get_rank() num_ranks = ctx.get_num_ranks() diff --git a/tests/unittests/test_optional_sdma.py b/tests/unittests/test_optional_sdma.py index fae409477..b6c91659d 100644 --- a/tests/unittests/test_optional_sdma.py +++ b/tests/unittests/test_optional_sdma.py @@ -8,8 +8,6 @@ import pytest import torch -import triton -import triton.language as tl import iris @@ -22,32 +20,13 @@ def test_import_iris_does_not_import_xio(): subprocess.check_call([sys.executable, "-c", code], env=env) -@triton.jit -def _store_to_peer(buffer, n_elements, cur_rank, peer_rank, heap_bases, BLOCK_SIZE: tl.constexpr): - offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - values = tl.full([BLOCK_SIZE], 1, tl.float32) * (cur_rank + 1) - iris.store(buffer + offsets, values, cur_rank, peer_rank, heap_bases, mask=mask) - - -def test_default_shader_store_without_sdma(): +def test_default_context_rejects_host_sdma(): shmem = iris.iris(1 << 20) assert shmem.get_copy_engine_ctx() is None rank = shmem.get_rank() - world_size = shmem.get_num_ranks() - peer = (rank + 1) % world_size - n_elements = 4096 - buffer = shmem.zeros(n_elements, dtype=torch.float32) - shmem.barrier() - - _store_to_peer[(triton.cdiv(n_elements, 1024),)]( - buffer, n_elements, rank, peer, shmem.get_heap_bases(), BLOCK_SIZE=1024 - ) - shmem.barrier() - - writer = (rank - 1) % world_size - torch.testing.assert_close(buffer, torch.full_like(buffer, writer + 1)) + 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)