From 700335484e59809202969fe4e86eeda0ab751f2c Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Thu, 17 Sep 2026 09:35:19 -0500 Subject: [PATCH 1/6] Move SymmetricAddressMap into a module of its own A second provider needs the same descriptor, and importing it from rocshmem_provider would make that provider depend on rocshmem4py. The new module imports nothing but torch, so no provider drags in another's runtime. Re-exported from rocshmem_provider so existing imports keep working. --- iris/experimental/rocshmem_provider.py | 23 ++------------ iris/experimental/symmetric_memory.py | 44 ++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 21 deletions(-) create mode 100644 iris/experimental/symmetric_memory.py diff --git a/iris/experimental/rocshmem_provider.py b/iris/experimental/rocshmem_provider.py index def3901a5..a3e9d2c88 100644 --- a/iris/experimental/rocshmem_provider.py +++ b/iris/experimental/rocshmem_provider.py @@ -60,33 +60,14 @@ from __future__ import annotations -from dataclasses import dataclass - import torch import rocshmem4py as rshmem from rocshmem4py.interop import torch as rshmem_torch +from iris.experimental.symmetric_memory import SymmetricAddressMap -@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] +__all__ = ["RocshmemProvider", "SymmetricAddressMap"] class RocshmemProvider: diff --git a/iris/experimental/symmetric_memory.py b/iris/experimental/symmetric_memory.py new file mode 100644 index 000000000..c8d1a79ab --- /dev/null +++ b/iris/experimental/symmetric_memory.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""Address metadata shared by allocation providers. + +Iris device code translates a pointer with + + remote = peer_bases[to] + (ptr - peer_bases[local_rank]) + +so any provider that can produce a table satisfying +``peer_bases[local_rank] == local allocation base`` drives iris.store, load and +copy unchanged. This module holds the descriptor that carries such a table plus +what the bare table cannot express, notably per-peer reachability. + +It deliberately imports nothing but torch, so a provider for one runtime never +drags in another's dependency. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import 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] From 22259124bf0c83265606d07ddee9ad6b0a0ff18b Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Thu, 17 Sep 2026 09:35:19 -0500 Subject: [PATCH 2/6] Add PyTorch symmetric memory as an allocation provider Lets Iris device kernels operate on tensors from torch.distributed._symmetric_memory. No Iris device code changes: iris.store, load and copy translate against a peer-base table, and torch's rendezvous handle exposes buffer_ptrs, which already is that table with the local rank's entry equal to the tensor's own data_ptr. The provider reads it off the handle and checks the invariant rather than computing anything. Unlike rocSHMEM, each allocation is a separate IPC mapping rather than a window onto one linear heap, so peer offsets are not shared between allocations and a table must only translate pointers into the allocation it describes. symmetric_address_map refuses tensors this provider did not allocate: rendezvous is collective, so re-deriving a handle on demand would hang the ranks that did not ask. Validated on gfx950 at 2, 4 and 8 ranks: 5 passed. --- iris/experimental/torch_symmmem_provider.py | 171 ++++++++++++++++++ .../unittests/test_torch_symmmem_provider.py | 164 +++++++++++++++++ 2 files changed, 335 insertions(+) create mode 100644 iris/experimental/torch_symmmem_provider.py create mode 100644 tests/unittests/test_torch_symmmem_provider.py diff --git a/iris/experimental/torch_symmmem_provider.py b/iris/experimental/torch_symmmem_provider.py new file mode 100644 index 000000000..a5c22cc09 --- /dev/null +++ b/iris/experimental/torch_symmmem_provider.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""PyTorch symmetric memory as an allocation provider for Iris device kernels. + +Lets Iris device code operate on tensors from ``torch.distributed._symmetric_memory`` +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. + +Torch hands that table over directly. Rendezvous returns a handle whose +``buffer_ptrs`` is a per-peer list of pointers into this process's address space, +with the local rank's entry equal to the tensor's own ``data_ptr()`` -- exactly +the invariant Iris needs. So this provider does no pointer arithmetic at all; it +reads the table off the handle per allocation. + +That is worth preferring over the alternative even though the alternative also +works. Peer offsets happen to be constant across allocations, so one table could +be reused for all of them, but that only holds while every rank allocates in +lockstep: ``symm_mem.empty`` is a local call and only ``rendezvous`` is +collective, so ranks can diverge. Reading ``buffer_ptrs`` per allocation is +correct either way. + +Scope is intra-node, set by the backend torch selects. The default backend +reports as ``'CUDA'`` and is the HIP IPC path on ROCm; it reaches peers sharing a +node. Do not call ``set_backend`` to try to change this -- the name +``get_backend`` returns is not one ``set_backend`` accepts, and forcing a +different backend selects one that fails at allocation. + +``SymmetricAddressMap.direct`` records per-peer reachability for symmetry with +the other providers; a peer that is not directly addressable gets a base of 0, +which would translate to a wild pointer rather than an error, so callers are +expected to check it before launching. + +This module is deliberately NOT imported by ``iris/experimental/__init__.py``. +Torch symmetric memory is a private torch API and its availability varies by +build, so importing it eagerly would make ``import iris`` fail on builds that +lack it. + +The caller owns bootstrap and tensor lifetime; the process group must already be +initialised on a device backend: + + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + provider = TorchSymmMemProvider() +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem + +from iris.experimental.symmetric_memory import SymmetricAddressMap + +__all__ = ["TorchSymmMemProvider", "SymmetricAddressMap"] + + +class TorchSymmMemProvider: + """Allocates torch symmetric tensors and describes them for Iris kernels.""" + + def __init__(self, group: dist.ProcessGroup | None = None, device: torch.device | str | None = None): + if not dist.is_initialized(): + raise RuntimeError( + "torch.distributed must be initialised before constructing " + "TorchSymmMemProvider; symmetric memory rendezvous is a collective." + ) + self._group = group if group is not None else dist.group.WORLD + self.cur_rank = dist.get_rank(self._group) + self.num_ranks = dist.get_world_size(self._group) + self._device = ( + torch.device(device) if device is not None else torch.device(f"cuda:{torch.cuda.current_device()}") + ) + # data_ptr -> handle, so an allocation can be described again later + # without a second rendezvous. Rendezvous is collective: calling it from + # one rank alone would hang the others. + self._handles: dict[int, object] = {} + + # ── 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. + """ + tensor, amap = self.allocate_symmetric_map(*size, dtype=dtype) + return tensor, amap.peer_bases + + # ── descriptor form ────────────────────────────────────────────────────── + + def allocate_symmetric_map(self, *size, dtype=None) -> tuple[torch.Tensor, SymmetricAddressMap]: + """As allocate_symmetric, but returning the full address descriptor. + + Every rank must call this the same number of times and in the same + order: the rendezvous inside is a collective. + """ + shape = tuple(size[0]) if len(size) == 1 and hasattr(size[0], "__iter__") else tuple(size) + dtype = dtype or torch.get_default_dtype() + + # State may be built from a forward running under inference_mode, which + # would otherwise mark the allocation inference-only. + with torch.inference_mode(False), torch.no_grad(): + tensor = symm_mem.empty(*shape, dtype=dtype, device=self._device) + handle = symm_mem.rendezvous(tensor, group=self._group) + self._handles[tensor.data_ptr()] = handle + return tensor, self._map_from_handle(tensor, handle) + + def symmetric_address_map(self, tensor: torch.Tensor) -> SymmetricAddressMap: + """Describe an already-allocated symmetric tensor. + + The tensor must have been allocated through this provider, whose handle + is reused. Rendezvousing again here would be a collective call made from + whichever rank happened to ask. + """ + handle = self._handles.get(tensor.data_ptr()) + if handle is None: + raise KeyError( + "tensor was not allocated by this provider, so its rendezvous " + "handle is unknown. Allocate through allocate_symmetric to have " + "the handle recorded; rendezvous cannot be repeated here because " + "it is collective." + ) + return self._map_from_handle(tensor, handle) + + def _map_from_handle(self, tensor: torch.Tensor, handle) -> SymmetricAddressMap: + """Build the descriptor from the handle's own peer-pointer table.""" + bases = [int(p) for p in handle.buffer_ptrs] + + # The invariant every Iris translation depends on. Checked rather than + # assumed: a silent mismatch here turns every remote address in a kernel + # into a wild pointer. + if bases[self.cur_rank] != tensor.data_ptr(): + raise RuntimeError( + f"buffer_ptrs[{self.cur_rank}]={bases[self.cur_rank]:#x} does not " + f"match the tensor base {tensor.data_ptr():#x}; Iris address " + "translation would produce wild pointers." + ) + + return SymmetricAddressMap( + peer_bases=torch.tensor(bases, dtype=torch.int64, device=tensor.device), + local_rank=self.cur_rank, + allocation_base=tensor.data_ptr(), + allocation_bytes=tensor.numel() * tensor.element_size(), + direct=tuple(b != 0 for b in bases), + ) + + # ── convenience ────────────────────────────────────────────────────────── + + def barrier(self): + dist.barrier(self._group) + + def free(self, tensor: torch.Tensor): + """Drop this provider's reference to the allocation. + + Torch symmetric memory is reference counted like any other tensor, so + the storage goes away when the caller's reference does too. This only + forgets the handle. + """ + self._handles.pop(tensor.data_ptr(), None) + + def get_rank(self) -> int: + return self.cur_rank + + def get_num_ranks(self) -> int: + return self.num_ranks diff --git a/tests/unittests/test_torch_symmmem_provider.py b/tests/unittests/test_torch_symmmem_provider.py new file mode 100644 index 000000000..500fc1b3c --- /dev/null +++ b/tests/unittests/test_torch_symmmem_provider.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""Iris device kernels driving PyTorch symmetric memory. + +Run under the usual launcher, which sets up torch.distributed and the device: + + python tests/run_tests_distributed.py tests/unittests/test_torch_symmmem_provider.py \ + --num_ranks 2 -v + +Skips when torch symmetric memory cannot allocate on this build, 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. +""" + +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 torch.""" + 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)") + if dist.get_backend() == "gloo": + pytest.skip("symmetric memory needs a device process group, not gloo") + + # 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. + from iris.experimental.torch_symmmem_provider import TorchSymmMemProvider + + p = TorchSymmMemProvider() + # Availability is not an import question: the module always imports, and + # whether an allocation backend exists is a property of the torch build. A + # throwaway allocation is the only reliable probe. It is collective, so + # every rank runs it and reaches the same verdict. + try: + probe, _ = p.allocate_symmetric(8, dtype=torch.float32) + except Exception as exc: + pytest.skip(f"torch symmetric memory cannot allocate on this build: {exc}") + p.free(probe) + del probe + return p + + +@pytest.fixture +def symmetric_pair(provider): + """Two allocations, each with its own table. + + Both tables are yielded because, unlike rocSHMEM, a table built for one + torch allocation does not translate another's pointers: each allocation is a + separate IPC mapping rather than a window onto one linear heap. A kernel must + translate against the table of the allocation it is addressing. + """ + data, data_bases = provider.allocate_symmetric(BLOCK_SIZE, dtype=torch.float32) + results, results_bases = provider.allocate_symmetric(BLOCK_SIZE, dtype=torch.float32) + yield provider, data, results, data_bases, results_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, data_bases, _results_bases = symmetric_pair + ws = provider.get_num_ranks() + + assert data_bases.numel() == ws + assert data_bases.dtype == torch.int64 + assert data_bases.is_cuda + # peer_bases[local_rank] is the base translation subtracts. + assert int(data_bases[provider.get_rank()].item()) == data.data_ptr() + + +def test_table_is_per_allocation(provider): + """Each allocation is described by its own table, read off its own handle. + + Anchoring on its own base is the whole guarantee. Peer offsets are NOT + guaranteed to be shared between allocations -- each one is a separate IPC + mapping, not a window onto a linear heap -- so a table built for one + allocation must not be used to translate another's pointers. + """ + a, bases_a = provider.allocate_symmetric(64, dtype=torch.float32) + b, bases_b = provider.allocate_symmetric(64, dtype=torch.float32) + try: + me = provider.get_rank() + assert int(bases_a[me].item()) == a.data_ptr() + assert int(bases_b[me].item()) == b.data_ptr() + assert a.data_ptr() != b.data_ptr(), "distinct allocations must not alias" + 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, _data_bases, _results_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_foreign_tensor_is_rejected(provider): + """A tensor this provider never allocated has no handle to describe it.""" + stray = torch.zeros(8, dtype=torch.float32, device="cuda") + with pytest.raises(KeyError): + provider.symmetric_address_map(stray) + + +def test_iris_store_over_torch_symmmem_memory(symmetric_pair): + """Unmodified iris.store, on memory Iris did not allocate.""" + provider, data, results, _data_bases, results_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: + # results_bases, not data_bases: the store addresses `results`, and + # each allocation has its own mapping. tl.load of `data` is local and + # needs no translation. + _broadcast_kernel[(1,)]( + data, results, results_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 5a62d9ff7bbdf5aa4b37f7e4047a08d11bb5dfbc Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Thu, 17 Sep 2026 09:35:19 -0500 Subject: [PATCH 3/6] Document the torch symmetric memory provider alongside rocSHMEM Covers that it needs no install beyond a torch build whose symmetric memory can allocate, which torch versions those are, why the tests probe by allocating rather than importing, the set_backend trap, and how its per-allocation mapping differs from rocSHMEM's uniform heap offsets. --- iris/experimental/README.md | 62 +++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/iris/experimental/README.md b/iris/experimental/README.md index 3fe75e0b8..f89f035d1 100644 --- a/iris/experimental/README.md +++ b/iris/experimental/README.md @@ -17,6 +17,7 @@ iris` never requires either dependency. | Provider | Dependency | Scope | | --- | --- | --- | | `rocshmem_provider.py` | `rocshmem4py` | intra-node (IPC) | +| `torch_symmmem_provider.py` | torch only | intra-node (HIP IPC) | ## rocSHMEM provider @@ -120,3 +121,64 @@ Allocation and free are both **collective** — `rocshmem_free` is documented as 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. + +## PyTorch symmetric memory provider + +### Installing + +Nothing to install. It uses `torch.distributed._symmetric_memory`, so the only +requirement is a torch build whose symmetric memory can allocate. + +That is not a given on ROCm, and the import is not a useful test of it — the +module imports successfully on builds where allocation then fails. Measured on +gfx950: + +| torch | symmetric memory | +| --- | --- | +| `2.10.0+rocm7.2.1` | works | +| `2.13.0+rocm7.2` | works | +| `2.9.1+rocm7.1` | no — lacks `set_signal_pad_size`, and `rendezvous` fails with `HIP error: invalid argument` | +| `2.13.0+rocm7.15.0a` (nightly) | no — `hipErrorOutOfMemory` on a 4 KB allocation | + +So the tests probe by attempting a throwaway allocation and skip if it fails, +rather than by importing. + +**Do not call `symm_mem.set_backend()`.** The name `get_backend()` returns on +ROCm is `'CUDA'` — the HIP IPC path — but `set_backend('CUDA')` is rejected with +"SymmetricMemory does not find allocation backend CUDA", and forcing any other +name selects a backend that then fails at allocation. The working sequence is +simply to allocate without ever setting one: + +```python +torch.cuda.set_device(device) +dist.init_process_group("nccl") +tensor = symm_mem.empty(shape, dtype=dtype, device=device) +handle = symm_mem.rendezvous(tensor, group=group) +``` + +### Verifying + +```bash +python tests/run_tests_distributed.py \ + tests/unittests/test_torch_symmmem_provider.py --num_ranks 2 -v +``` + +Passes at 2, 4 and 8 ranks on `torch 2.10.0+rocm7.2.1`. Both allocation and +`rendezvous` are collective, so every rank must call `allocate_symmetric` the +same number of times and in the same order. + +### How it differs from the rocSHMEM provider + +The rendezvous handle's `buffer_ptrs` is already the table Iris needs, with the +local rank's entry equal to the tensor's own `data_ptr()`, so this provider does +no pointer arithmetic — it reads the table off the handle. + +Peer offsets are **not** shared between allocations here. Each allocation is a +separate IPC mapping rather than a window onto one linear heap, so a table must +only translate pointers into the allocation it describes. This is the opposite of +rocSHMEM, where offsets are uniform across the symmetric heap; borrowing another +allocation's table there is legitimate and here it produces wild pointers. + +`symmetric_address_map` only accepts tensors this provider allocated. Rendezvous +is collective, so deriving a handle on demand would hang whichever ranks did not +ask; handles are recorded at allocation time instead. From 14ba27c6e0d662c3a898163ca72c9cc894810bd5 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Thu, 17 Sep 2026 10:02:11 -0500 Subject: [PATCH 4/6] Review: rename to match torch's abbreviation, check table length, sync all ranks Rename torch_symmmem_provider to torch_symm_mem_provider. The triple m reads as a typo; torch itself abbreviates torch.distributed._symmetric_memory as symm_mem, so follow that rather than the symmem spelling suggested in review. Validate that handle.buffer_ptrs has one entry per rank. Iris indexes this table by rank inside the kernel, so a short table reads past the end of the allocation instead of raising. Observed builds return exactly world_size ints, so this is insurance rather than a fix. Synchronise on every rank after the barrier in the store test, not only on the writer before it. The barrier orders the ranks; it does not drain the reading rank's stream, so inbound writes were not strictly ordered against work queued next on that rank. --- iris/experimental/README.md | 4 ++-- ..._symmmem_provider.py => torch_symm_mem_provider.py} | 9 +++++++++ ...mem_provider.py => test_torch_symm_mem_provider.py} | 10 +++++++--- 3 files changed, 18 insertions(+), 5 deletions(-) rename iris/experimental/{torch_symmmem_provider.py => torch_symm_mem_provider.py} (94%) rename tests/unittests/{test_torch_symmmem_provider.py => test_torch_symm_mem_provider.py} (93%) diff --git a/iris/experimental/README.md b/iris/experimental/README.md index f89f035d1..638b4aeed 100644 --- a/iris/experimental/README.md +++ b/iris/experimental/README.md @@ -17,7 +17,7 @@ iris` never requires either dependency. | Provider | Dependency | Scope | | --- | --- | --- | | `rocshmem_provider.py` | `rocshmem4py` | intra-node (IPC) | -| `torch_symmmem_provider.py` | torch only | intra-node (HIP IPC) | +| `torch_symm_mem_provider.py` | torch only | intra-node (HIP IPC) | ## rocSHMEM provider @@ -160,7 +160,7 @@ handle = symm_mem.rendezvous(tensor, group=group) ```bash python tests/run_tests_distributed.py \ - tests/unittests/test_torch_symmmem_provider.py --num_ranks 2 -v + tests/unittests/test_torch_symm_mem_provider.py --num_ranks 2 -v ``` Passes at 2, 4 and 8 ranks on `torch 2.10.0+rocm7.2.1`. Both allocation and diff --git a/iris/experimental/torch_symmmem_provider.py b/iris/experimental/torch_symm_mem_provider.py similarity index 94% rename from iris/experimental/torch_symmmem_provider.py rename to iris/experimental/torch_symm_mem_provider.py index a5c22cc09..eab7bb85e 100644 --- a/iris/experimental/torch_symmmem_provider.py +++ b/iris/experimental/torch_symm_mem_provider.py @@ -132,6 +132,15 @@ def _map_from_handle(self, tensor: torch.Tensor, handle) -> SymmetricAddressMap: """Build the descriptor from the handle's own peer-pointer table.""" bases = [int(p) for p in handle.buffer_ptrs] + # Length is checked, not assumed. Iris indexes this table by rank inside + # the kernel, so a short table reads past the end of the allocation + # rather than raising. + if len(bases) != self.num_ranks: + raise RuntimeError( + f"handle.buffer_ptrs has {len(bases)} entries, expected " + f"{self.num_ranks}; the table Iris indexes by rank would be short." + ) + # The invariant every Iris translation depends on. Checked rather than # assumed: a silent mismatch here turns every remote address in a kernel # into a wild pointer. diff --git a/tests/unittests/test_torch_symmmem_provider.py b/tests/unittests/test_torch_symm_mem_provider.py similarity index 93% rename from tests/unittests/test_torch_symmmem_provider.py rename to tests/unittests/test_torch_symm_mem_provider.py index 500fc1b3c..8ad821cd7 100644 --- a/tests/unittests/test_torch_symmmem_provider.py +++ b/tests/unittests/test_torch_symm_mem_provider.py @@ -5,7 +5,7 @@ Run under the usual launcher, which sets up torch.distributed and the device: - python tests/run_tests_distributed.py tests/unittests/test_torch_symmmem_provider.py \ + python tests/run_tests_distributed.py tests/unittests/test_torch_symm_mem_provider.py \ --num_ranks 2 -v Skips when torch symmetric memory cannot allocate on this build, when fewer than @@ -48,7 +48,7 @@ def provider(): # 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. - from iris.experimental.torch_symmmem_provider import TorchSymmMemProvider + from iris.experimental.torch_symm_mem_provider import TorchSymmMemProvider p = TorchSymmMemProvider() # Availability is not an import question: the module always imports, and @@ -136,7 +136,7 @@ def test_foreign_tensor_is_rejected(provider): provider.symmetric_address_map(stray) -def test_iris_store_over_torch_symmmem_memory(symmetric_pair): +def test_iris_store_over_torch_symm_mem_memory(symmetric_pair): """Unmodified iris.store, on memory Iris did not allocate.""" provider, data, results, _data_bases, results_bases = symmetric_pair me, ws = provider.get_rank(), provider.get_num_ranks() @@ -159,6 +159,10 @@ def test_iris_store_over_torch_symmmem_memory(symmetric_pair): ) torch.cuda.synchronize() provider.barrier() + # Every rank, not just the writer: the barrier orders the ranks, but the + # reading rank still has its own stream to drain before the inbound writes + # are guaranteed visible to work it queues next. + torch.cuda.synchronize() # Rank 0 pushed its value to every rank, including this one. assert torch.allclose(results, torch.full_like(results, 1.0)) From b0797f8d51cae35655feec65c0b957f36d188351 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Fri, 9 Oct 2026 12:27:34 -0500 Subject: [PATCH 5/6] Review: join the unified provider test, scope backend claims to what was measured Register the torch provider in test_provider_unified.py (#552). Its probe and skips move into a session fixture in conftest so both modules share them, and the two torch tests the unified test now covers -- the table invariant and per-allocation anchoring -- are dropped. The backend guidance overreached. "Do not call set_backend" and "forcing any other name fails at allocation" generalised from forcing NCCL on 2.13 builds; on the 2.15 rocm10.0 nightly symmetric memory reportedly works only with set_backend("NVSHMEM"), i.e. rocSHMEM. The provider never chose a backend anyway, so the docs now say the choice is the caller's, keep the one real trap (set_backend('CUDA') is rejected even though get_backend reports it), and record the nightly report as reported, not measured. Likewise "peer offsets are not shared between allocations" is a property of the default IPC backend, not of torch symmetric memory; reworded as "need not be". Fixes the README saying allocation is collective while the docstring said it is local: only rendezvous is, on the default backend. The module docstring is cut to the points that matter when reading the code; the rest already lives in the README. Co-Authored-By: Claude Opus 5.5 (1M context) --- iris/experimental/README.md | 59 ++++++++------- iris/experimental/torch_symm_mem_provider.py | 57 +++------------ tests/unittests/conftest.py | 28 +++++++ tests/unittests/test_provider_unified.py | 11 ++- .../unittests/test_torch_symm_mem_provider.py | 73 +++---------------- 5 files changed, 91 insertions(+), 137 deletions(-) diff --git a/iris/experimental/README.md b/iris/experimental/README.md index 638b4aeed..8f4b6ecfb 100644 --- a/iris/experimental/README.md +++ b/iris/experimental/README.md @@ -17,7 +17,7 @@ iris` never requires either dependency. | Provider | Dependency | Scope | | --- | --- | --- | | `rocshmem_provider.py` | `rocshmem4py` | intra-node (IPC) | -| `torch_symm_mem_provider.py` | torch only | intra-node (HIP IPC) | +| `torch_symm_mem_provider.py` | torch only | intra-node (HIP IPC) on the default backend | ## rocSHMEM provider @@ -130,42 +130,46 @@ Nothing to install. It uses `torch.distributed._symmetric_memory`, so the only requirement is a torch build whose symmetric memory can allocate. That is not a given on ROCm, and the import is not a useful test of it — the -module imports successfully on builds where allocation then fails. Measured on -gfx950: +module imports successfully on builds where allocation then fails. On gfx950: -| torch | symmetric memory | -| --- | --- | -| `2.10.0+rocm7.2.1` | works | -| `2.13.0+rocm7.2` | works | -| `2.9.1+rocm7.1` | no — lacks `set_signal_pad_size`, and `rendezvous` fails with `HIP error: invalid argument` | -| `2.13.0+rocm7.15.0a` (nightly) | no — `hipErrorOutOfMemory` on a 4 KB allocation | +| torch | backend | symmetric memory | +| --- | --- | --- | +| `2.10.0+rocm7.2.1` | default | works | +| `2.13.0+rocm7.2` | default | works | +| `2.9.1+rocm7.1` | default | no — lacks `set_signal_pad_size`, and `rendezvous` fails with `HIP error: invalid argument` | +| `2.13.0+rocm7.15.0a` (nightly) | default | no — `hipErrorOutOfMemory` on a 4 KB allocation | +| `2.15.0.dev20260930+rocm10.0` (nightly) | `NVSHMEM` | works, and only with this backend — reported in review of #551, not run here | So the tests probe by attempting a throwaway allocation and skip if it fails, -rather than by importing. +rather than by importing. Only the default backend has been run through them. -**Do not call `symm_mem.set_backend()`.** The name `get_backend()` returns on -ROCm is `'CUDA'` — the HIP IPC path — but `set_backend('CUDA')` is rejected with -"SymmetricMemory does not find allocation backend CUDA", and forcing any other -name selects a backend that then fails at allocation. The working sequence is -simply to allocate without ever setting one: +**The backend is the caller's choice.** The provider never calls +`symm_mem.set_backend()`; it reads the table off whatever handle `rendezvous` +returns. `NVSHMEM` is torch's name for rocSHMEM on ROCm. One trap if you want +the default: on ROCm `get_backend()` reports it as `'CUDA'` (the HIP IPC path), +but `set_backend('CUDA')` is rejected with "SymmetricMemory does not find +allocation backend CUDA", so keep it by not calling `set_backend` at all. ```python torch.cuda.set_device(device) dist.init_process_group("nccl") -tensor = symm_mem.empty(shape, dtype=dtype, device=device) -handle = symm_mem.rendezvous(tensor, group=group) +# symm_mem.set_backend("NVSHMEM") # on builds where the default cannot allocate +provider = TorchSymmMemProvider() +tensor, peer_bases = provider.allocate_symmetric(n, dtype=torch.float32) ``` ### Verifying ```bash python tests/run_tests_distributed.py \ - tests/unittests/test_torch_symm_mem_provider.py --num_ranks 2 -v + tests/unittests/test_torch_symm_mem_provider.py \ + tests/unittests/test_provider_unified.py --num_ranks 2 -v ``` -Passes at 2, 4 and 8 ranks on `torch 2.10.0+rocm7.2.1`. Both allocation and -`rendezvous` are collective, so every rank must call `allocate_symmetric` the -same number of times and in the same order. +Passes at 2, 4 and 8 ranks on `torch 2.10.0+rocm7.2.1` with the default backend. +`rendezvous` is collective (allocation itself is local on the default backend), +so every rank must call `allocate_symmetric` the same number of times and in the +same order. ### How it differs from the rocSHMEM provider @@ -173,12 +177,13 @@ The rendezvous handle's `buffer_ptrs` is already the table Iris needs, with the local rank's entry equal to the tensor's own `data_ptr()`, so this provider does no pointer arithmetic — it reads the table off the handle. -Peer offsets are **not** shared between allocations here. Each allocation is a -separate IPC mapping rather than a window onto one linear heap, so a table must -only translate pointers into the allocation it describes. This is the opposite of -rocSHMEM, where offsets are uniform across the symmetric heap; borrowing another -allocation's table there is legitimate and here it produces wild pointers. +Peer offsets need not be shared between allocations. On the default backend each +allocation is a separate IPC mapping rather than a window onto one linear heap, +so a table must only translate pointers into the allocation it describes. With +rocSHMEM, offsets are uniform across the symmetric heap and borrowing another +allocation's table happens to work; here it can mistranslate. `symmetric_address_map` only accepts tensors this provider allocated. Rendezvous is collective, so deriving a handle on demand would hang whichever ranks did not -ask; handles are recorded at allocation time instead. +ask; handles are recorded at allocation time instead. Describing tensors that +were already rendezvoused outside the provider is tracked in #577. diff --git a/iris/experimental/torch_symm_mem_provider.py b/iris/experimental/torch_symm_mem_provider.py index eab7bb85e..2d4329e97 100644 --- a/iris/experimental/torch_symm_mem_provider.py +++ b/iris/experimental/torch_symm_mem_provider.py @@ -3,51 +3,18 @@ """PyTorch symmetric memory as an allocation provider for Iris device kernels. -Lets Iris device code operate on tensors from ``torch.distributed._symmetric_memory`` -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. - -Torch hands that table over directly. Rendezvous returns a handle whose -``buffer_ptrs`` is a per-peer list of pointers into this process's address space, -with the local rank's entry equal to the tensor's own ``data_ptr()`` -- exactly -the invariant Iris needs. So this provider does no pointer arithmetic at all; it -reads the table off the handle per allocation. - -That is worth preferring over the alternative even though the alternative also -works. Peer offsets happen to be constant across allocations, so one table could -be reused for all of them, but that only holds while every rank allocates in -lockstep: ``symm_mem.empty`` is a local call and only ``rendezvous`` is -collective, so ranks can diverge. Reading ``buffer_ptrs`` per allocation is -correct either way. - -Scope is intra-node, set by the backend torch selects. The default backend -reports as ``'CUDA'`` and is the HIP IPC path on ROCm; it reaches peers sharing a -node. Do not call ``set_backend`` to try to change this -- the name -``get_backend`` returns is not one ``set_backend`` accepts, and forcing a -different backend selects one that fails at allocation. - -``SymmetricAddressMap.direct`` records per-peer reachability for symmetry with -the other providers; a peer that is not directly addressable gets a base of 0, -which would translate to a wild pointer rather than an error, so callers are -expected to check it before launching. - -This module is deliberately NOT imported by ``iris/experimental/__init__.py``. -Torch symmetric memory is a private torch API and its availability varies by -build, so importing it eagerly would make ``import iris`` fail on builds that -lack it. - -The caller owns bootstrap and tensor lifetime; the process group must already be -initialised on a device backend: - - torch.cuda.set_device(local_rank) - dist.init_process_group("nccl") - provider = TorchSymmMemProvider() +Setup, measured torch builds and how this differs from the rocSHMEM provider are +in iris/experimental/README.md. What matters in the code: + +- The rendezvous handle's ``buffer_ptrs`` already is the table Iris translates + against, with the local entry equal to the tensor's ``data_ptr()``. The + provider reads it per allocation and checks that invariant; it computes + nothing. +- Use each table only for its own allocation. Whether peer offsets are shared + between allocations depends on the backend, and on the default one they need + not be. +- The provider never calls ``symm_mem.set_backend``; it uses whatever backend + the caller selected, or torch's default. """ from __future__ import annotations diff --git a/tests/unittests/conftest.py b/tests/unittests/conftest.py index 8a9e7d3e2..683ddfb04 100644 --- a/tests/unittests/conftest.py +++ b/tests/unittests/conftest.py @@ -25,3 +25,31 @@ def rocshmem_runtime(): rshmem = pytest.importorskip("rocshmem4py", reason="rocSHMEM tests need rocshmem4py installed") rshmem.init_rocshmem_by_uniqueid(dist.group.WORLD) return rshmem + + +@pytest.fixture(scope="session") +def torch_symm_mem_provider(): + """A TorchSymmMemProvider, or a skip if this torch build cannot allocate. + + Availability is not an import question: the module always imports, and + whether an allocation backend exists is a property of the torch build. A + throwaway allocation is the only reliable probe. It is collective, so every + rank runs it and reaches the same verdict. + """ + 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)") + if dist.get_backend() == "gloo": + pytest.skip("symmetric memory needs a device process group, not gloo") + + import torch + from iris.experimental.torch_symm_mem_provider import TorchSymmMemProvider + + provider = TorchSymmMemProvider() + try: + probe, _ = provider.allocate_symmetric(8, dtype=torch.float32) + except Exception as exc: + pytest.skip(f"torch symmetric memory cannot allocate on this build: {exc}") + provider.free(probe) + return provider diff --git a/tests/unittests/test_provider_unified.py b/tests/unittests/test_provider_unified.py index 4cfd61c44..388352b5b 100644 --- a/tests/unittests/test_provider_unified.py +++ b/tests/unittests/test_provider_unified.py @@ -6,8 +6,8 @@ Each tensor translates against its own table. Iris and rocSHMEM allocate from one heap, so the peer delta is shared and any table happens to translate any -pointer; Torch Symmetric Memory allocates per tensor, so the deltas differ and -reusing a table mistranslates. +pointer; Torch Symmetric Memory's default backend maps each tensor separately, +so the deltas need not match and reusing a table can mistranslate. Run: python tests/run_tests_distributed.py tests/unittests/test_provider_unified.py --num_ranks 2 @@ -89,11 +89,18 @@ def _make_rocshmem(request): return RocshmemProvider() +def _make_torch_symm_mem(request): + # Session-scoped too: it probes with a collective allocation and skips if + # this torch build cannot allocate. + return request.getfixturevalue("torch_symm_mem_provider") + + # Needs allocate_symmetric / get_rank / get_num_ranks / barrier. An Iris # context already has all four, so it goes in unwrapped. PROVIDERS = { "iris": lambda request: iris.iris(1 << 24), "rocshmem": _make_rocshmem, + "torch_symm_mem": _make_torch_symm_mem, } diff --git a/tests/unittests/test_torch_symm_mem_provider.py b/tests/unittests/test_torch_symm_mem_provider.py index 8ad821cd7..3bec4e228 100644 --- a/tests/unittests/test_torch_symm_mem_provider.py +++ b/tests/unittests/test_torch_symm_mem_provider.py @@ -11,11 +11,13 @@ Skips when torch symmetric memory cannot allocate on this build, 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. + +The table invariant and per-allocation anchoring are covered for every provider +by test_provider_unified.py; this module holds what is specific to this one. """ import pytest import torch -import torch.distributed as dist import triton import triton.language as tl @@ -37,41 +39,19 @@ def _broadcast_kernel( @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)") - if dist.get_backend() == "gloo": - pytest.skip("symmetric memory needs a device process group, not gloo") - - # 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. - from iris.experimental.torch_symm_mem_provider import TorchSymmMemProvider - - p = TorchSymmMemProvider() - # Availability is not an import question: the module always imports, and - # whether an allocation backend exists is a property of the torch build. A - # throwaway allocation is the only reliable probe. It is collective, so - # every rank runs it and reaches the same verdict. - try: - probe, _ = p.allocate_symmetric(8, dtype=torch.float32) - except Exception as exc: - pytest.skip(f"torch symmetric memory cannot allocate on this build: {exc}") - p.free(probe) - del probe - return p +def provider(torch_symm_mem_provider): + # Probe and skips live in conftest, shared with test_provider_unified.py. + return torch_symm_mem_provider @pytest.fixture def symmetric_pair(provider): """Two allocations, each with its own table. - Both tables are yielded because, unlike rocSHMEM, a table built for one - torch allocation does not translate another's pointers: each allocation is a - separate IPC mapping rather than a window onto one linear heap. A kernel must - translate against the table of the allocation it is addressing. + Both tables are yielded because a table built for one torch allocation need + not translate another's pointers: on the default backend each allocation is + a separate IPC mapping rather than a window onto one linear heap. A kernel + must translate against the table of the allocation it is addressing. """ data, data_bases = provider.allocate_symmetric(BLOCK_SIZE, dtype=torch.float32) results, results_bases = provider.allocate_symmetric(BLOCK_SIZE, dtype=torch.float32) @@ -81,39 +61,6 @@ def symmetric_pair(provider): provider.free(results) -def test_peer_bases_shape_and_invariant(symmetric_pair): - """The invariant Iris device code translates against.""" - provider, data, _results, data_bases, _results_bases = symmetric_pair - ws = provider.get_num_ranks() - - assert data_bases.numel() == ws - assert data_bases.dtype == torch.int64 - assert data_bases.is_cuda - # peer_bases[local_rank] is the base translation subtracts. - assert int(data_bases[provider.get_rank()].item()) == data.data_ptr() - - -def test_table_is_per_allocation(provider): - """Each allocation is described by its own table, read off its own handle. - - Anchoring on its own base is the whole guarantee. Peer offsets are NOT - guaranteed to be shared between allocations -- each one is a separate IPC - mapping, not a window onto a linear heap -- so a table built for one - allocation must not be used to translate another's pointers. - """ - a, bases_a = provider.allocate_symmetric(64, dtype=torch.float32) - b, bases_b = provider.allocate_symmetric(64, dtype=torch.float32) - try: - me = provider.get_rank() - assert int(bases_a[me].item()) == a.data_ptr() - assert int(bases_b[me].item()) == b.data_ptr() - assert a.data_ptr() != b.data_ptr(), "distinct allocations must not alias" - 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, _data_bases, _results_bases = symmetric_pair From f0ba95f7b6b8b3b1331f89a9336a4dd113f1e896 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram Date: Tue, 13 Oct 2026 15:35:09 -0500 Subject: [PATCH 6/6] Say that callers should free(), and record the rocm10 results from review free()'s docstring said dropping the tensor releases the storage and free() only forgets the handle. Backwards: the recorded rendezvous handle keeps the allocation and its peer mappings alive, so nothing is released until free() is called or the provider is dropped. The docstring and the README now say so, and the README example calls it. Also adds the 2.13.0+rocm10.1 build from review, where the provider tests pass on both the default backend and NVSHMEM, and drops the line saying only the default backend has been run through them. Co-Authored-By: Claude Opus 5.5 (1M context) --- iris/experimental/README.md | 16 +++++++++++++--- iris/experimental/torch_symm_mem_provider.py | 10 ++++++---- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/iris/experimental/README.md b/iris/experimental/README.md index 8f4b6ecfb..78d305544 100644 --- a/iris/experimental/README.md +++ b/iris/experimental/README.md @@ -130,7 +130,8 @@ Nothing to install. It uses `torch.distributed._symmetric_memory`, so the only requirement is a torch build whose symmetric memory can allocate. That is not a given on ROCm, and the import is not a useful test of it — the -module imports successfully on builds where allocation then fails. On gfx950: +module imports successfully on builds where allocation then fails. Measured on +gfx950, except the rows reported in review: | torch | backend | symmetric memory | | --- | --- | --- | @@ -138,10 +139,11 @@ module imports successfully on builds where allocation then fails. On gfx950: | `2.13.0+rocm7.2` | default | works | | `2.9.1+rocm7.1` | default | no — lacks `set_signal_pad_size`, and `rendezvous` fails with `HIP error: invalid argument` | | `2.13.0+rocm7.15.0a` (nightly) | default | no — `hipErrorOutOfMemory` on a 4 KB allocation | -| `2.15.0.dev20260930+rocm10.0` (nightly) | `NVSHMEM` | works, and only with this backend — reported in review of #551, not run here | +| `2.15.0.dev20260930+rocm10.0` (nightly) | `NVSHMEM` | works, and only with this backend — standalone check, reported in review of #551 | +| `2.13.0+rocm10.1.0a2026082` | default and `NVSHMEM` | provider tests pass on both — reported in review of #551 | So the tests probe by attempting a throwaway allocation and skip if it fails, -rather than by importing. Only the default backend has been run through them. +rather than by importing. **The backend is the caller's choice.** The provider never calls `symm_mem.set_backend()`; it reads the table off whatever handle `rendezvous` @@ -156,6 +158,8 @@ dist.init_process_group("nccl") # symm_mem.set_backend("NVSHMEM") # on builds where the default cannot allocate provider = TorchSymmMemProvider() tensor, peer_bases = provider.allocate_symmetric(n, dtype=torch.float32) +# ... launch Iris kernels with peer_bases ... +provider.free(tensor) ``` ### Verifying @@ -187,3 +191,9 @@ allocation's table happens to work; here it can mistranslate. is collective, so deriving a handle on demand would hang whichever ranks did not ask; handles are recorded at allocation time instead. Describing tensors that were already rendezvoused outside the provider is tracked in #577. + +**Call `free()` when done with an allocation.** The provider records each +allocation's rendezvous handle, and that handle keeps the memory and its peer +mappings alive, so dropping the tensor alone releases nothing. Unlike +rocSHMEM's, this `free()` is local rather than collective, but barrier first if +peers may still be accessing the allocation. diff --git a/iris/experimental/torch_symm_mem_provider.py b/iris/experimental/torch_symm_mem_provider.py index 2d4329e97..ce6cefaa8 100644 --- a/iris/experimental/torch_symm_mem_provider.py +++ b/iris/experimental/torch_symm_mem_provider.py @@ -132,11 +132,13 @@ def barrier(self): dist.barrier(self._group) def free(self, tensor: torch.Tensor): - """Drop this provider's reference to the allocation. + """Release the provider's hold on the allocation. Callers should call it. - Torch symmetric memory is reference counted like any other tensor, so - the storage goes away when the caller's reference does too. This only - forgets the handle. + The rendezvous handle recorded at allocation keeps the memory and its + peer mappings alive, so dropping the tensor alone releases nothing: + memory is freed only once this is called (or the provider is dropped) + and the caller's own references are gone. It is local, not collective, + but barrier first if peers may still be accessing the allocation. """ self._handles.pop(tensor.data_ptr(), None)