diff --git a/iris/experimental/README.md b/iris/experimental/README.md index 3fe75e0b8..78d305544 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_symm_mem_provider.py` | torch only | intra-node (HIP IPC) on the default backend | ## rocSHMEM provider @@ -120,3 +121,79 @@ 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, except the rows reported in review: + +| 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 — 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. + +**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") +# 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 + +```bash +python tests/run_tests_distributed.py \ + 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` 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 + +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 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. 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/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] diff --git a/iris/experimental/torch_symm_mem_provider.py b/iris/experimental/torch_symm_mem_provider.py new file mode 100644 index 000000000..ce6cefaa8 --- /dev/null +++ b/iris/experimental/torch_symm_mem_provider.py @@ -0,0 +1,149 @@ +# 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. + +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 + +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] + + # 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. + 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): + """Release the provider's hold on the allocation. Callers should call it. + + 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) + + def get_rank(self) -> int: + return self.cur_rank + + def get_num_ranks(self) -> int: + return self.num_ranks 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 new file mode 100644 index 000000000..3bec4e228 --- /dev/null +++ b/tests/unittests/test_torch_symm_mem_provider.py @@ -0,0 +1,115 @@ +# 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_symm_mem_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. + +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 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(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 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) + yield provider, data, results, data_bases, results_bases + provider.barrier() + provider.free(data) + provider.free(results) + + +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_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() + + 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() + # 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))