[experimental] PyTorch symmetric memory as an allocation provider

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,69 @@ 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. On gfx950:
+
+| 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. Only the default backend has been run through them.
+
+**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)
+```
+
+### 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.

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:

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]

iris/experimental/torch_symm_mem_provider.py

@@ -0,0 +1,147 @@
+# 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):
+        """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

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

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,
 }