diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 000000000..4ba30cc66 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,15 @@ +# The Docker build context is the repo root (see .github/scripts/container_build.sh) +# so that the Dockerfile can COPY in .github/scripts/install_rocshmem.sh. Nothing +# else in the tree is needed at image build time, and .git alone is tens of MB. +.git +.github/workflows +**/__pycache__ +**/*.pyc +.pytest_cache +*.egg-info +build +dist +docs +examples +tests +.claude diff --git a/.github/scripts/container_build.sh b/.github/scripts/container_build.sh index 5e8bda7bd..83583e57d 100755 --- a/.github/scripts/container_build.sh +++ b/.github/scripts/container_build.sh @@ -89,8 +89,11 @@ elif [ "$CONTAINER_RUNTIME" = "docker" ]; then echo "[INFO] Using existing Docker image: $IMAGE_NAME" else echo "[INFO] Docker image $IMAGE_NAME not found, building..." - DOCKER_DIR="$(dirname "$(realpath "$0")")/../../docker" - if docker build -t "$IMAGE_NAME" "$DOCKER_DIR"; then + REPO_ROOT="$(dirname "$(realpath "$0")")/../.." + # Build from the repo root, not docker/, so the Dockerfile can COPY in + # .github/scripts/install_rocshmem.sh -- the same installer the Apptainer + # def file pulls in via %files. A docker/-only context cannot see it. + if docker build -t "$IMAGE_NAME" -f "$REPO_ROOT/docker/Dockerfile" "$REPO_ROOT"; then echo "[INFO] Built Docker image: $IMAGE_NAME" else echo "[ERROR] Docker build failed" diff --git a/.github/scripts/install_rocshmem.sh b/.github/scripts/install_rocshmem.sh new file mode 100755 index 000000000..e2f510e56 --- /dev/null +++ b/.github/scripts/install_rocshmem.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# Install rocSHMEM and its Python bindings, for tests/unittests/test_rocshmem_provider.py. +# +# Without this the provider tests skip rather than run: rocshmem4py has no +# prebuilt wheel anywhere (not PyPI, not the ROCm wheel indexes), and its +# bindings do not build rocSHMEM themselves -- python/rocshmem does +# find_package(rocshmem 3.5.0 CONFIG REQUIRED) with no FetchContent. So rocSHMEM +# has to be built first, and then the bindings against it. +# +# Scope is deliberately IPC-only: that is what the provider uses, and it is all +# a single-node CI runner can exercise. Upstream already defaults USE_IPC=ON and +# USE_GDA=OFF, so no conduit flags are passed -- which also keeps MPI and the +# RDMA provider libraries out of the picture entirely. +set -euo pipefail + +ROCSHMEM_PREFIX="${ROCSHMEM_PREFIX:-/opt/rocshmem}" +# MI325X runners are gfx942. Semicolon-separated for more than one. +ROCSHMEM_GPU_TARGETS="${ROCSHMEM_GPU_TARGETS:-gfx942}" +ROCSHMEM_REPO="${ROCSHMEM_REPO:-https://github.com/ROCm/rocm-systems.git}" +ROCSHMEM_REF="${ROCSHMEM_REF:-develop}" +ROCM_PATH="${ROCM_PATH:-/opt/rocm}" +SRC="$(mktemp -d)" + +# rocSHMEM is built from source because the CI bases are ROCm 7.2.1 (apptainer) +# and 7.1 (docker). ROCm 7.14 artifacts onward ship rocSHMEM's static library and +# headers, so once a base image is that new this build can be dropped and only +# the bindings below are needed. rocshmem4py has to be built either way until its +# TheRock packaging lands. +echo "==> rocSHMEM ${ROCSHMEM_REF} -> ${ROCSHMEM_PREFIX} (GPU_TARGETS=${ROCSHMEM_GPU_TARGETS})" + +# rocm-systems is a large monorepo and we need two directories out of it. Sparse +# checkout keeps this from dominating image build time and size. +git clone --depth 1 --branch "${ROCSHMEM_REF}" --filter=blob:none --sparse \ + "${ROCSHMEM_REPO}" "${SRC}" +git -C "${SRC}" sparse-checkout set projects/rocshmem python/rocshmem + +[ -f "${SRC}/projects/rocshmem/CMakeLists.txt" ] || { + echo "ERROR: projects/rocshmem missing after sparse checkout" >&2; exit 1; } + +# rocSHMEM's cmake/setup_project.cmake does a REQUIRED find_file for +# .info/version under ROCM_PATH. Images that lack that file fail to configure +# with "Could not find rocm_version_file", so use the documented escape hatch and +# read the version from rocm_version.h, which is authoritative. hipconfig +# --version is not used: it reports a build number that parses as the patch level. +EXPLICIT_ROCM_VERSION="${EXPLICIT_ROCM_VERSION:-}" +if [ -z "${EXPLICIT_ROCM_VERSION}" ] && [ ! -f "${ROCM_PATH}/.info/version" ]; then + _vh="$(find "${ROCM_PATH}" -name rocm_version.h 2>/dev/null | head -1)" + if [ -n "${_vh}" ]; then + EXPLICIT_ROCM_VERSION="$(awk ' + /ROCM_VERSION_MAJOR/ {maj=$3} /ROCM_VERSION_MINOR/ {min=$3} + /ROCM_VERSION_PATCH/ {pat=$3} + END {if (maj != "") printf "%s.%s.%s", maj, min, pat}' "${_vh}")" + echo "==> ROCm ${EXPLICIT_ROCM_VERSION} detected from ${_vh}" + fi +fi + +cmake -S "${SRC}/projects/rocshmem" -B "${SRC}/build" -G Ninja \ + ${EXPLICIT_ROCM_VERSION:+-DEXPLICIT_ROCM_VERSION="${EXPLICIT_ROCM_VERSION}"} \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_PREFIX}" \ + -DCMAKE_POSITION_INDEPENDENT_CODE=ON \ + -DGPU_TARGETS="${ROCSHMEM_GPU_TARGETS}" +cmake --build "${SRC}/build" --parallel "$(nproc)" +cmake --install "${SRC}/build" + +# USE_IPC must be ON or rocshmem_ptr returns NULL for every peer and the provider +# refuses to build a table. It is the upstream default, so this asserts rather +# than sets it -- a silent flip would otherwise surface much later as skipped tests. +if ! grep -qi "define ROCSHMEM_USE_IPC\|USE_IPC" \ + "${ROCSHMEM_PREFIX}"/include/rocshmem/*.hpp 2>/dev/null; then + echo "==> note: could not confirm USE_IPC from headers; provider will report at run time" +fi + +# A pip install from source, same as the one-liner in iris/experimental/README.md +# but pointed at the checkout above instead of a git+ URL. That is deliberate: a +# git+ URL makes pip clone the monorepo again, independently, at whatever HEAD +# develop happens to be at -- so the bindings could be built from a different +# revision than the core installed above. find_package would not catch it, since +# it only compares versions, and the bindings statically link the core. One +# checkout for both makes the skew impossible, and saves a second clone. +# +# CMAKE_PREFIX_PATH is the documented way to point the bindings at an install; +# setup.py forwards it to CMake as a cache variable so a rocSHMEM shipped under +# /opt/rocm cannot shadow it. ROCSHMEM_HOME is no longer required. +echo "==> building rocshmem4py against ${ROCSHMEM_PREFIX}" +CMAKE_PREFIX_PATH="${ROCSHMEM_PREFIX}" ROCM_PATH="${ROCM_PATH}" \ + pip3 install --no-cache-dir "${SRC}/python/rocshmem" + +python3 -c " +import rocshmem4py, importlib.metadata as md +print(' rocshmem4py', md.version('rocshmem4py'), '->', rocshmem4py.__file__) +for n in ('rocshmem_my_pe', 'rocshmem_n_pes', 'rocshmem_ptr'): + assert hasattr(rocshmem4py, n), f'missing {n}' +print(' provider API present') +" + +rm -rf "${SRC}" +echo "==> rocSHMEM install complete" diff --git a/apptainer/iris.def b/apptainer/iris.def index bea9e2d67..7a9d307ba 100644 --- a/apptainer/iris.def +++ b/apptainer/iris.def @@ -4,6 +4,13 @@ Bootstrap: docker From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0 +# The rocSHMEM installer is copied in rather than inlined, so it stays one +# implementation shared with docker/Dockerfile. Caveat: container_build.sh caches +# the image on the checksum of THIS def file only, so editing +# install_rocshmem.sh alone reuses a stale image -- touch this file too. +%files + .github/scripts/install_rocshmem.sh /opt/install_rocshmem.sh + %post /bin/bash -c " # Set environment variables @@ -37,6 +44,13 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0 # Make the venv writable by all chmod -R 777 /opt/venv + + # rocSHMEM + rocshmem4py, for tests/unittests/test_rocshmem_provider.py. + # Without it those tests skip: rocshmem4py has no prebuilt wheel on any + # index, and its bindings need an existing rocSHMEM to build against. + # gfx942 matches the MI325X CI runners. + ROCSHMEM_GPU_TARGETS=gfx942 ROCSHMEM_PREFIX=/opt/rocshmem \ + bash /opt/install_rocshmem.sh " %environment @@ -52,6 +66,7 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0 export OMPI_ALLOW_RUN_AS_ROOT=1 # Set required RCCL environment variable for ROCm export HSA_NO_SCRATCH_RECLAIM=1 + export ROCSHMEM_PREFIX=/opt/rocshmem %runscript echo "Welcome to the ROCm-aware Apptainer image!" diff --git a/docker/Dockerfile b/docker/Dockerfile index 04126529a..944b5c254 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -47,6 +47,14 @@ RUN git checkout bcbcabdd0cff6539c7168299075992b2a23ff38e RUN pip3 install -e . ENV PYTHONPATH=$TRITON_PATH +# rocSHMEM + rocshmem4py, for tests/unittests/test_rocshmem_provider.py. Without +# it those tests skip: there is no prebuilt rocshmem4py wheel on any index, and +# its bindings require an existing rocSHMEM install to build against. +# gfx942 matches the MI325X CI runners; override ROCSHMEM_GPU_TARGETS for others. +ENV ROCSHMEM_PREFIX=/opt/rocshmem +COPY .github/scripts/install_rocshmem.sh /tmp/install_rocshmem.sh +RUN ROCSHMEM_GPU_TARGETS=gfx942 bash /tmp/install_rocshmem.sh && rm /tmp/install_rocshmem.sh + # Set up workspace WORKDIR /workspace diff --git a/iris/experimental/README.md b/iris/experimental/README.md new file mode 100644 index 000000000..dc0b543cd --- /dev/null +++ b/iris/experimental/README.md @@ -0,0 +1,123 @@ +# Experimental allocation providers + +Adapters that let Iris device kernels run on memory allocated by another +runtime. Iris device code is unchanged: `iris.store`, `load` and `copy` take +`heap_bases` as a plain pointer argument and translate with + +``` +remote = peer_bases[to] + (ptr - peer_bases[local_rank]) +``` + +so a provider's whole job is to hand Iris one `int64[num_ranks]` table per +allocation whose `local_rank` entry is that allocation's own base. + +Neither provider is imported by `iris/experimental/__init__.py`, so `import +iris` never requires either dependency. + +| Provider | Dependency | Scope | +| --- | --- | --- | +| `rocshmem_provider.py` | `rocshmem4py` | intra-node (IPC) | + +## rocSHMEM provider + +### Installing `rocshmem4py` + +There is no prebuilt wheel — `rocshmem4py` is not on PyPI, not in the ROCm +nightly wheel indexes, and the `rocm-systems` release assets are source +tarballs. But pip builds it from source in a single command, given a rocSHMEM +install to build against: + +```bash +CMAKE_PREFIX_PATH= pip install \ + "rocshmem4py @ git+https://github.com/ROCm/rocm-systems.git#subdirectory=python/rocshmem" +``` + +Verified against a rocSHMEM 3.7.0 install: builds and installs +`rocshmem4py-0.1.0+rocshmem3.7.0-cp312-cp312-linux_x86_64.whl` with no other +environment set. `ROCSHMEM_HOME` is accepted as a convenience but is not +required; `CMAKE_PREFIX_PATH` is the documented mechanism and takes precedence +over both it and `ROCM_PATH`. It is forwarded to CMake as a cache variable +specifically so a rocSHMEM shipped under `/opt/rocm` cannot shadow the one you +asked for. + +Two properties of the result are worth knowing: + +- It **contains** rocSHMEM rather than depending on it at run time — rocSHMEM is + statically linked into the extension module, and the version records which one + (`0.1.0+rocshmem3.7.0`). Nothing needs to be on `LD_LIBRARY_PATH` afterwards. +- Consequently rocSHMEM's **build options are fixed when `rocshmem4py` is + built**, not when it is used. + +The wheel is CPython-ABI-tagged (`cp312` above), so build it with the interpreter +that will run it. + +CI does the same pip install from source, but from a checkout it already has +rather than a `git+` URL — see `.github/scripts/install_rocshmem.sh`. Since it +has to build rocSHMEM itself anyway, taking both from one checkout keeps the core +and the bindings at the same revision; a `git+` URL would clone independently and +could drift, which `find_package` would not catch because it only compares +versions. If you are building both by hand, prefer the same: point +`pip install` at your `python/rocshmem` directory rather than at the URL. + +### Building rocSHMEM first + +The bindings do not build rocSHMEM: `CMakeLists.txt` does +`find_package(rocshmem 3.5.0 CONFIG REQUIRED)` with no `FetchContent`, so +**rocSHMEM 3.5.0 or newer must already be installed**. + +**`USE_IPC` must be `ON`.** This provider gets its peer addresses from +`rocshmem_ptr`, which returns NULL unconditionally when IPC is compiled out. The +provider raises with that hint if every peer comes back NULL, rather than handing +Iris a table of zeros that would become wild pointers inside a kernel. Upstream +defaults it ON. + +```bash +cmake -S "$ROCSHMEM_SRC" -B "$BUILD" -G Ninja \ + -DCMAKE_INSTALL_PREFIX="$ROCSHMEM_HOME" \ + -DCMAKE_POSITION_INDEPENDENT_CODE=ON \ + -DGPU_TARGETS=gfx950 \ + -DUSE_IPC=ON +cmake --build "$BUILD" --parallel +cmake --install "$BUILD" +``` + +If `GPU_TARGETS` is rejected as `invalid offload arch combinations: 'gfx950' and +'gfx950:sramecc+:xnack-'`, rocSHMEM's auto-detected arch and the one +`find_package(hip)` derives from the build host disagree; set +`ROCSHMEM_GPU_TARGETS='gfx950:sramecc+:xnack-'` to match. + +### Verifying + +```bash +python -c "import rocshmem4py; print(rocshmem4py.__file__)" +python tests/run_tests_distributed.py \ + tests/unittests/test_rocshmem_provider.py --num_ranks 2 -v +``` + +The tests skip rather than fail when `rocshmem4py` is absent, when fewer than 2 +ranks are present, or when peers are not directly addressable, so they are inert +in an environment without rocSHMEM. + +`tests/manual_rocshmem_provider.py` covers what the unit tests structurally +cannot: `run_tests_distributed.py` launches `torchrun` with `--nnodes=1`, so the +unit tests only ever see intra-node peers, where `rocshmem_ptr` resolves every +one. The manual script exercises the multi-node case, where `rocshmem_ptr` +returns NULL for remote peers and `SymmetricAddressMap.direct` is the thing under +test (`EXPECT_INDIRECT=1`). + +### Running + +rocSHMEM must be initialised before the provider is constructed; the caller owns +bootstrap and tensor lifetime: + +```python +dist.init_process_group(backend="gloo") +rocshmem4py.init_rocshmem_by_uniqueid(dist.group.WORLD) +provider = RocshmemProvider() +``` + +Allocation and free are both **collective** — `rocshmem_free` is documented as +"a collective operation and must be called by all PEs" — so every rank must make +the same calls in the same order. That is why `free()` is explicit rather than +driven by garbage collection: `__del__` would run at whatever moment each rank +happened to collect, and ranks would hang instead of raising. diff --git a/iris/experimental/rocshmem_provider.py b/iris/experimental/rocshmem_provider.py new file mode 100644 index 000000000..c7a4cd341 --- /dev/null +++ b/iris/experimental/rocshmem_provider.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""rocSHMEM as an allocation provider for Iris device kernels. + +Lets Iris device code operate on tensors allocated by rocSHMEM rather than from +Iris's own symmetric heap. No Iris device code changes are needed: iris.store, +load and copy take ``heap_bases`` as a plain pointer argument and translate with + + remote = peer_bases[to] + (ptr - peer_bases[local_rank]) + +so any table satisfying ``peer_bases[local_rank] == local allocation base`` +drives them. This module builds that table for rocSHMEM memory. + +The table comes from ``rocshmem_ptr(base, peer)``, OpenSHMEM's ``shmem_ptr``: an +address in this process's own address space for the peer's counterpart of a +symmetric object, or NULL when that peer is not reachable by direct load/store. + +The per-peer offsets are queried once. rocSHMEM's peer mapping is a single +linear translation of the whole symmetric heap, so the offset from a local +address to its counterpart on a given peer is the same constant everywhere in +the heap. Only those offsets are cached; each allocation's table is materialised +from its own base, so ``peer_bases[local_rank]`` is always that allocation's +base. rocSHMEM's heap base, which it does not expose publicly, is never needed. + +Because the offsets are shared, a table built for one allocation still +translates pointers belonging to another. iris.copy relies on that: it takes one +``heap_bases`` and translates two pointers against it. + +Scope is intra-node. A peer not reachable by direct load/store gets a base of 0, +which would translate to a wild pointer rather than an error, so +``SymmetricAddressMap.direct`` records reachability per peer and callers are +expected to check it before launching. Inter-node peers need a transport this +module does not provide. + +TODO: settle where this belongs. It sits in Iris on the assumption that Iris +hosts provider adapters; the alternative is for it to live alongside rocSHMEM, +which owns the allocation and the tensor lifetime. It is one file either way. + +This module is deliberately NOT imported by ``iris/experimental/__init__.py``, +so ``import iris`` does not require rocshmem4py. Keep it that way: adding it to +that package's eager imports would make a rocSHMEM install mandatory for every +Iris user. + +The dependency is ``rocshmem4py``, a standalone Python package from the +ROCm/rocm-systems repository rather than something a ROCm install provides. It +does not link rocSHMEM at run time; it statically links it into its extension +module, and its version records which rocSHMEM that was (e.g. +``0.1.0+rocshmem3.7.0``). So installing it needs no separate rocSHMEM on the +system, and the rocSHMEM build options it was compiled with -- ``USE_IPC`` in +particular -- are fixed at its build time, not selectable later. + +The caller owns bootstrap and tensor lifetime; rocSHMEM must already be +initialised: + + dist.init_process_group(backend="gloo") + rocshmem4py.init_rocshmem_by_uniqueid(dist.group.WORLD) + provider = RocshmemProvider() +""" +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +import rocshmem4py as rshmem +from rocshmem4py.interop import torch as rshmem_torch + + +@dataclass(frozen=True) +class SymmetricAddressMap: + """Address metadata for one symmetric allocation. + + ``allocate_symmetric`` returns only ``(tensor, peer_bases)``; this carries + what that pair cannot, notably ``direct``. + """ + + peer_bases: torch.Tensor # int64[world_size], device-resident + local_rank: int + allocation_base: int + allocation_bytes: int + direct: tuple[bool, ...] # per peer: reachable by load/store? + + def all_direct(self) -> bool: + return all(self.direct) + + def indirect_peers(self) -> list[int]: + return [r for r, d in enumerate(self.direct) if not d] + + +class RocshmemProvider: + """Allocates rocSHMEM symmetric tensors and describes them for Iris kernels.""" + + def __init__(self, device: str | None = None): + self.cur_rank = rshmem.rocshmem_my_pe() + self.num_ranks = rshmem.rocshmem_n_pes() + self._device = device + # peer -> byte offset from a local address to its counterpart on that + # peer, or None when the peer is not reachable by load/store. Constant + # across the heap, so it is computed once from the first allocation. + self._deltas: list[int | None] | None = None + + # ── table form ─────────────────────────────────────────────────────────── + + def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Tensor]: + """Allocate a symmetric tensor and return it with its peer-base table. + + Returns the provider-facing shape symmetric allocation is converging on, + ``(tensor, peer_bases)``, so the same device kernels drive any provider. + + Collective: rocSHMEM allocation is, so every PE must call this the same + number of times and in the same order. + + Returns ``(tensor, peer_bases)`` where ``peer_bases`` is an + ``int64[num_ranks]`` tensor on the same device as ``tensor``, holding for + each peer the address of that peer's counterpart of this allocation, in + this process's address space. Its ``local_rank`` entry is this + allocation's own base, which is what Iris translation subtracts. A peer + not reachable by direct load/store is 0; ``allocate_symmetric_map`` + returns the same thing plus the ``direct`` mask that says which, and + callers that may run inter-node should check it rather than launching + against a 0. + + Hold the returned table for as long as the allocation lives rather than + re-deriving it per launch; it is built once here and does not change. + """ + tensor, amap = self.allocate_symmetric_map(*size, dtype=dtype) + return tensor, amap.peer_bases + + # ── descriptor form ────────────────────────────────────────────────────── + + def allocate_symmetric_map(self, *size, dtype=None) -> tuple[torch.Tensor, SymmetricAddressMap]: + """As allocate_symmetric, but returning the full address descriptor.""" + shape = tuple(size[0]) if len(size) == 1 and hasattr(size[0], "__iter__") else tuple(size) + dtype = dtype or torch.get_default_dtype() + tensor = rshmem_torch.create_tensor(shape, dtype) + return tensor, self.symmetric_address_map(tensor) + + def symmetric_address_map(self, tensor: torch.Tensor) -> SymmetricAddressMap: + """Describe an already-allocated rocSHMEM tensor. + + Materialises a fresh ``int64[num_ranks]`` table on each call. That is one + small device tensor per allocation on the normal path, since + ``allocate_symmetric`` calls this once; it is not meant to be called per + kernel launch. The table is not memoised on purpose: keying a cache by + ``data_ptr()`` would alias once an allocation is freed and its address + reused, and the result would be a silently wrong table rather than an + error. + """ + base = tensor.data_ptr() + deltas = self._peer_deltas(tensor) + bases = [0 if d is None else base + d for d in deltas] + return SymmetricAddressMap( + peer_bases=torch.tensor(bases, dtype=torch.int64, device=tensor.device), + local_rank=self.cur_rank, + allocation_base=base, + allocation_bytes=tensor.numel() * tensor.element_size(), + direct=tuple(d is not None for d in deltas), + ) + + def _peer_deltas(self, anchor: torch.Tensor) -> list[int | None]: + """Per-peer byte offsets, queried once and reused. + + rocshmem_ptr is a linear translation of the whole symmetric heap, so the + offset to a peer's counterpart is the same for every address in it. Only + the offsets are cached; each allocation's table is materialised from its + own base, which keeps peer_bases[local_rank] == that allocation's base. + """ + if self._deltas is not None: + return self._deltas + + base = anchor.data_ptr() + deltas: list[int | None] = [] + for peer in range(self.num_ranks): + if peer == self.cur_rank: + deltas.append(0) + continue + p = int(rshmem.rocshmem_ptr(base, peer)) + deltas.append(p - base if p else None) + + # Every peer unreachable usually means rocSHMEM was built with + # USE_IPC=OFF rather than that every peer is remote: with IPC compiled + # out rocshmem_ptr returns NULL unconditionally. Upstream defaults it + # ON. Failing here beats handing back a table whose zeros would + # translate to wild pointers inside a kernel. + peers = [r for r in range(self.num_ranks) if r != self.cur_rank] + if peers and all(deltas[r] is None for r in peers): + raise RuntimeError( + "rocshmem_ptr returned NULL for every peer. If any peer shares " + "this node, rocSHMEM was likely built with USE_IPC=OFF (upstream " + "defaults ON); check the USE_IPC line in the rocSHMEM banner." + ) + + self._deltas = deltas + return deltas + + # ── convenience ────────────────────────────────────────────────────────── + + def barrier(self): + rshmem_torch.barrier_all() + + def free(self, tensor: torch.Tensor): + """Release a symmetric allocation. Collective. + + Explicit by necessity, not by preference. rocshmem_free is documented as + "a collective operation and must be called by all PEs", so it cannot be + driven from ``__del__`` or a weakref finalizer: Python decides when to + collect per process, and ranks that collect in different orders, or at + different times, would diverge and hang instead of raising. Freeing has + to stay where the caller can order it across ranks. + """ + rshmem_torch.free_tensor(tensor) + + def get_rank(self) -> int: + return self.cur_rank + + def get_num_ranks(self) -> int: + return self.num_ranks diff --git a/tests/manual_rocshmem_provider.py b/tests/manual_rocshmem_provider.py new file mode 100644 index 000000000..6addfb8b7 --- /dev/null +++ b/tests/manual_rocshmem_provider.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""Iris device kernels driving rocSHMEM-allocated buffers. + +Intra-node (IPC). Run on one node with 2+ ranks: + + torchrun --nproc_per_node=2 tests/manual_rocshmem_provider.py + +Set EXPECT_INDIRECT=1 and run across 2 nodes to check that peers which are not +directly addressable are reported rather than translated. +""" + +import os +import sys + +import torch +import torch.distributed as dist +import triton +import triton.language as tl + +import iris +import rocshmem4py as rshmem + +from iris.experimental.rocshmem_provider import RocshmemProvider + +BLOCK_SIZE = 1024 + + +@triton.jit +def _broadcast_kernel( + data, + results, + peer_bases, + n_elements, + cur_rank, + num_ranks: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Push this rank's values into `results` on every rank. + + Deliberately identical in shape to tests/unittests/test_store_triton.py: + the whole point is that this is ordinary Iris device code, unaware that + `peer_bases` came from rocSHMEM rather than an Iris heap. + """ + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + value = tl.load(data + offsets, mask=mask) + for dst_rank in range(num_ranks): + iris.store(results + offsets, value, cur_rank, dst_rank, peer_bases, mask=mask) + + +def main(): + dist.init_process_group(backend="gloo") + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", "0"))) + rshmem.init_rocshmem_by_uniqueid(dist.group.WORLD) + + provider = RocshmemProvider() + me, ws = provider.get_rank(), provider.get_num_ranks() + assert ws >= 2, "need at least 2 ranks" + + # Two allocations from a non-Iris allocator. + data, data_bases = provider.allocate_symmetric(BLOCK_SIZE, dtype=torch.float32) + results, peer_bases = provider.allocate_symmetric(BLOCK_SIZE, dtype=torch.float32) + + amap = provider.symmetric_address_map(results) + print(f"[rank{me}] direct={amap.direct} base={amap.allocation_base:#x} " + f"bases={[hex(int(b)) for b in peer_bases.tolist()]}", flush=True) + + # A non-direct peer's base is 0, which would translate to a wild pointer + # rather than error, so refuse instead. EXPECT_INDIRECT=1 tests that path. + if os.environ.get("EXPECT_INDIRECT") == "1": + detected = not amap.all_direct() + print(f"[rank{me}] EXPECT_INDIRECT: indirect peers={amap.indirect_peers()} " + f"detected={detected}", flush=True) + res = [None] * ws + dist.all_gather_object(res, detected) + if me == 0: + print("ROCSHMEM_PROVIDER_INDIRECT_RESULT:", + "PASS" if all(res) else "FAIL", flush=True) + provider.barrier() + provider.free(data) + provider.free(results) + dist.destroy_process_group() + return 0 + + assert amap.all_direct(), ( + f"[rank{me}] peers {amap.indirect_peers()} are not directly addressable; " + "this prototype is IPC-only -- run all ranks on one node") + + # The invariant the device code actually depends on. + assert int(peer_bases[me].item()) == results.data_ptr() + + # Rank 0 broadcasts its values; every rank should end up with them. + data.fill_(float(me + 1)) + results.fill_(-1.0) + torch.cuda.synchronize() + provider.barrier() + + if me == 0: + _broadcast_kernel[(1,)]( + data, results, peer_bases, BLOCK_SIZE, me, + num_ranks=ws, BLOCK_SIZE=BLOCK_SIZE, num_warps=4, + ) + torch.cuda.synchronize() + provider.barrier() + + want = 1.0 # rank 0's fill value + ok = bool(torch.allclose(results, torch.full_like(results, want))) + got = torch.unique(results)[:4].tolist() + print(f"[rank{me}] results want={want} got={got} match={ok}", flush=True) + + res = [None] * ws + dist.all_gather_object(res, ok) + if me == 0: + print("ROCSHMEM_PROVIDER_RESULT:", "PASS" if all(res) else "FAIL", flush=True) + + # Translate pointers in `results` using the table built from `data`: one + # table should be valid for every allocation. + provider.barrier() + results.fill_(-1.0) + torch.cuda.synchronize() + provider.barrier() + + if me == 0: + _broadcast_kernel[(1,)]( + data, results, data_bases, BLOCK_SIZE, me, # data's table, results' pointers + num_ranks=ws, BLOCK_SIZE=BLOCK_SIZE, num_warps=4, + ) + torch.cuda.synchronize() + provider.barrier() + + xok = bool(torch.allclose(results, torch.full_like(results, want))) + xgot = torch.unique(results)[:4].tolist() + print(f"[rank{me}] cross-alloc want={want} got={xgot} match={xok}", flush=True) + + xres = [None] * ws + dist.all_gather_object(xres, xok) + if me == 0: + print("ROCSHMEM_CROSS_ALLOC_RESULT:", "PASS" if all(xres) else "FAIL", flush=True) + + provider.barrier() + provider.free(data) + provider.free(results) + dist.destroy_process_group() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unittests/test_rocshmem_provider.py b/tests/unittests/test_rocshmem_provider.py new file mode 100644 index 000000000..62107cb55 --- /dev/null +++ b/tests/unittests/test_rocshmem_provider.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""Iris device kernels driving rocSHMEM-allocated memory. + +Run under the usual launcher, which sets up torch.distributed and the device: + + python tests/run_tests_distributed.py tests/unittests/test_rocshmem_provider.py \ + --num_ranks 2 -v + +Skips when rocshmem4py is absent, when fewer than 2 ranks are present, or when +peers are not directly addressable, so it is inert rather than failing in a +normal CI run. tests/manual_rocshmem_provider.py covers the multi-node case. +""" + +import pytest +import torch +import torch.distributed as dist +import triton +import triton.language as tl + +import iris + +BLOCK_SIZE = 1024 + + +@triton.jit +def _broadcast_kernel(data, results, peer_bases, n_elements, cur_rank, + num_ranks: tl.constexpr, BLOCK_SIZE: tl.constexpr): + """Ordinary Iris device code -- unaware the table came from rocSHMEM.""" + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + value = tl.load(data + offsets, mask=mask) + for dst_rank in range(num_ranks): + iris.store(results + offsets, value, cur_rank, dst_rank, peer_bases, mask=mask) + + +@pytest.fixture(scope="module") +def provider(): + if not dist.is_initialized(): + pytest.skip("needs torch.distributed; run via tests/run_tests_distributed.py") + if dist.get_world_size() < 2: + pytest.skip("needs at least 2 ranks (--num_ranks 2)") + + # Imported here rather than at module scope so the tests are collected and + # individually skipped. A module-level importorskip collects zero items, + # which makes pytest exit 5 (NO_TESTS_COLLECTED) and fails the whole run. + rshmem = pytest.importorskip( + "rocshmem4py", reason="rocSHMEM provider tests need rocshmem4py installed" + ) + from iris.experimental.rocshmem_provider import RocshmemProvider + + # rocSHMEM initialises once per process, hence module scope. No finalize in + # teardown: it would pull the runtime out from under anything else running. + rshmem.init_rocshmem_by_uniqueid(dist.group.WORLD) + return RocshmemProvider() + + +@pytest.fixture +def symmetric_pair(provider): + data, peer_bases = provider.allocate_symmetric(BLOCK_SIZE, dtype=torch.float32) + results, _ = provider.allocate_symmetric(BLOCK_SIZE, dtype=torch.float32) + yield provider, data, results, peer_bases + provider.barrier() + provider.free(data) + provider.free(results) + + +def test_peer_bases_shape_and_invariant(symmetric_pair): + """The invariant Iris device code translates against.""" + provider, data, _results, peer_bases = symmetric_pair + ws = provider.get_num_ranks() + + assert peer_bases.numel() == ws + assert peer_bases.dtype == torch.int64 + assert peer_bases.is_cuda + # peer_bases[local_rank] is the base translation subtracts. + assert int(peer_bases[provider.get_rank()].item()) == data.data_ptr() + + +def test_peer_offsets_are_shared(provider): + """Each allocation gets its own table, built from shared per-peer offsets. + + Equal offsets are what make a table from one allocation able to translate + another's pointers, which iris.copy depends on. + """ + a, bases_a = provider.allocate_symmetric(64, dtype=torch.float32) + b, bases_b = provider.allocate_symmetric(64, dtype=torch.float32) + try: + assert int(bases_a[provider.get_rank()].item()) == a.data_ptr() + assert int(bases_b[provider.get_rank()].item()) == b.data_ptr() + direct = provider.symmetric_address_map(a).direct + for peer in range(provider.get_num_ranks()): + if not direct[peer]: + # Unreachable peers are 0 in every table, not base + offset. + assert int(bases_a[peer].item()) == 0 + assert int(bases_b[peer].item()) == 0 + continue + da = int(bases_a[peer].item()) - a.data_ptr() + db = int(bases_b[peer].item()) - b.data_ptr() + assert da == db, f"peer {peer}: offset {da} != {db}" + finally: + provider.barrier() + provider.free(a) + provider.free(b) + + +def test_address_map_reports_reachability(symmetric_pair): + """Per-peer reachability, and the 0 base that goes with it.""" + provider, _data, results, _peer_bases = symmetric_pair + amap = provider.symmetric_address_map(results) + ws = provider.get_num_ranks() + + assert len(amap.direct) == ws + assert amap.direct[provider.get_rank()], "a rank must be able to reach itself" + assert amap.allocation_base == results.data_ptr() + assert amap.allocation_bytes == results.numel() * results.element_size() + # A non-direct peer's base is 0. + for peer, is_direct in enumerate(amap.direct): + assert (int(amap.peer_bases[peer].item()) != 0) == is_direct + + +def test_iris_store_over_rocshmem_memory(symmetric_pair): + """Unmodified iris.store, on memory Iris did not allocate.""" + provider, data, results, peer_bases = symmetric_pair + me, ws = provider.get_rank(), provider.get_num_ranks() + + amap = provider.symmetric_address_map(results) + if not amap.all_direct(): + pytest.skip(f"peers {amap.indirect_peers()} are not directly addressable; " + "this path is intra-node only") + + data.fill_(float(me + 1)) + results.fill_(-1.0) + torch.cuda.synchronize() + provider.barrier() + + if me == 0: + _broadcast_kernel[(1,)](data, results, peer_bases, BLOCK_SIZE, me, + num_ranks=ws, BLOCK_SIZE=BLOCK_SIZE, num_warps=4) + torch.cuda.synchronize() + provider.barrier() + + # Rank 0 pushed its value to every rank, including this one. + assert torch.allclose(results, torch.full_like(results, 1.0))