From 87112af0e350e4f8af98641c2fb88408078abc9e Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Tue, 8 Sep 2026 11:13:15 -0500 Subject: [PATCH 01/16] Add rocSHMEM as an allocation provider for Iris device kernels Lets Iris device code operate on tensors allocated by rocSHMEM instead of from Iris's own symmetric heap. No Iris device code changes are required: store, load and copy take heap_bases as a plain pointer argument, so any table satisfying peer_bases[local_rank] == local allocation base drives them. iris/experimental/rocshmem_provider.py builds that table from rocshmem_ptr(base, peer), which returns an address in this process's own space for a peer's counterpart of a symmetric object, or NULL when that peer is not reachable by direct load/store. allocate_symmetric(*size, dtype) -> (tensor, peer_bases) allocate_symmetric_map(*size, dtype) -> (tensor, SymmetricAddressMap) symmetric_address_map(tensor) -> SymmetricAddressMap The first matches Iris.allocate_symmetric's shape so the same kernels drive either provider. The descriptor form adds local_rank, allocation_base, allocation_bytes and a per-peer `direct` mask; callers check that mask before launching, since a peer that is not directly addressable has a base of 0 and would translate to a wild pointer rather than an error. One table serves every allocation. rocSHMEM's peer mapping is a linear translation of the whole symmetric heap, so any symmetric address anchors a table valid for all allocations, and rocSHMEM's heap base -- which it does not expose publicly -- is never needed. That also keeps iris.copy usable, since it translates two pointers against a single heap_bases. Scope is intra-node. Inter-node peers are reported as unreachable rather than driven; they need a transport this module does not provide. Tests: tests/unittests/test_rocshmem_provider.py pytest under the repo launcher, skipping when rocshmem4py is absent, when fewer than 2 ranks are present, or when peers are not directly addressable tests/manual_rocshmem_provider.py multi-node script, including the non-addressable-peer path via EXPECT_INDIRECT=1 The provider module is not imported by iris/experimental/__init__.py, so `import iris` does not require rocshmem4py. Co-Authored-By: Claude Opus 5 (1M context) --- iris/experimental/rocshmem_provider.py | 171 ++++++++++++++++++++++ tests/manual_rocshmem_provider.py | 150 +++++++++++++++++++ tests/unittests/test_rocshmem_provider.py | 130 ++++++++++++++++ 3 files changed, 451 insertions(+) create mode 100644 iris/experimental/rocshmem_provider.py create mode 100644 tests/manual_rocshmem_provider.py create mode 100644 tests/unittests/test_rocshmem_provider.py diff --git a/iris/experimental/rocshmem_provider.py b/iris/experimental/rocshmem_provider.py new file mode 100644 index 000000000..c94f16fb2 --- /dev/null +++ b/iris/experimental/rocshmem_provider.py @@ -0,0 +1,171 @@ +# 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. + +One table serves every allocation. rocSHMEM's peer mapping is a single linear +translation of the whole symmetric heap, so the offset between a local address +and its counterpart on a given peer is the same constant everywhere in the heap, +whatever allocation it belongs to. Any symmetric address therefore anchors a +table valid for all of them -- which also means rocSHMEM's heap base, which it +does not expose publicly, is never needed. That property matters because +iris.copy takes one ``heap_bases`` and translates two pointers against it; a +provider handing out per-allocation tables could not drive 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 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 or f"cuda:{torch.cuda.current_device()}" + self._context_bases: torch.Tensor | 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. + + Same signature and return shape as Iris.allocate_symmetric, so the same + device kernels drive either provider. + + The table is context-wide: it is built once from the first symmetric + allocation and shared by every later one. See the module docstring for + why a single anchor suffices, and test_table_is_context_wide for the + check that it holds. + """ + tensor, _ = self.allocate_symmetric_map(*size, dtype=dtype) + return tensor, self.context_peer_bases(tensor) + + # ── 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. + + Yields both the base table and, from the same call, rocSHMEM's own + answer to whether each peer is reachable by direct load/store. + """ + base = tensor.data_ptr() + bases, direct = [], [] + for peer in range(self.num_ranks): + p = base if peer == self.cur_rank else int(rshmem.rocshmem_ptr(base, peer)) + bases.append(p) + direct.append(p != 0) + + # An all-zero table (bar our own entry) almost always 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 USE_IPC=ON. Failing here beats handing back a table whose + # zeros translate to wild pointers inside a kernel. + peers = [r for r in range(self.num_ranks) if r != self.cur_rank] + if peers and not any(direct[r] 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." + ) + + return SymmetricAddressMap( + peer_bases=torch.tensor(bases, dtype=torch.int64, device=self.device), + local_rank=self.cur_rank, + allocation_base=base, + allocation_bytes=tensor.numel() * tensor.element_size(), + direct=tuple(direct), + ) + + def context_peer_bases(self, anchor: torch.Tensor) -> torch.Tensor: + """One peer-base table valid for every symmetric allocation. + + Built from the first symmetric tensor seen and cached. See + allocate_symmetric for why a single anchor suffices. + """ + if self._context_bases is None: + self._context_bases = self.symmetric_address_map(anchor).peer_bases + return self._context_bases + + # ── convenience ────────────────────────────────────────────────────────── + + def barrier(self): + rshmem_torch.barrier_all() + + def free(self, tensor: torch.Tensor): + 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..b0739a8fb --- /dev/null +++ b/tests/unittests/test_rocshmem_provider.py @@ -0,0 +1,130 @@ +# 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_table_is_context_wide(provider): + """A table built from one allocation is valid for any other.""" + a, bases_a = provider.allocate_symmetric(64, dtype=torch.float32) + b, bases_b = provider.allocate_symmetric(64, dtype=torch.float32) + try: + assert torch.equal(bases_a, bases_b) + 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)) From 14a5fd9d4c3dcc9d85228771b53aa865467e45ef Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Tue, 8 Sep 2026 16:18:27 -0500 Subject: [PATCH 02/16] Materialise each allocation's peer table from cached per-peer offsets allocate_symmetric() cached the table built from the first allocation and returned it for every later one. Translation still worked, because the per-peer offset is constant across the heap, but peer_bases[local_rank] was the first allocation's base rather than the current tensor's -- contradicting the stated invariant and breaking the assertion in tests/manual_rocshmem_provider.py, which checks it on a second allocation. Cache the per-peer offsets instead and build each allocation's table from its own base. peer_bases[local_rank] is now that allocation's base for every allocation, while the shared offsets keep a table from one allocation able to translate another's pointers, which iris.copy relies on. This also removes a duplicated rocshmem_ptr sweep: the first allocation previously queried every peer twice, once to build the map and once to seed the cache. It is now queried once per process. peer_bases is created on tensor.device rather than a device captured when the provider was constructed, so the table cannot end up on a different device than the memory it describes. test_table_is_context_wide asserted the two tables were equal, which no longer holds and was the weaker property anyway. It is now test_peer_offsets_are_shared and checks what actually matters: each table's local entry is its own allocation's base, and the per-peer offsets agree. Unreachable peers are excluded from that comparison, since their entry is 0 rather than base + offset. Verified on 2 ranks: pytest 4 passed, manual test PASS including the cross-allocation check. Co-Authored-By: Claude Opus 5 (1M context) --- iris/experimental/rocshmem_provider.py | 103 ++++++++++++---------- tests/unittests/test_rocshmem_provider.py | 21 ++++- 2 files changed, 72 insertions(+), 52 deletions(-) diff --git a/iris/experimental/rocshmem_provider.py b/iris/experimental/rocshmem_provider.py index c94f16fb2..870961b07 100644 --- a/iris/experimental/rocshmem_provider.py +++ b/iris/experimental/rocshmem_provider.py @@ -16,14 +16,16 @@ 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. -One table serves every allocation. rocSHMEM's peer mapping is a single linear -translation of the whole symmetric heap, so the offset between a local address -and its counterpart on a given peer is the same constant everywhere in the heap, -whatever allocation it belongs to. Any symmetric address therefore anchors a -table valid for all of them -- which also means rocSHMEM's heap base, which it -does not expose publicly, is never needed. That property matters because -iris.copy takes one ``heap_bases`` and translates two pointers against it; a -provider handing out per-allocation tables could not drive it. +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 @@ -84,8 +86,11 @@ class RocshmemProvider: def __init__(self, device: str | None = None): self.cur_rank = rshmem.rocshmem_my_pe() self.num_ranks = rshmem.rocshmem_n_pes() - self.device = device or f"cuda:{torch.cuda.current_device()}" - self._context_bases: torch.Tensor | None = None + 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 ─────────────────────────────────────────────────────────── @@ -94,14 +99,9 @@ def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Ten Same signature and return shape as Iris.allocate_symmetric, so the same device kernels drive either provider. - - The table is context-wide: it is built once from the first symmetric - allocation and shared by every later one. See the module docstring for - why a single anchor suffices, and test_table_is_context_wide for the - check that it holds. """ - tensor, _ = self.allocate_symmetric_map(*size, dtype=dtype) - return tensor, self.context_peer_bases(tensor) + tensor, amap = self.allocate_symmetric_map(*size, dtype=dtype) + return tensor, amap.peer_bases # ── descriptor form ────────────────────────────────────────────────────── @@ -113,48 +113,53 @@ def allocate_symmetric_map(self, *size, dtype=None) -> tuple[torch.Tensor, Symme return tensor, self.symmetric_address_map(tensor) def symmetric_address_map(self, tensor: torch.Tensor) -> SymmetricAddressMap: - """Describe an already-allocated rocSHMEM tensor. + """Describe an already-allocated rocSHMEM tensor.""" + 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), + ) - Yields both the base table and, from the same call, rocSHMEM's own - answer to whether each peer is reachable by direct load/store. + 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. """ - base = tensor.data_ptr() - bases, direct = [], [] + if self._deltas is not None: + return self._deltas + + base = anchor.data_ptr() + deltas: list[int | None] = [] for peer in range(self.num_ranks): - p = base if peer == self.cur_rank else int(rshmem.rocshmem_ptr(base, peer)) - bases.append(p) - direct.append(p != 0) - - # An all-zero table (bar our own entry) almost always 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 USE_IPC=ON. Failing here beats handing back a table whose - # zeros translate to wild pointers inside a kernel. + 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 not any(direct[r] for r in peers): + 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." ) - return SymmetricAddressMap( - peer_bases=torch.tensor(bases, dtype=torch.int64, device=self.device), - local_rank=self.cur_rank, - allocation_base=base, - allocation_bytes=tensor.numel() * tensor.element_size(), - direct=tuple(direct), - ) - - def context_peer_bases(self, anchor: torch.Tensor) -> torch.Tensor: - """One peer-base table valid for every symmetric allocation. - - Built from the first symmetric tensor seen and cached. See - allocate_symmetric for why a single anchor suffices. - """ - if self._context_bases is None: - self._context_bases = self.symmetric_address_map(anchor).peer_bases - return self._context_bases + self._deltas = deltas + return deltas # ── convenience ────────────────────────────────────────────────────────── diff --git a/tests/unittests/test_rocshmem_provider.py b/tests/unittests/test_rocshmem_provider.py index b0739a8fb..62107cb55 100644 --- a/tests/unittests/test_rocshmem_provider.py +++ b/tests/unittests/test_rocshmem_provider.py @@ -78,12 +78,27 @@ def test_peer_bases_shape_and_invariant(symmetric_pair): assert int(peer_bases[provider.get_rank()].item()) == data.data_ptr() -def test_table_is_context_wide(provider): - """A table built from one allocation is valid for any other.""" +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 torch.equal(bases_a, bases_b) + 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) From e3c2955f266bbb501808762650c904760891c07a Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Fri, 11 Sep 2026 17:27:15 -0500 Subject: [PATCH 03/16] Document collective semantics, table lifetime, and packaging Review follow-ups, all documentation: - allocate_symmetric: say it is collective, describe the returned table's shape, dtype, device and invariant, and point at allocate_symmetric_map for the direct mask. - symmetric_address_map: note it builds a fresh table per call, why that is once per allocation rather than per launch, and why it is deliberately not memoised by data_ptr. - free: explain it cannot be automated. rocshmem_free is documented as collective and must be called by all PEs, so a __del__ or weakref finalizer would let ranks diverge on GC timing and hang. - module: record that rocshmem4py is a standalone package that statically links rocSHMEM rather than linking it at run time, so its USE_IPC setting is fixed at its build time. --- iris/experimental/rocshmem_provider.py | 44 +++++++++++++++++++++++++- 1 file changed, 43 insertions(+), 1 deletion(-) diff --git a/iris/experimental/rocshmem_provider.py b/iris/experimental/rocshmem_provider.py index 870961b07..3fa8e9feb 100644 --- a/iris/experimental/rocshmem_provider.py +++ b/iris/experimental/rocshmem_provider.py @@ -42,6 +42,14 @@ 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: @@ -99,6 +107,22 @@ def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Ten Same signature and return shape as Iris.allocate_symmetric, so the same device kernels drive either 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 @@ -113,7 +137,16 @@ def allocate_symmetric_map(self, *size, dtype=None) -> tuple[torch.Tensor, Symme return tensor, self.symmetric_address_map(tensor) def symmetric_address_map(self, tensor: torch.Tensor) -> SymmetricAddressMap: - """Describe an already-allocated rocSHMEM tensor.""" + """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] @@ -167,6 +200,15 @@ 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: From 2c709f705dadc0116f1c4301d5c5b4e426731288 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Fri, 11 Sep 2026 17:55:27 -0500 Subject: [PATCH 04/16] Install rocSHMEM in the CI images so the provider tests run The provider tests have only ever skipped in CI, because rocshmem4py was not installed. Add it, so they actually execute. Installing it is a source build but a short one. There is no prebuilt rocshmem4py wheel on any index -- not PyPI, not the ROCm nightly indexes -- and the bindings do not build rocSHMEM themselves (find_package(rocshmem 3.5.0 CONFIG REQUIRED), no FetchContent), so rocSHMEM is built first and the bindings are pip-installed against it with CMAKE_PREFIX_PATH. Scope is IPC only, which is what the provider uses and all a single-node runner can exercise. Upstream already defaults USE_IPC=ON and USE_GDA=OFF, so no conduit flags are passed, which keeps MPI and the RDMA provider libraries out of it. Built for gfx942 to match the MI325X runners. One installer shared by both image definitions rather than two copies. That required widening the Docker build context from docker/ to the repo root so the Dockerfile can COPY it, hence the .dockerignore; the Apptainer def pulls the same file in via %files. Also adds iris/experimental/README.md covering install, verification and the collective-call contract, as requested in review. --- .dockerignore | 15 ++++ .github/scripts/container_build.sh | 7 +- .github/scripts/install_rocshmem.sh | 88 +++++++++++++++++++++ apptainer/iris.def | 15 ++++ docker/Dockerfile | 8 ++ iris/experimental/README.md | 115 ++++++++++++++++++++++++++++ 6 files changed, 246 insertions(+), 2 deletions(-) create mode 100644 .dockerignore create mode 100755 .github/scripts/install_rocshmem.sh create mode 100644 iris/experimental/README.md 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..f07598559 --- /dev/null +++ b/.github/scripts/install_rocshmem.sh @@ -0,0 +1,88 @@ +#!/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)" + +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 + +# 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..d8caa71a1 --- /dev/null +++ b/iris/experimental/README.md @@ -0,0 +1,115 @@ +# 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`), so build it with the interpreter that +will run it. + +### 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. From 8daf164f641063cdbac1c9de1251e5d74f9ed37f Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Fri, 11 Sep 2026 18:00:37 -0500 Subject: [PATCH 05/16] Say why the CI install uses the checkout rather than a git+ URL Both are pip installs from source, but the README documented the git+ form while the CI script used a local path, with nothing explaining the difference. One checkout serves the core build and the bindings, so they are guaranteed to be the same revision. A git+ URL would have pip clone the monorepo again at whatever develop is at by then, and find_package would not catch the skew because it only compares versions while the bindings statically link the core. --- .github/scripts/install_rocshmem.sh | 8 ++++++++ iris/experimental/README.md | 12 ++++++++++-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/.github/scripts/install_rocshmem.sh b/.github/scripts/install_rocshmem.sh index f07598559..b1d0c568d 100755 --- a/.github/scripts/install_rocshmem.sh +++ b/.github/scripts/install_rocshmem.sh @@ -69,6 +69,14 @@ if ! grep -qi "define ROCSHMEM_USE_IPC\|USE_IPC" \ 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. diff --git a/iris/experimental/README.md b/iris/experimental/README.md index d8caa71a1..dc0b543cd 100644 --- a/iris/experimental/README.md +++ b/iris/experimental/README.md @@ -48,8 +48,16 @@ Two properties of the result are worth knowing: - Consequently rocSHMEM's **build options are fixed when `rocshmem4py` is built**, not when it is used. -The wheel is CPython-ABI-tagged (`cp312`), so build it with the interpreter that -will run it. +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 From dfe2c8fcdf286048ebf3090f47a86f55c462e209 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Fri, 11 Sep 2026 18:20:58 -0500 Subject: [PATCH 06/16] Stop citing an Iris.allocate_symmetric method that does not exist The docstring claimed the same signature and return shape as Iris.allocate_symmetric. Iris has no such method -- its allocation API is zeros/ones/full/uniform/as_symmetric, with a context-wide table from get_heap_bases(). The (tensor, peer_bases) shape comes from the provider interface being proposed, not from an existing method. --- iris/experimental/rocshmem_provider.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/iris/experimental/rocshmem_provider.py b/iris/experimental/rocshmem_provider.py index 3fa8e9feb..c7a4cd341 100644 --- a/iris/experimental/rocshmem_provider.py +++ b/iris/experimental/rocshmem_provider.py @@ -105,8 +105,8 @@ def __init__(self, device: str | None = None): def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Tensor]: """Allocate a symmetric tensor and return it with its peer-base table. - Same signature and return shape as Iris.allocate_symmetric, so the same - device kernels drive either provider. + 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. From a69502db3f1972eccd31d9ecb94837a98afa95f2 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Fri, 11 Sep 2026 18:36:08 -0500 Subject: [PATCH 07/16] Note that ROCm 7.14+ ships rocSHMEM, so this build is temporary Comment only. The CI bases are ROCm 7.2.1 (apptainer) and 7.1 (docker), both older than the 7.14 artifacts that ship rocSHMEM's static library and headers, so the source build is still required. Records where to cut it when a base image bumps, and that rocshmem4py stays a source build either way until its TheRock packaging lands. --- .github/scripts/install_rocshmem.sh | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/scripts/install_rocshmem.sh b/.github/scripts/install_rocshmem.sh index b1d0c568d..e2f510e56 100755 --- a/.github/scripts/install_rocshmem.sh +++ b/.github/scripts/install_rocshmem.sh @@ -24,6 +24,11 @@ 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 From dd73aa81c678d74a1599871dd4053a15b3c5730a Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Mon, 14 Sep 2026 11:26:13 -0500 Subject: [PATCH 08/16] Make the rocSHMEM init failure in CI diagnosable The provider tests now run in CI rather than skip, and rocSHMEM aborts inside init with no message: a bare SIGABRT through library_init -> IPCBackend -> Backend::init, 40ms in, with nothing to say why. Set ROCSHMEM_DEBUG_LEVEL=info in both images so it prints its config banner, the env vars it actually saw, and the backend it selected. Also write the rank, world size and visible GPU count straight to fd 2 before the call, which survives both pytest's capture and the abort, so the crash says what it was attempting. Diagnostic only; drop both once the tests pass. --- apptainer/iris.def | 5 +++++ docker/Dockerfile | 7 ++++++- tests/unittests/test_rocshmem_provider.py | 9 +++++++++ 3 files changed, 20 insertions(+), 1 deletion(-) diff --git a/apptainer/iris.def b/apptainer/iris.def index 7a9d307ba..e944a57a9 100644 --- a/apptainer/iris.def +++ b/apptainer/iris.def @@ -67,6 +67,11 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0 # Set required RCCL environment variable for ROCm export HSA_NO_SCRATCH_RECLAIM=1 export ROCSHMEM_PREFIX=/opt/rocshmem + # Diagnostic, remove once the provider tests pass in CI: rocSHMEM aborts inside + # init with no message at all, and at info level it prints its config banner, the + # env vars it actually saw, and the backend it selected. Without this the failure + # is a bare SIGABRT. + export ROCSHMEM_DEBUG_LEVEL=info %runscript echo "Welcome to the ROCm-aware Apptainer image!" diff --git a/docker/Dockerfile b/docker/Dockerfile index 944b5c254..10a2ce6f1 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -51,7 +51,12 @@ ENV PYTHONPATH=$TRITON_PATH # 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 +# Diagnostic, remove once the provider tests pass in CI: rocSHMEM aborts inside +# init with no message at all, and at info level it prints its config banner, the +# env vars it actually saw, and the backend it selected. Without this the failure +# is a bare SIGABRT. +ENV ROCSHMEM_PREFIX=/opt/rocshmem \ + ROCSHMEM_DEBUG_LEVEL=info 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 diff --git a/tests/unittests/test_rocshmem_provider.py b/tests/unittests/test_rocshmem_provider.py index 62107cb55..4da15dedc 100644 --- a/tests/unittests/test_rocshmem_provider.py +++ b/tests/unittests/test_rocshmem_provider.py @@ -13,6 +13,8 @@ normal CI run. tests/manual_rocshmem_provider.py covers the multi-node case. """ +import os + import pytest import torch import torch.distributed as dist @@ -50,6 +52,13 @@ def provider(): ) from iris.experimental.rocshmem_provider import RocshmemProvider + # init_rocshmem_by_uniqueid can abort the process rather than raise, and it + # does so before pytest can attribute the failure to anything. Write the + # context straight to fd 2 first, which survives both the capture and the + # abort, so a crash says what it was attempting. + os.write(2, f"[rank {dist.get_rank()}/{dist.get_world_size()}] " + f"rocshmem init, {torch.cuda.device_count()} visible GPUs\n".encode()) + # 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) From c1c820fe9055748c963f1e3275f013d42e1786f1 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Tue, 15 Sep 2026 11:23:11 -0500 Subject: [PATCH 09/16] Add the ROCm 7.14 TODO and stop repeating the same comment three times Review feedback. The rationale for the source build, and for ROCSHMEM_DEBUG_LEVEL, was spelled out in install_rocshmem.sh and then again in both image definitions. Keep it in the script and leave a pointer at each call site. Mark the ROCm 7.14 note as a TODO so it shows up when someone greps for work to drop, rather than reading as background. --- .github/scripts/install_rocshmem.sh | 15 ++++++++++----- apptainer/iris.def | 11 +++-------- docker/Dockerfile | 11 +++-------- 3 files changed, 16 insertions(+), 21 deletions(-) diff --git a/.github/scripts/install_rocshmem.sh b/.github/scripts/install_rocshmem.sh index e2f510e56..9e1fec588 100755 --- a/.github/scripts/install_rocshmem.sh +++ b/.github/scripts/install_rocshmem.sh @@ -24,11 +24,16 @@ 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. +# TODO: drop this source build once the CI base images reach ROCm 7.14+, whose +# artifacts ship rocSHMEM's static library and headers -- only the bindings below +# would still be needed. The bases are ROCm 7.2.1 (apptainer) and 7.1 (docker) +# today. rocshmem4py stays a source build either way until its TheRock packaging +# lands. +# +# The images set ROCSHMEM_DEBUG_LEVEL=info alongside this. That is diagnostic +# only, for a rocSHMEM init that aborts in CI with no message: at info level it +# prints its config banner, the env vars it saw, and the backend it chose. Drop +# it once the provider tests pass. 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 diff --git a/apptainer/iris.def b/apptainer/iris.def index e944a57a9..3d49d0517 100644 --- a/apptainer/iris.def +++ b/apptainer/iris.def @@ -45,10 +45,8 @@ 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 + rocshmem4py, so the provider tests run instead of skipping. + # See install_rocshmem.sh for the rationale. gfx942 = the MI325X runners. ROCSHMEM_GPU_TARGETS=gfx942 ROCSHMEM_PREFIX=/opt/rocshmem \ bash /opt/install_rocshmem.sh " @@ -67,10 +65,7 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0 # Set required RCCL environment variable for ROCm export HSA_NO_SCRATCH_RECLAIM=1 export ROCSHMEM_PREFIX=/opt/rocshmem - # Diagnostic, remove once the provider tests pass in CI: rocSHMEM aborts inside - # init with no message at all, and at info level it prints its config banner, the - # env vars it actually saw, and the backend it selected. Without this the failure - # is a bare SIGABRT. + # Diagnostic; see install_rocshmem.sh. Drop once the provider tests pass. export ROCSHMEM_DEBUG_LEVEL=info %runscript diff --git a/docker/Dockerfile b/docker/Dockerfile index 10a2ce6f1..fba2ce9d2 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -47,14 +47,9 @@ 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. -# Diagnostic, remove once the provider tests pass in CI: rocSHMEM aborts inside -# init with no message at all, and at info level it prints its config banner, the -# env vars it actually saw, and the backend it selected. Without this the failure -# is a bare SIGABRT. +# rocSHMEM + rocshmem4py, so tests/unittests/test_rocshmem_provider.py runs +# instead of skipping. See install_rocshmem.sh for why it is a source build and +# why ROCSHMEM_DEBUG_LEVEL is set. gfx942 matches the MI325X CI runners. ENV ROCSHMEM_PREFIX=/opt/rocshmem \ ROCSHMEM_DEBUG_LEVEL=info COPY .github/scripts/install_rocshmem.sh /tmp/install_rocshmem.sh From 6e8a77f2e27d149176c35e30681660a5c5b3b733 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Tue, 15 Sep 2026 15:04:29 -0500 Subject: [PATCH 10/16] Suspend pytest capture so the rocSHMEM init abort says something The previous diagnostic round produced nothing: neither rocSHMEM's banner nor a write to fd 2 reached the CI log. pytest captures at the fd level and has already redirected fd 2 by the time a fixture runs, so both landed in a buffer that is discarded when the process aborts. faulthandler's output survived only because it dups the original fd 2 at interpreter startup, which is why "Fatal Python error" came through and nothing else did. Suspend capture around init via the capturemanager plugin so rocSHMEM's own logging, and the rank/GPU-count line, reach the log. Also fold two overlapping "Returns (tensor, peer_bases)" paragraphs in allocate_symmetric's docstring into one; they came from two separate edits. Diagnostic; drop with the debug level once the tests pass. --- iris/experimental/rocshmem_provider.py | 9 ++++----- tests/unittests/test_rocshmem_provider.py | 24 +++++++++++++++-------- 2 files changed, 20 insertions(+), 13 deletions(-) diff --git a/iris/experimental/rocshmem_provider.py b/iris/experimental/rocshmem_provider.py index c7a4cd341..80ebef923 100644 --- a/iris/experimental/rocshmem_provider.py +++ b/iris/experimental/rocshmem_provider.py @@ -105,15 +105,14 @@ def __init__(self, device: str | None = None): 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. + This is the provider-facing shape symmetric allocation is converging on, + 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 + ``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`` diff --git a/tests/unittests/test_rocshmem_provider.py b/tests/unittests/test_rocshmem_provider.py index 4da15dedc..70211d411 100644 --- a/tests/unittests/test_rocshmem_provider.py +++ b/tests/unittests/test_rocshmem_provider.py @@ -13,6 +13,7 @@ normal CI run. tests/manual_rocshmem_provider.py covers the multi-node case. """ +import contextlib import os import pytest @@ -38,7 +39,7 @@ def _broadcast_kernel(data, results, peer_bases, n_elements, cur_rank, @pytest.fixture(scope="module") -def provider(): +def provider(request): if not dist.is_initialized(): pytest.skip("needs torch.distributed; run via tests/run_tests_distributed.py") if dist.get_world_size() < 2: @@ -52,16 +53,23 @@ def provider(): ) from iris.experimental.rocshmem_provider import RocshmemProvider - # init_rocshmem_by_uniqueid can abort the process rather than raise, and it - # does so before pytest can attribute the failure to anything. Write the - # context straight to fd 2 first, which survives both the capture and the - # abort, so a crash says what it was attempting. - os.write(2, f"[rank {dist.get_rank()}/{dist.get_world_size()}] " - f"rocshmem init, {torch.cuda.device_count()} visible GPUs\n".encode()) + # init_rocshmem_by_uniqueid can abort the process rather than raise, taking + # pytest with it before anything is attributed. Getting a diagnosis out of + # that needs capture suspended: pytest captures at the fd level and has + # already redirected fd 2 by the time a fixture runs, so rocSHMEM's own + # logging -- and a plain write to fd 2 -- land in a buffer that is discarded + # when the process aborts. faulthandler's output survives only because it + # dups the original fd 2 at interpreter startup. + capman = request.config.pluginmanager.getplugin("capturemanager") + suspended = (capman.global_and_fixture_disabled() + if capman is not None else contextlib.nullcontext()) # 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) + with suspended: + os.write(2, f"[rank {dist.get_rank()}/{dist.get_world_size()}] " + f"rocshmem init, {torch.cuda.device_count()} visible GPUs\n".encode()) + rshmem.init_rocshmem_by_uniqueid(dist.group.WORLD) return RocshmemProvider() From f0ced7e7df418ed4d77cc67a139744715419f402 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Tue, 15 Sep 2026 15:55:30 -0500 Subject: [PATCH 11/16] Build rocSHMEM for gfx950 too; the runner label understates the hardware With capture suspended, rocSHMEM's banner finally reached the log and named the problem outright: # Compiled Arch(s) : gfx942 # System Arch : gfx950:sramecc+:xnack- # System Arch is supported : No The runner label is linux-mi325-8gpu-ossci-rad, and MI325X is gfx942, but the hardware reports gfx950. rocSHMEM's device code has to match the GPU it runs on; built for the wrong arch, the loaded code object has none of its device globals and init aborts in HIP: hip_global.cpp:70 : Cannot create GlobalVar Obj for symbol: _ZN8rocshmem14logd_constantsE Build both arches rather than swapping one guess for another, since the label cannot be trusted to track the pool. --- .github/scripts/install_rocshmem.sh | 8 ++++++-- apptainer/iris.def | 4 ++-- docker/Dockerfile | 4 ++-- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/.github/scripts/install_rocshmem.sh b/.github/scripts/install_rocshmem.sh index 9e1fec588..1db5c69a4 100755 --- a/.github/scripts/install_rocshmem.sh +++ b/.github/scripts/install_rocshmem.sh @@ -17,8 +17,12 @@ 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}" +# Both arches on purpose. The runner label says mi325 (gfx942) but the hardware +# reports gfx950, and rocSHMEM's device code must match the GPU it runs on: build +# for the wrong one and hipModuleGetGlobal fails on rocSHMEM's device globals +# ("Cannot create GlobalVar Obj for symbol: _ZN8rocshmem14logd_constantsE") and +# init aborts. Semicolon-separated cmake list. +ROCSHMEM_GPU_TARGETS="${ROCSHMEM_GPU_TARGETS:-gfx942;gfx950}" ROCSHMEM_REPO="${ROCSHMEM_REPO:-https://github.com/ROCm/rocm-systems.git}" ROCSHMEM_REF="${ROCSHMEM_REF:-develop}" ROCM_PATH="${ROCM_PATH:-/opt/rocm}" diff --git a/apptainer/iris.def b/apptainer/iris.def index 3d49d0517..34d8c1dfa 100644 --- a/apptainer/iris.def +++ b/apptainer/iris.def @@ -46,8 +46,8 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0 chmod -R 777 /opt/venv # rocSHMEM + rocshmem4py, so the provider tests run instead of skipping. - # See install_rocshmem.sh for the rationale. gfx942 = the MI325X runners. - ROCSHMEM_GPU_TARGETS=gfx942 ROCSHMEM_PREFIX=/opt/rocshmem \ + # See install_rocshmem.sh for the rationale, including the two arches. + ROCSHMEM_GPU_TARGETS='gfx942;gfx950' ROCSHMEM_PREFIX=/opt/rocshmem \ bash /opt/install_rocshmem.sh " diff --git a/docker/Dockerfile b/docker/Dockerfile index fba2ce9d2..92b9cd5cc 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -49,11 +49,11 @@ ENV PYTHONPATH=$TRITON_PATH # rocSHMEM + rocshmem4py, so tests/unittests/test_rocshmem_provider.py runs # instead of skipping. See install_rocshmem.sh for why it is a source build and -# why ROCSHMEM_DEBUG_LEVEL is set. gfx942 matches the MI325X CI runners. +# why ROCSHMEM_DEBUG_LEVEL is set, and why it builds for two arches. ENV ROCSHMEM_PREFIX=/opt/rocshmem \ ROCSHMEM_DEBUG_LEVEL=info 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 +RUN ROCSHMEM_GPU_TARGETS='gfx942;gfx950' bash /tmp/install_rocshmem.sh && rm /tmp/install_rocshmem.sh # Set up workspace WORKDIR /workspace From 3851e4ae4c8817a99095ddae60f7d982fa4bfb8b Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Tue, 15 Sep 2026 16:00:51 -0500 Subject: [PATCH 12/16] Remove the rocSHMEM CI diagnostics They did their job: the banner named the arch mismatch that 578394b9 fixes. Removes ROCSHMEM_DEBUG_LEVEL=info from both images, and the fd-2 write plus the pytest capture suspension from the provider fixture. The debug level printed a config banner on every run, and the capture suspension is scaffolding that should not outlive the bug it was added for. Kept as its own commit so `git revert` restores the whole apparatus in one step if a rocSHMEM init failure ever needs diagnosing again. Note that without it such a failure is mute: pytest captures at the fd level, so rocSHMEM's output is buffered and discarded when the process aborts. --- .github/scripts/install_rocshmem.sh | 5 ----- apptainer/iris.def | 2 -- docker/Dockerfile | 5 ++--- tests/unittests/test_rocshmem_provider.py | 21 ++------------------- 4 files changed, 4 insertions(+), 29 deletions(-) diff --git a/.github/scripts/install_rocshmem.sh b/.github/scripts/install_rocshmem.sh index 1db5c69a4..04bace8a0 100755 --- a/.github/scripts/install_rocshmem.sh +++ b/.github/scripts/install_rocshmem.sh @@ -33,11 +33,6 @@ SRC="$(mktemp -d)" # would still be needed. The bases are ROCm 7.2.1 (apptainer) and 7.1 (docker) # today. rocshmem4py stays a source build either way until its TheRock packaging # lands. -# -# The images set ROCSHMEM_DEBUG_LEVEL=info alongside this. That is diagnostic -# only, for a rocSHMEM init that aborts in CI with no message: at info level it -# prints its config banner, the env vars it saw, and the backend it chose. Drop -# it once the provider tests pass. 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 diff --git a/apptainer/iris.def b/apptainer/iris.def index 34d8c1dfa..2fd229477 100644 --- a/apptainer/iris.def +++ b/apptainer/iris.def @@ -65,8 +65,6 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0 # Set required RCCL environment variable for ROCm export HSA_NO_SCRATCH_RECLAIM=1 export ROCSHMEM_PREFIX=/opt/rocshmem - # Diagnostic; see install_rocshmem.sh. Drop once the provider tests pass. - export ROCSHMEM_DEBUG_LEVEL=info %runscript echo "Welcome to the ROCm-aware Apptainer image!" diff --git a/docker/Dockerfile b/docker/Dockerfile index 92b9cd5cc..013a2b643 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -49,9 +49,8 @@ ENV PYTHONPATH=$TRITON_PATH # rocSHMEM + rocshmem4py, so tests/unittests/test_rocshmem_provider.py runs # instead of skipping. See install_rocshmem.sh for why it is a source build and -# why ROCSHMEM_DEBUG_LEVEL is set, and why it builds for two arches. -ENV ROCSHMEM_PREFIX=/opt/rocshmem \ - ROCSHMEM_DEBUG_LEVEL=info +# why it builds for two arches. +ENV ROCSHMEM_PREFIX=/opt/rocshmem COPY .github/scripts/install_rocshmem.sh /tmp/install_rocshmem.sh RUN ROCSHMEM_GPU_TARGETS='gfx942;gfx950' bash /tmp/install_rocshmem.sh && rm /tmp/install_rocshmem.sh diff --git a/tests/unittests/test_rocshmem_provider.py b/tests/unittests/test_rocshmem_provider.py index 70211d411..62107cb55 100644 --- a/tests/unittests/test_rocshmem_provider.py +++ b/tests/unittests/test_rocshmem_provider.py @@ -13,9 +13,6 @@ normal CI run. tests/manual_rocshmem_provider.py covers the multi-node case. """ -import contextlib -import os - import pytest import torch import torch.distributed as dist @@ -39,7 +36,7 @@ def _broadcast_kernel(data, results, peer_bases, n_elements, cur_rank, @pytest.fixture(scope="module") -def provider(request): +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: @@ -53,23 +50,9 @@ def provider(request): ) from iris.experimental.rocshmem_provider import RocshmemProvider - # init_rocshmem_by_uniqueid can abort the process rather than raise, taking - # pytest with it before anything is attributed. Getting a diagnosis out of - # that needs capture suspended: pytest captures at the fd level and has - # already redirected fd 2 by the time a fixture runs, so rocSHMEM's own - # logging -- and a plain write to fd 2 -- land in a buffer that is discarded - # when the process aborts. faulthandler's output survives only because it - # dups the original fd 2 at interpreter startup. - capman = request.config.pluginmanager.getplugin("capturemanager") - suspended = (capman.global_and_fixture_disabled() - if capman is not None else contextlib.nullcontext()) - # rocSHMEM initialises once per process, hence module scope. No finalize in # teardown: it would pull the runtime out from under anything else running. - with suspended: - os.write(2, f"[rank {dist.get_rank()}/{dist.get_world_size()}] " - f"rocshmem init, {torch.cuda.device_count()} visible GPUs\n".encode()) - rshmem.init_rocshmem_by_uniqueid(dist.group.WORLD) + rshmem.init_rocshmem_by_uniqueid(dist.group.WORLD) return RocshmemProvider() From 233c840811702f45971b66dee8d744f58cc63785 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Tue, 15 Sep 2026 16:02:08 -0500 Subject: [PATCH 13/16] Point back at Iris.allocate_symmetric now that it exists #549 landed it on main, returning the same (tensor, peer_bases) pair with peer_bases[cur_rank] == data_ptr(). An earlier commit dropped this reference because the method did not exist yet; matching it is the whole point of the provider, so say so. --- iris/experimental/rocshmem_provider.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/iris/experimental/rocshmem_provider.py b/iris/experimental/rocshmem_provider.py index 80ebef923..52cddb00c 100644 --- a/iris/experimental/rocshmem_provider.py +++ b/iris/experimental/rocshmem_provider.py @@ -105,16 +105,16 @@ def __init__(self, device: str | None = None): def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Tensor]: """Allocate a symmetric tensor and return it with its peer-base table. - This is the provider-facing shape symmetric allocation is converging on, - so the same device kernels drive any provider. + Same signature and return shape as ``Iris.allocate_symmetric``, so the + same device kernels drive either. Collective: rocSHMEM allocation is, so every PE must call this the same number of times and in the same order. ``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 + ``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 From ed0c69ee8570410e81afeb6955de15d1d909547f Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Wed, 16 Sep 2026 10:19:48 -0500 Subject: [PATCH 14/16] Fail the Apptainer build on error instead of caching a broken image A transient network failure during the Triton clone produced an image that Apptainer reported as built, container_build.sh cached by def-file checksum, and every later job reused: Cloning into '/opt/triton'... fetch-pack: unexpected disconnect while reading sideband packet fatal: early EOF /bin/bash: line 22: cd: /opt/triton: No such file or directory ERROR: file:///opt does not appear to be a Python project The %post block had no `set -e`, so the failed clone, the failed cd, and the failed editable install all continued. With /opt/triton absent and the pinned 3.7.0 checkout never installed, `import triton` fell back to the older pytorch-triton-rocm in site-packages, which does not understand Python 3.14's __annotate__. Every Triton and Gluon test then failed far downstream with "Unsupported function referenced: ", pointing nowhere near the cause. Add `set -e`, and retry the clone three times since it is large and this is a transient failure that will recur. The Dockerfile is unaffected: each RUN already fails on error. --- apptainer/iris.def | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/apptainer/iris.def b/apptainer/iris.def index 2fd229477..b2b2d10df 100644 --- a/apptainer/iris.def +++ b/apptainer/iris.def @@ -13,6 +13,16 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0 %post /bin/bash -c " + # Fail the build on the first error. Without this a step can fail silently + # and still produce an image that Apptainer reports as built -- which + # container_build.sh then caches by def-file checksum and every later job + # reuses. That happened: a transient 'fetch-pack: unexpected disconnect' + # during the Triton clone left /opt/triton absent, so the pinned checkout and + # its editable install never ran, 'import triton' silently fell back to the + # older pytorch-triton-rocm in site-packages, and every Triton test failed + # far downstream with an unrelated-looking error. + set -e + # Set environment variables export TRITON_PATH=/opt/triton export ROCM_PATH=/opt/rocm @@ -35,9 +45,16 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0 pip3 install --upgrade pip && \ pip3 install wheel jupyter - # Clone and install Triton + # Clone and install Triton. Retried: this is a large clone and a transient + # disconnect here used to poison the cached image rather than fail the build. cd /opt - git clone https://github.com/triton-lang/triton.git \$TRITON_PATH + for attempt in 1 2 3; do + rm -rf \$TRITON_PATH + git clone https://github.com/triton-lang/triton.git \$TRITON_PATH && break + echo \"triton clone failed (attempt \$attempt)\" >&2 + [ \$attempt -lt 3 ] || exit 1 + sleep 10 + done cd \$TRITON_PATH git checkout f7c1d69401e9f09050451f30776562954b05e850 pip3 install -e . From 51d0b33f7bb3b1ae513e8f5d93e3180d76c23e33 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Wed, 16 Sep 2026 10:42:26 -0500 Subject: [PATCH 15/16] Serialize Apptainer image builds and publish atomically $HOME is shared across the runners, so every job in a run uses one cache directory. Jobs that start together all find no image and all build concurrently into the same path -- observed in run 35023261876, where jobs on iris-mi350x-0 and iris-mi350x-2 both reported "Image or checksum not found" and built, while jobs on -1, -2 and -3 read a checksum one of them had written. Two problems with that. `apptainer build --force` writes IMAGE_PATH in place, so a builder truncates the image other jobs are currently executing. And every job that loses the race still pays for a full redundant build. Take an flock around the check-and-build, and re-check inside it so whoever waits usually finds the image already there. Build to a private temp path and rename into place: rename is atomic and leaves the inode alone, so a job already running the old image is unaffected. Verified with two concurrent builders against a stub: one build, the other waited and reused it, no temp files left. On build failure: exit 1, no checksum written, no image published, temp removed. --- .github/scripts/container_build.sh | 73 +++++++++++++++++++++--------- 1 file changed, 51 insertions(+), 22 deletions(-) diff --git a/.github/scripts/container_build.sh b/.github/scripts/container_build.sh index 83583e57d..10eacf932 100755 --- a/.github/scripts/container_build.sh +++ b/.github/scripts/container_build.sh @@ -43,40 +43,69 @@ if [ "$CONTAINER_RUNTIME" = "apptainer" ]; then DEF_CHECKSUM=$(sha256sum "$DEF_FILE" | awk '{print $1}') # Create persistent Apptainer directory with checksum subdirectory - mkdir -p "${HOME}/iris-apptainer-images/${DEF_CHECKSUM}" + CACHE_DIR="${HOME}/iris-apptainer-images/${DEF_CHECKSUM}" + mkdir -p "$CACHE_DIR" - # Define paths - IMAGE_PATH="${HOME}/iris-apptainer-images/${DEF_CHECKSUM}/iris-dev.sif" - CHECKSUM_FILE="${HOME}/iris-apptainer-images/${DEF_CHECKSUM}/iris-dev.sif.checksum" + # Define paths. $HOME is shared across the runners, so every job in a run + # reads and writes this one directory. + IMAGE_PATH="$CACHE_DIR/iris-dev.sif" + CHECKSUM_FILE="$CACHE_DIR/iris-dev.sif.checksum" + LOCK_FILE="$CACHE_DIR/.build.lock" - # Check if image exists and has a valid checksum - REBUILD_NEEDED=true - if [ -f "$IMAGE_PATH" ] && [ -f "$CHECKSUM_FILE" ]; then - OLD_CHECKSUM=$(head -n1 "$CHECKSUM_FILE" 2>/dev/null) + image_is_current() { + [ -f "$IMAGE_PATH" ] && [ -f "$CHECKSUM_FILE" ] || return 1 + local old + old=$(head -n1 "$CHECKSUM_FILE" 2>/dev/null) # Validate checksum format (64 hex characters for SHA256) - if [[ "$OLD_CHECKSUM" =~ ^[a-f0-9]{64}$ ]] && [ "$OLD_CHECKSUM" = "$DEF_CHECKSUM" ]; then - echo "[INFO] Def file unchanged (checksum: $DEF_CHECKSUM)" - echo "[INFO] Skipping rebuild, using existing image at $IMAGE_PATH" - REBUILD_NEEDED=false - else - echo "[INFO] Def file changed (old: ${OLD_CHECKSUM:-}, new: $DEF_CHECKSUM)" - echo "[INFO] Rebuilding Apptainer image..." - fi - else - echo "[INFO] Image or checksum not found, building new Apptainer image..." - fi + [[ "$old" =~ ^[a-f0-9]{64}$ ]] && [ "$old" = "$DEF_CHECKSUM" ] + } - # Build the image if needed - if [ "$REBUILD_NEEDED" = true ]; then - if apptainer build --force "$IMAGE_PATH" "$DEF_FILE"; then + # Build to a private path and rename into place. Renaming is atomic and + # leaves the inode alone, so a job already executing the old image keeps + # running against it; `apptainer build --force` straight to IMAGE_PATH would + # truncate the file out from under it. + build_image() { + local tmp + tmp=$(mktemp -u "$CACHE_DIR/.iris-dev.XXXXXXXX.sif") + if apptainer build --force "$tmp" "$DEF_FILE"; then + mv -f "$tmp" "$IMAGE_PATH" # Store the checksum only if build succeeded echo "$DEF_CHECKSUM" > "$CHECKSUM_FILE" echo "[INFO] Built image: $IMAGE_PATH" echo "[INFO] Checksum saved: $DEF_CHECKSUM" else + rm -f "$tmp" echo "[ERROR] Apptainer build failed" exit 1 fi + } + + if image_is_current; then + echo "[INFO] Def file unchanged (checksum: $DEF_CHECKSUM)" + echo "[INFO] Skipping rebuild, using existing image at $IMAGE_PATH" + else + echo "[INFO] Image or checksum not found, building new Apptainer image..." + # Serialize builders. Without this, jobs that start together all see no + # image and all build concurrently into the same path -- observed, with + # two runners building at once. The re-check inside the lock is the + # point: whoever waits usually finds the image already built and skips a + # redundant half-hour build. + if command -v flock > /dev/null 2>&1; then + exec 9> "$LOCK_FILE" + if ! flock -w 5400 9; then + echo "[ERROR] Timed out waiting for the image build lock" + exit 1 + fi + if image_is_current; then + echo "[INFO] Another job built it while we waited; using $IMAGE_PATH" + else + build_image + fi + exec 9>&- + else + echo "[WARN] flock not available; building without a lock" + build_image + fi fi elif [ "$CONTAINER_RUNTIME" = "docker" ]; then From 6d6da4dfb7c9a6ede6057376f6646c370d8a7c27 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Wed, 16 Sep 2026 14:03:32 -0500 Subject: [PATCH 16/16] Remove the manual multi-node script Review feedback: it duplicated the unit tests closely enough not to earn its keep, and nothing ran it automatically -- it lives outside the directories CI collects, its name does not match the test_*.py glob, and the launcher is --nnodes=1 regardless. Note what goes with it: the inter-node case, where rocshmem_ptr returns NULL and SymmetricAddressMap.direct reports the peer as unreachable, now has no coverage at all. The mask is still needed for correctness on more than one node. Said so in the test module docstring and the README rather than leaving it implicit. Also fix the build log while touching the file: every job printed "building new Apptainer image" before taking the lock, including the ones that then found the image already built. Now it says what each job actually did. --- .github/scripts/container_build.sh | 4 +- iris/experimental/README.md | 11 +- tests/manual_rocshmem_provider.py | 150 ---------------------- tests/unittests/test_rocshmem_provider.py | 4 +- 4 files changed, 11 insertions(+), 158 deletions(-) delete mode 100644 tests/manual_rocshmem_provider.py diff --git a/.github/scripts/container_build.sh b/.github/scripts/container_build.sh index 10eacf932..47f5c2d04 100755 --- a/.github/scripts/container_build.sh +++ b/.github/scripts/container_build.sh @@ -84,7 +84,7 @@ if [ "$CONTAINER_RUNTIME" = "apptainer" ]; then echo "[INFO] Def file unchanged (checksum: $DEF_CHECKSUM)" echo "[INFO] Skipping rebuild, using existing image at $IMAGE_PATH" else - echo "[INFO] Image or checksum not found, building new Apptainer image..." + echo "[INFO] No current image for this def file" # Serialize builders. Without this, jobs that start together all see no # image and all build concurrently into the same path -- observed, with # two runners building at once. The re-check inside the lock is the @@ -99,11 +99,13 @@ if [ "$CONTAINER_RUNTIME" = "apptainer" ]; then if image_is_current; then echo "[INFO] Another job built it while we waited; using $IMAGE_PATH" else + echo "[INFO] Building new Apptainer image..." build_image fi exec 9>&- else echo "[WARN] flock not available; building without a lock" + echo "[INFO] Building new Apptainer image..." build_image fi fi diff --git a/iris/experimental/README.md b/iris/experimental/README.md index dc0b543cd..3fe75e0b8 100644 --- a/iris/experimental/README.md +++ b/iris/experimental/README.md @@ -98,12 +98,11 @@ 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`). +`run_tests_distributed.py` launches `torchrun` with `--nnodes=1`, so these tests +only ever see intra-node peers, where `rocshmem_ptr` resolves every one. The +inter-node case -- where `rocshmem_ptr` returns NULL and +`SymmetricAddressMap.direct` reports the peer as unreachable -- has no automated +coverage, since it needs two nodes. ### Running diff --git a/tests/manual_rocshmem_provider.py b/tests/manual_rocshmem_provider.py deleted file mode 100644 index 6addfb8b7..000000000 --- a/tests/manual_rocshmem_provider.py +++ /dev/null @@ -1,150 +0,0 @@ -# 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 index 62107cb55..d68665dbc 100644 --- a/tests/unittests/test_rocshmem_provider.py +++ b/tests/unittests/test_rocshmem_provider.py @@ -10,7 +10,9 @@ 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. +normal CI run. The launcher runs torchrun with --nnodes=1, so these only ever +see intra-node peers; the inter-node case, where rocshmem_ptr returns NULL and +SymmetricAddressMap.direct reports it, has no automated coverage. """ import pytest