From 96372f51c3ae2ffedb1e5ef2772411a8c00b82fe Mon Sep 17 00:00:00 2001
From: Muhammad Awad <muhaawad@amd.com>
Date: Mon, 31 Aug 2026 14:03:17 -0700
Subject: [PATCH 1/5] Add allocator-agnostic symmetric address descriptor

Introduce SymmetricAddressMap, the normalized per-allocation address
metadata that device-side translation needs, plus allocate_symmetric()
and get_symmetric_address_map() on the Iris context.

Device translation only ever needs to subtract a local backing-allocation
base and add a peer base. Tying that metadata to the allocation rather
than to the Iris context is what lets the same kernel run over tensors
from different providers - Iris today, rocSHMEM or Torch Symmetric Memory
behind the same descriptor later.

The tensor and its map are returned as two values and passed to kernels
as two arguments rather than bound into one struct, so this works on the
older Triton releases pinned in several test environments.

Additive: get_heap_bases() and the existing RMA APIs are unchanged.

Refs #546
---
 iris/__init__.py                           |   8 ++
 iris/host/iris.py                          |  80 +++++++++++++
 iris/host/memory/address_map.py            | 125 ++++++++++++++++++++
 tests/unittests/test_allocate_symmetric.py | 128 +++++++++++++++++++++
 4 files changed, 341 insertions(+)
 create mode 100644 iris/host/memory/address_map.py
 create mode 100644 tests/unittests/test_allocate_symmetric.py

diff --git a/iris/__init__.py b/iris/__init__.py
index 7e4047ec4..0917f48ba 100644
--- a/iris/__init__.py
+++ b/iris/__init__.py
@@ -43,6 +43,11 @@
 """
 
 from iris.host.iris import Iris, iris
+from iris.host.memory.address_map import (
+    SymmetricAddressMap,
+    CAP_REMOTE_LOAD_STORE,
+    CAP_REMOTE_ATOMICS,
+)
 from iris.mem.triton.context import Context, Context as DeviceContext
 from iris.host.tracing.events import TraceEvent
 from iris.mem.triton.types import (
@@ -97,6 +102,9 @@
 __all__ = [
     "Iris",
     "iris",
+    "SymmetricAddressMap",
+    "CAP_REMOTE_LOAD_STORE",
+    "CAP_REMOTE_ATOMICS",
     "get_device_id_for_rank",
     "Context",
     "DeviceContext",
diff --git a/iris/host/iris.py b/iris/host/iris.py
index 1c9843fbe..d80c8e61b 100644
--- a/iris/host/iris.py
+++ b/iris/host/iris.py
@@ -53,6 +53,7 @@
     count_devices,
 )
 from iris.host.memory.symmetric_heap import SymmetricHeap
+from iris.host.memory.address_map import SymmetricAddressMap
 import numpy as np
 from typing import Any
 import torch
@@ -918,6 +919,85 @@ def get_heap_bases(self):
         """
         return self.heap_bases
 
+    def get_symmetric_address_map(self, tensor: torch.Tensor) -> SymmetricAddressMap:
+        """
+        Return the normalized address descriptor for a symmetric tensor.
+
+        This is the provider-facing half of the allocator-agnostic interface: it
+        turns an Iris-allocated tensor into the same :class:`SymmetricAddressMap`
+        that rocSHMEM or Torch Symmetric Memory adapters will produce, so device
+        code can translate addresses without knowing which allocator was used.
+
+        Iris currently maps every symmetric tensor through one context-wide heap,
+        so the descriptor for any Iris tensor describes that heap. The shape is
+        per-allocation regardless, which is what lets a single kernel consume
+        tensors backed by different providers.
+
+        Args:
+            tensor (torch.Tensor): Tensor on the Iris symmetric heap.
+
+        Returns:
+            SymmetricAddressMap: Descriptor for the tensor's backing allocation.
+
+        Raises:
+            ValueError: If the tensor is not on the Iris symmetric heap.
+
+        Example:
+            >>> ctx = iris.iris(1 << 20)
+            >>> tensor = ctx.zeros(1024, dtype=torch.float32)
+            >>> address_map = ctx.get_symmetric_address_map(tensor)
+            >>> address_map.peer_bases[address_map.local_rank] == address_map.allocation_base
+        """
+        if not self.is_symmetric(tensor):
+            raise ValueError(
+                "tensor is not on the Iris symmetric heap; allocate it with an Iris "
+                "creation op or import it with as_symmetric()"
+            )
+
+        return SymmetricAddressMap(
+            peer_bases=self.heap_bases,
+            local_rank=self.cur_rank,
+            allocation_base=int(self.heap_bases[self.cur_rank].item()),
+            allocation_bytes=self.heap_size,
+        )
+
+    def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, SymmetricAddressMap]:
+        """
+        Allocate a symmetric tensor together with its address descriptor.
+
+        This is the allocator-agnostic allocation entry point. Every provider
+        exposes the same call, so kernels written against the returned pair run
+        unchanged on tensors from any of them::
+
+            tensor, address_map = provider.allocate_symmetric(1024, dtype=torch.float32)
+            kernel[grid](tensor, address_map.peer_bases, target_rank, ...)
+
+        The tensor and its map are passed to kernels as two separate arguments
+        -- a pointer and a tensor -- rather than bound into one struct, so that
+        this works on the older Triton releases pinned in several environments.
+
+        Args:
+            *size (int...): Shape of the tensor, as a sequence of integers or a
+                single collection.
+            dtype (torch.dtype, optional): Element type. Defaults to the torch
+                default dtype.
+
+        Returns:
+            tuple[torch.Tensor, SymmetricAddressMap]: The tensor and the
+            descriptor for its backing allocation.
+
+        Note:
+            Collective. All ranks must call this together, as with the other
+            Iris allocation ops.
+
+        Example:
+            >>> ctx = iris.iris(1 << 20)
+            >>> tensor, address_map = ctx.allocate_symmetric(1024, dtype=torch.float32)
+            >>> tensor.shape  # torch.Size([1024])
+        """
+        tensor = self.zeros(*size, dtype=dtype)
+        return tensor, self.get_symmetric_address_map(tensor)
+
     def _build_device_context(self):
         """
         Build and cache the device context tensor.
diff --git a/iris/host/memory/address_map.py b/iris/host/memory/address_map.py
new file mode 100644
index 000000000..88f8fdbdb
--- /dev/null
+++ b/iris/host/memory/address_map.py
@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+
+"""
+Allocator-agnostic symmetric address metadata.
+
+Iris device-side translation only needs to answer one question: given a pointer
+that is local to this rank, what address reaches the same location on a peer?
+The answer is the same arithmetic regardless of which host allocator produced
+the tensor::
+
+    offset = local_pointer - local_allocation_base
+    remote_pointer = peer_bases[target_rank] + offset
+
+:class:`SymmetricAddressMap` is the normalized form of the metadata that
+arithmetic requires. Providers (the Iris allocators today, rocSHMEM or Torch
+Symmetric Memory later) own allocation, lifetime, handle exchange and peer-base
+production; Iris consumes only this descriptor.
+
+The descriptor is attached to a tensor's *backing allocation*, not to the Iris
+context, so a single kernel can consume tensors from different providers with
+different peer-base tables while running identical device code.
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+#: The provider guarantees remote loads and stores through translated pointers.
+CAP_REMOTE_LOAD_STORE = 1 << 0
+
+#: The provider guarantees remote atomics through translated pointers.
+CAP_REMOTE_ATOMICS = 1 << 1
+
+
+@dataclass(frozen=True)
+class SymmetricAddressMap:
+    """
+    Normalized peer address metadata for one backing allocation.
+
+    Args:
+        peer_bases (torch.Tensor): ``int64[world_size]`` device-resident tensor
+            of backing-allocation base addresses, indexed by rank. This is the
+            only field device code needs; pass it to kernels alongside the
+            tensor pointer. Device-side translation casts the pointer to
+            ``tl.uint64`` before subtracting, so signed storage is fine.
+        local_rank (int): Rank of the calling process.
+        allocation_base (int): Base address of the backing allocation on this
+            rank. Views translate against this, not against the view pointer.
+        allocation_bytes (int): Size of the backing allocation in bytes.
+        capabilities (int): Bitmask of ``CAP_*`` flags the provider guarantees.
+
+    Invariant:
+        ``peer_bases[local_rank] == allocation_base``. Translation subtracts the
+        local base and adds the peer base, so a descriptor that violates this
+        produces silently wrong remote addresses.
+
+    Example:
+        >>> tensor, address_map = ctx.allocate_symmetric(1024, dtype=torch.float32)
+        >>> kernel[grid](tensor, address_map.peer_bases, target_rank, ...)
+    """
+
+    peer_bases: torch.Tensor
+    local_rank: int
+    allocation_base: int
+    allocation_bytes: int
+    capabilities: int = CAP_REMOTE_LOAD_STORE | CAP_REMOTE_ATOMICS
+
+    def __post_init__(self):
+        if self.peer_bases.dtype not in (torch.int64, torch.uint64):
+            raise ValueError(f"peer_bases must be int64 or uint64, got {self.peer_bases.dtype}")
+        if not 0 <= self.local_rank < self.peer_bases.numel():
+            raise ValueError(f"local_rank {self.local_rank} out of range for {self.peer_bases.numel()} ranks")
+
+        local_base = int(self.peer_bases[self.local_rank].item())
+        if local_base != self.allocation_base:
+            raise ValueError(
+                f"peer_bases[{self.local_rank}]={hex(local_base)} does not match "
+                f"allocation_base={hex(self.allocation_base)}; translation would "
+                "produce wrong remote addresses"
+            )
+
+    @property
+    def world_size(self) -> int:
+        """Number of ranks in the peer-base table."""
+        return self.peer_bases.numel()
+
+    def owns(self, tensor: torch.Tensor) -> bool:
+        """
+        Check that ``tensor`` lies inside the allocation this map describes.
+
+        Passing a tensor and an address map as separate kernel arguments allows
+        them to be paired by mistake -- a tensor from one provider with another
+        provider's map. Such a pairing translates against the wrong base and
+        corrupts memory silently. Different providers hand out disjoint address
+        ranges, so a bounds check against the backing allocation catches it.
+
+        This is a host-side check against metadata that is fully known before
+        launch; it costs nothing at runtime. It cannot catch a mispairing
+        between two allocations that share a backing range, which under the
+        context-wide heap means two Iris tensors -- but those share a base, so
+        translating one against the other's map is harmless today and becomes
+        detectable once per-allocation bases replace the shared heap.
+
+        Args:
+            tensor (torch.Tensor): Tensor to check.
+
+        Returns:
+            bool: True if the tensor's storage lies within the backing allocation.
+        """
+        start = tensor.data_ptr()
+        end = start + tensor.numel() * tensor.element_size()
+        return self.allocation_base <= start and end <= self.allocation_base + self.allocation_bytes
+
+    def supports(self, capability: int) -> bool:
+        """
+        Check whether the provider guarantees a capability.
+
+        Args:
+            capability (int): One of the ``CAP_*`` flags.
+
+        Returns:
+            bool: True if the flag is set.
+        """
+        return bool(self.capabilities & capability)
diff --git a/tests/unittests/test_allocate_symmetric.py b/tests/unittests/test_allocate_symmetric.py
new file mode 100644
index 000000000..caf3b2e73
--- /dev/null
+++ b/tests/unittests/test_allocate_symmetric.py
@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+
+"""
+Test the allocator-agnostic allocate_symmetric() API.
+
+The kernel here deliberately does not call any Iris translation helper. It
+takes a pointer and a peer-base table as two ordinary kernel arguments and
+inlines the translation, which is the form the allocator-agnostic interface is
+specified against: identical device code has to work for a tensor from any
+provider, and hoisting the base loads out of the access path is only possible
+when the translation is written by hand.
+"""
+
+import gc
+
+import pytest
+import torch
+import triton
+import triton.language as tl
+
+import iris
+
+
+@triton.jit
+def _put_translated_kernel(
+    src,
+    dst,
+    dst_peer_bases,
+    n_elements,
+    target_rank,
+    CUR_RANK: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Copy src into dst on target_rank, translating dst by hand."""
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    # Hoisted: both base loads happen once, not once per access. target_rank is
+    # a runtime argument because a collective loops over peers; specializing on
+    # it would compile a separate kernel per destination.
+    local_base = tl.load(dst_peer_bases + CUR_RANK)
+    remote_base = tl.load(dst_peer_bases + target_rank)
+
+    # Translate the allocation, then index into it.
+    offset = tl.cast(dst, tl.uint64) - local_base
+    remote_base_byte = tl.cast(remote_base, tl.pointer_type(tl.int8))
+    remote_dst = tl.cast(remote_base_byte + offset, dst.dtype)
+
+    values = tl.load(src + offsets, mask=mask)
+    tl.store(remote_dst + offsets, values, mask=mask)
+
+
+def test_allocate_symmetric_descriptor():
+    """The descriptor describes the tensor's backing allocation."""
+    shmem = iris.iris(1 << 20)
+
+    try:
+        tensor, address_map = shmem.allocate_symmetric(1024, dtype=torch.float32)
+
+        assert tensor.shape == (1024,)
+        assert tensor.dtype == torch.float32
+        assert shmem.is_symmetric(tensor)
+
+        # The invariant device translation depends on: subtracting the local
+        # base and adding a peer base is only correct if these agree.
+        local_base = int(address_map.peer_bases[address_map.local_rank].item())
+        assert local_base == address_map.allocation_base
+
+        assert address_map.local_rank == shmem.get_rank()
+        assert address_map.world_size == shmem.get_num_ranks()
+        assert address_map.owns(tensor)
+
+        # A tensor outside the heap is not covered by this allocation.
+        external = torch.zeros(1024, dtype=torch.float32, device=shmem.get_device())
+        assert not address_map.owns(external)
+    finally:
+        shmem.barrier()
+        del shmem
+        gc.collect()
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+def test_allocate_symmetric_remote_put(dtype):
+    """A kernel given (pointer, peer_bases) reaches the right peer allocation."""
+    shmem = iris.iris(1 << 24)
+    rank = shmem.get_rank()
+    world_size = shmem.get_num_ranks()
+    target_rank = (rank + 1) % world_size
+
+    n_elements = 256
+    block_size = 256
+
+    try:
+        src, _ = shmem.allocate_symmetric(n_elements, dtype=dtype)
+        dst, dst_map = shmem.allocate_symmetric(n_elements, dtype=dtype)
+
+        # Each rank stamps a distinct value so a write landing on the wrong
+        # rank produces the wrong answer rather than a plausible one.
+        src.fill_(rank + 1)
+        dst.fill_(-1)
+        shmem.barrier()
+
+        _put_translated_kernel[(1,)](
+            src,
+            dst,
+            dst_map.peer_bases,
+            n_elements,
+            target_rank,
+            CUR_RANK=rank,
+            BLOCK_SIZE=block_size,
+        )
+        shmem.barrier()
+
+        # We received from the rank that targets us, not from the rank we target.
+        source_rank = (rank - 1) % world_size
+        expected = torch.full_like(dst, source_rank + 1)
+        torch.testing.assert_close(dst, expected)
+
+        if world_size > 1:
+            # Pin the failure mode: a translation that silently resolved to the
+            # local allocation would leave this rank's own value here, and
+            # assert_close above would still pass at world_size 1.
+            assert dst[0].item() != rank + 1
+    finally:
+        shmem.barrier()
+        del shmem
+        gc.collect()

From 8b87c46ff671ff4877b891851e24fe821d4740aa Mon Sep 17 00:00:00 2001
From: Muhammad Awad <muhaawad@amd.com>
Date: Mon, 31 Aug 2026 14:26:06 -0700
Subject: [PATCH 2/5] Re-apply the contiguity hint that inlined translation
 drops

iris.load/store pass a hint through __translate which applies
tl.multiple_of/tl.max_contiguous to the translated pointer, and every
production collective uses it. Inlining the translation loses that
unless the kernel puts it back, so the example puts it back.

It has to go on the indexed pointer, not the translated base:
translating the allocation once and indexing after leaves the pointer
scalar, and max_contiguous requires a block matching the hint shape.

Refs #546
---
 tests/unittests/test_allocate_symmetric.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/tests/unittests/test_allocate_symmetric.py b/tests/unittests/test_allocate_symmetric.py
index caf3b2e73..ce9ef4f76 100644
--- a/tests/unittests/test_allocate_symmetric.py
+++ b/tests/unittests/test_allocate_symmetric.py
@@ -47,8 +47,19 @@ def _put_translated_kernel(
     remote_base_byte = tl.cast(remote_base, tl.pointer_type(tl.int8))
     remote_dst = tl.cast(remote_base_byte + offset, dst.dtype)
 
+    # Re-apply the contiguity hint by hand. iris.load/store take a `hint` that
+    # does this inside __translate, and every production collective passes one;
+    # inlining the translation drops it unless the author puts it back. Manual
+    # translation buys control over hoisting, not a free win over the helper.
+    #
+    # It has to go here rather than on remote_dst: translating the base once and
+    # indexing after leaves the translated pointer scalar, and max_contiguous
+    # takes a block whose shape matches the hint.
+    remote_ptrs = remote_dst + offsets
+    remote_ptrs = tl.max_contiguous(tl.multiple_of(remote_ptrs, BLOCK_SIZE), BLOCK_SIZE)
+
     values = tl.load(src + offsets, mask=mask)
-    tl.store(remote_dst + offsets, values, mask=mask)
+    tl.store(remote_ptrs, values, mask=mask)
 
 
 def test_allocate_symmetric_descriptor():

From eb071ddc8247c7e3f9cf35a3c74be10b8cbe95ee Mon Sep 17 00:00:00 2001
From: Muhammad Awad <muhaawad@amd.com>
Date: Mon, 31 Aug 2026 15:27:46 -0700
Subject: [PATCH 3/5] Drop the descriptor struct and return the peer base table
 directly

Every field on SymmetricAddressMap except peer_bases was something the
caller already had: local_rank was get_rank(), allocation_base was
peer_bases[local_rank] by the stated invariant, allocation_bytes was the
heap size the caller passed to iris(), and capabilities was a constant.
A struct that carries no information the caller lacks is a struct to
delete, and deleting it removes a device sync and a validator whose main
check compared a value against itself.

allocate_symmetric() now returns (tensor, peer_bases) and
get_peer_bases(tensor) replaces get_symmetric_address_map(). The contract
gets simpler for a second provider to satisfy: produce a device-resident
int64 table indexed by rank, rather than fill in a dataclass with an
invariant and a capability mask.

Note this diverges from #546, which names SymmetricAddressMap as the
provider-facing contract. The descriptor is worth introducing when a
second provider exists and it has something to normalize.

Refs #546
---
 iris/__init__.py                           |   8 --
 iris/host/iris.py                          |  59 +++++-----
 iris/host/memory/address_map.py            | 125 ---------------------
 tests/unittests/test_allocate_symmetric.py |  84 +++++++++++---
 4 files changed, 95 insertions(+), 181 deletions(-)
 delete mode 100644 iris/host/memory/address_map.py

diff --git a/iris/__init__.py b/iris/__init__.py
index 0917f48ba..7e4047ec4 100644
--- a/iris/__init__.py
+++ b/iris/__init__.py
@@ -43,11 +43,6 @@
 """
 
 from iris.host.iris import Iris, iris
-from iris.host.memory.address_map import (
-    SymmetricAddressMap,
-    CAP_REMOTE_LOAD_STORE,
-    CAP_REMOTE_ATOMICS,
-)
 from iris.mem.triton.context import Context, Context as DeviceContext
 from iris.host.tracing.events import TraceEvent
 from iris.mem.triton.types import (
@@ -102,9 +97,6 @@
 __all__ = [
     "Iris",
     "iris",
-    "SymmetricAddressMap",
-    "CAP_REMOTE_LOAD_STORE",
-    "CAP_REMOTE_ATOMICS",
     "get_device_id_for_rank",
     "Context",
     "DeviceContext",
diff --git a/iris/host/iris.py b/iris/host/iris.py
index d80c8e61b..3badba7fb 100644
--- a/iris/host/iris.py
+++ b/iris/host/iris.py
@@ -53,7 +53,6 @@
     count_devices,
 )
 from iris.host.memory.symmetric_heap import SymmetricHeap
-from iris.host.memory.address_map import SymmetricAddressMap
 import numpy as np
 from typing import Any
 import torch
@@ -919,25 +918,28 @@ def get_heap_bases(self):
         """
         return self.heap_bases
 
-    def get_symmetric_address_map(self, tensor: torch.Tensor) -> SymmetricAddressMap:
+    def get_peer_bases(self, tensor: torch.Tensor) -> torch.Tensor:
         """
-        Return the normalized address descriptor for a symmetric tensor.
+        Return the peer base-address table for a symmetric tensor.
 
-        This is the provider-facing half of the allocator-agnostic interface: it
-        turns an Iris-allocated tensor into the same :class:`SymmetricAddressMap`
-        that rocSHMEM or Torch Symmetric Memory adapters will produce, so device
-        code can translate addresses without knowing which allocator was used.
+        This is the tensor-scoped form of :meth:`get_heap_bases`, and the host
+        half of the allocator-agnostic interface. Device translation needs one
+        thing -- a device-resident table of backing-allocation bases indexed by
+        rank -- and asking a tensor for its table rather than asking the context
+        for a global one is what lets a single kernel consume tensors from
+        different providers.
 
-        Iris currently maps every symmetric tensor through one context-wide heap,
-        so the descriptor for any Iris tensor describes that heap. The shape is
-        per-allocation regardless, which is what lets a single kernel consume
-        tensors backed by different providers.
+        Iris maps every symmetric tensor through one context-wide heap today, so
+        every Iris tensor returns the same table. The call shape is what matters:
+        kernels written against it stop depending on the Iris context.
 
         Args:
             tensor (torch.Tensor): Tensor on the Iris symmetric heap.
 
         Returns:
-            SymmetricAddressMap: Descriptor for the tensor's backing allocation.
+            torch.Tensor: ``int64[world_size]`` device-resident base addresses,
+            indexed by rank. ``peer_bases[cur_rank]`` is this rank's base, which
+            is the value device-side translation subtracts.
 
         Raises:
             ValueError: If the tensor is not on the Iris symmetric heap.
@@ -945,36 +947,29 @@ def get_symmetric_address_map(self, tensor: torch.Tensor) -> SymmetricAddressMap
         Example:
             >>> ctx = iris.iris(1 << 20)
             >>> tensor = ctx.zeros(1024, dtype=torch.float32)
-            >>> address_map = ctx.get_symmetric_address_map(tensor)
-            >>> address_map.peer_bases[address_map.local_rank] == address_map.allocation_base
+            >>> peer_bases = ctx.get_peer_bases(tensor)
         """
         if not self.is_symmetric(tensor):
             raise ValueError(
                 "tensor is not on the Iris symmetric heap; allocate it with an Iris "
                 "creation op or import it with as_symmetric()"
             )
+        return self.heap_bases
 
-        return SymmetricAddressMap(
-            peer_bases=self.heap_bases,
-            local_rank=self.cur_rank,
-            allocation_base=int(self.heap_bases[self.cur_rank].item()),
-            allocation_bytes=self.heap_size,
-        )
-
-    def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, SymmetricAddressMap]:
+    def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Tensor]:
         """
-        Allocate a symmetric tensor together with its address descriptor.
+        Allocate a symmetric tensor together with its peer base table.
 
         This is the allocator-agnostic allocation entry point. Every provider
         exposes the same call, so kernels written against the returned pair run
         unchanged on tensors from any of them::
 
-            tensor, address_map = provider.allocate_symmetric(1024, dtype=torch.float32)
-            kernel[grid](tensor, address_map.peer_bases, target_rank, ...)
+            tensor, peer_bases = provider.allocate_symmetric(1024, dtype=torch.float32)
+            kernel[grid](tensor, peer_bases, target_rank, CUR_RANK=rank, ...)
 
-        The tensor and its map are passed to kernels as two separate arguments
-        -- a pointer and a tensor -- rather than bound into one struct, so that
-        this works on the older Triton releases pinned in several environments.
+        The tensor and its table are passed to kernels as two separate arguments
+        -- a pointer and a tensor -- rather than bound into one struct, so this
+        works on the older Triton releases pinned in several environments.
 
         Args:
             *size (int...): Shape of the tensor, as a sequence of integers or a
@@ -983,8 +978,8 @@ def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, Symmetric
                 default dtype.
 
         Returns:
-            tuple[torch.Tensor, SymmetricAddressMap]: The tensor and the
-            descriptor for its backing allocation.
+            tuple[torch.Tensor, torch.Tensor]: The tensor, and the ``int64``
+            peer base table for its backing allocation.
 
         Note:
             Collective. All ranks must call this together, as with the other
@@ -992,11 +987,11 @@ def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, Symmetric
 
         Example:
             >>> ctx = iris.iris(1 << 20)
-            >>> tensor, address_map = ctx.allocate_symmetric(1024, dtype=torch.float32)
+            >>> tensor, peer_bases = ctx.allocate_symmetric(1024, dtype=torch.float32)
             >>> tensor.shape  # torch.Size([1024])
         """
         tensor = self.zeros(*size, dtype=dtype)
-        return tensor, self.get_symmetric_address_map(tensor)
+        return tensor, self.get_peer_bases(tensor)
 
     def _build_device_context(self):
         """
diff --git a/iris/host/memory/address_map.py b/iris/host/memory/address_map.py
deleted file mode 100644
index 88f8fdbdb..000000000
--- a/iris/host/memory/address_map.py
+++ /dev/null
@@ -1,125 +0,0 @@
-# SPDX-License-Identifier: MIT
-# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
-
-"""
-Allocator-agnostic symmetric address metadata.
-
-Iris device-side translation only needs to answer one question: given a pointer
-that is local to this rank, what address reaches the same location on a peer?
-The answer is the same arithmetic regardless of which host allocator produced
-the tensor::
-
-    offset = local_pointer - local_allocation_base
-    remote_pointer = peer_bases[target_rank] + offset
-
-:class:`SymmetricAddressMap` is the normalized form of the metadata that
-arithmetic requires. Providers (the Iris allocators today, rocSHMEM or Torch
-Symmetric Memory later) own allocation, lifetime, handle exchange and peer-base
-production; Iris consumes only this descriptor.
-
-The descriptor is attached to a tensor's *backing allocation*, not to the Iris
-context, so a single kernel can consume tensors from different providers with
-different peer-base tables while running identical device code.
-"""
-
-from dataclasses import dataclass
-
-import torch
-
-#: The provider guarantees remote loads and stores through translated pointers.
-CAP_REMOTE_LOAD_STORE = 1 << 0
-
-#: The provider guarantees remote atomics through translated pointers.
-CAP_REMOTE_ATOMICS = 1 << 1
-
-
-@dataclass(frozen=True)
-class SymmetricAddressMap:
-    """
-    Normalized peer address metadata for one backing allocation.
-
-    Args:
-        peer_bases (torch.Tensor): ``int64[world_size]`` device-resident tensor
-            of backing-allocation base addresses, indexed by rank. This is the
-            only field device code needs; pass it to kernels alongside the
-            tensor pointer. Device-side translation casts the pointer to
-            ``tl.uint64`` before subtracting, so signed storage is fine.
-        local_rank (int): Rank of the calling process.
-        allocation_base (int): Base address of the backing allocation on this
-            rank. Views translate against this, not against the view pointer.
-        allocation_bytes (int): Size of the backing allocation in bytes.
-        capabilities (int): Bitmask of ``CAP_*`` flags the provider guarantees.
-
-    Invariant:
-        ``peer_bases[local_rank] == allocation_base``. Translation subtracts the
-        local base and adds the peer base, so a descriptor that violates this
-        produces silently wrong remote addresses.
-
-    Example:
-        >>> tensor, address_map = ctx.allocate_symmetric(1024, dtype=torch.float32)
-        >>> kernel[grid](tensor, address_map.peer_bases, target_rank, ...)
-    """
-
-    peer_bases: torch.Tensor
-    local_rank: int
-    allocation_base: int
-    allocation_bytes: int
-    capabilities: int = CAP_REMOTE_LOAD_STORE | CAP_REMOTE_ATOMICS
-
-    def __post_init__(self):
-        if self.peer_bases.dtype not in (torch.int64, torch.uint64):
-            raise ValueError(f"peer_bases must be int64 or uint64, got {self.peer_bases.dtype}")
-        if not 0 <= self.local_rank < self.peer_bases.numel():
-            raise ValueError(f"local_rank {self.local_rank} out of range for {self.peer_bases.numel()} ranks")
-
-        local_base = int(self.peer_bases[self.local_rank].item())
-        if local_base != self.allocation_base:
-            raise ValueError(
-                f"peer_bases[{self.local_rank}]={hex(local_base)} does not match "
-                f"allocation_base={hex(self.allocation_base)}; translation would "
-                "produce wrong remote addresses"
-            )
-
-    @property
-    def world_size(self) -> int:
-        """Number of ranks in the peer-base table."""
-        return self.peer_bases.numel()
-
-    def owns(self, tensor: torch.Tensor) -> bool:
-        """
-        Check that ``tensor`` lies inside the allocation this map describes.
-
-        Passing a tensor and an address map as separate kernel arguments allows
-        them to be paired by mistake -- a tensor from one provider with another
-        provider's map. Such a pairing translates against the wrong base and
-        corrupts memory silently. Different providers hand out disjoint address
-        ranges, so a bounds check against the backing allocation catches it.
-
-        This is a host-side check against metadata that is fully known before
-        launch; it costs nothing at runtime. It cannot catch a mispairing
-        between two allocations that share a backing range, which under the
-        context-wide heap means two Iris tensors -- but those share a base, so
-        translating one against the other's map is harmless today and becomes
-        detectable once per-allocation bases replace the shared heap.
-
-        Args:
-            tensor (torch.Tensor): Tensor to check.
-
-        Returns:
-            bool: True if the tensor's storage lies within the backing allocation.
-        """
-        start = tensor.data_ptr()
-        end = start + tensor.numel() * tensor.element_size()
-        return self.allocation_base <= start and end <= self.allocation_base + self.allocation_bytes
-
-    def supports(self, capability: int) -> bool:
-        """
-        Check whether the provider guarantees a capability.
-
-        Args:
-            capability (int): One of the ``CAP_*`` flags.
-
-        Returns:
-            bool: True if the flag is set.
-        """
-        return bool(self.capabilities & capability)
diff --git a/tests/unittests/test_allocate_symmetric.py b/tests/unittests/test_allocate_symmetric.py
index ce9ef4f76..be4f03f06 100644
--- a/tests/unittests/test_allocate_symmetric.py
+++ b/tests/unittests/test_allocate_symmetric.py
@@ -62,49 +62,101 @@ def _put_translated_kernel(
     tl.store(remote_ptrs, values, mask=mask)
 
 
-def test_allocate_symmetric_descriptor():
-    """The descriptor describes the tensor's backing allocation."""
+def test_allocate_symmetric_returns_peer_bases():
+    """The table is device-resident, rank-indexed, and holds our own base."""
     shmem = iris.iris(1 << 20)
 
     try:
-        tensor, address_map = shmem.allocate_symmetric(1024, dtype=torch.float32)
+        tensor, peer_bases = shmem.allocate_symmetric(1024, dtype=torch.float32)
 
         assert tensor.shape == (1024,)
         assert tensor.dtype == torch.float32
         assert shmem.is_symmetric(tensor)
 
-        # The invariant device translation depends on: subtracting the local
-        # base and adding a peer base is only correct if these agree.
-        local_base = int(address_map.peer_bases[address_map.local_rank].item())
-        assert local_base == address_map.allocation_base
+        assert peer_bases.numel() == shmem.get_num_ranks()
+        assert peer_bases.dtype in (torch.int64, torch.uint64)
+        assert peer_bases.is_cuda
 
-        assert address_map.local_rank == shmem.get_rank()
-        assert address_map.world_size == shmem.get_num_ranks()
-        assert address_map.owns(tensor)
+        # peer_bases[cur_rank] is the base translation subtracts, so it has to
+        # be this rank's own heap base and the tensor has to sit inside it.
+        local_base = int(peer_bases[shmem.get_rank()].item())
+        assert local_base == int(shmem.get_heap_bases()[shmem.get_rank()].item())
+        assert tensor.data_ptr() >= local_base
 
-        # A tensor outside the heap is not covered by this allocation.
+        # An external tensor has no table to hand out.
         external = torch.zeros(1024, dtype=torch.float32, device=shmem.get_device())
-        assert not address_map.owns(external)
+        with pytest.raises(ValueError, match="not on the Iris symmetric heap"):
+            shmem.get_peer_bases(external)
     finally:
         shmem.barrier()
         del shmem
         gc.collect()
 
 
+def test_view_translates_against_allocation_root():
+    """A view translates against the allocation root, not its own pointer.
+
+    The kernel subtracts peer_bases[local_rank] -- the allocation base -- so a
+    view's offset within the allocation survives translation. Subtracting the
+    view pointer instead would land at the start of the peer's allocation.
+    """
+    shmem = iris.iris(1 << 24)
+    rank = shmem.get_rank()
+    world_size = shmem.get_num_ranks()
+    target_rank = (rank + 1) % world_size
+
+    n_elements = 512
+    offset = 128
+
+    try:
+        src, _ = shmem.allocate_symmetric(n_elements, dtype=torch.float32)
+        dst, dst_peer_bases = shmem.allocate_symmetric(n_elements, dtype=torch.float32)
+
+        # A view into the middle of the allocation. Its data_ptr is not the
+        # allocation base, but it translates against the same table.
+        view = dst[offset:]
+        assert view.data_ptr() != int(dst_peer_bases[rank].item())
+
+        src.fill_(rank + 1)
+        dst.fill_(-1)
+        shmem.barrier()
+
+        _put_translated_kernel[(1,)](
+            src,
+            view,
+            dst_peer_bases,
+            view.numel(),
+            target_rank,
+            CUR_RANK=rank,
+            BLOCK_SIZE=512,
+        )
+        shmem.barrier()
+
+        source_rank = (rank - 1) % world_size
+        # The view's region received the peer's data; everything before it did not.
+        torch.testing.assert_close(dst[offset:], torch.full_like(dst[offset:], source_rank + 1))
+        torch.testing.assert_close(dst[:offset], torch.full_like(dst[:offset], -1.0))
+    finally:
+        shmem.barrier()
+        del shmem
+        gc.collect()
+
+
+@pytest.mark.parametrize("n_elements", [256, 200])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
-def test_allocate_symmetric_remote_put(dtype):
+def test_allocate_symmetric_remote_put(dtype, n_elements):
     """A kernel given (pointer, peer_bases) reaches the right peer allocation."""
     shmem = iris.iris(1 << 24)
     rank = shmem.get_rank()
     world_size = shmem.get_num_ranks()
     target_rank = (rank + 1) % world_size
 
-    n_elements = 256
+    # 200 is not a multiple of the block, so the mask is actually exercised.
     block_size = 256
 
     try:
         src, _ = shmem.allocate_symmetric(n_elements, dtype=dtype)
-        dst, dst_map = shmem.allocate_symmetric(n_elements, dtype=dtype)
+        dst, dst_peer_bases = shmem.allocate_symmetric(n_elements, dtype=dtype)
 
         # Each rank stamps a distinct value so a write landing on the wrong
         # rank produces the wrong answer rather than a plausible one.
@@ -115,7 +167,7 @@ def test_allocate_symmetric_remote_put(dtype):
         _put_translated_kernel[(1,)](
             src,
             dst,
-            dst_map.peer_bases,
+            dst_peer_bases,
             n_elements,
             target_rank,
             CUR_RANK=rank,

From a7fe0bb1e3c18d855b494e41b5023035a9543707 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <muhaawad@amd.com>
Date: Mon, 31 Aug 2026 16:11:34 -0700
Subject: [PATCH 4/5] Reduce allocate_symmetric to one method

get_peer_bases() returned self.heap_bases and ignored its tensor argument
except to validate it, and that guard could not fail on the only path that
called it: allocate_symmetric passed a tensor it had just allocated. A
caller who wants the check has is_symmetric().

Allocate with empty() rather than zeros(): the API says nothing about
contents, so zeroing is work nobody asked for.

Use ctx over shmem in the new code, matching the docstrings.
get_heap_bases() and its 76 existing call sites are untouched.

Refs #546
---
 iris/host/iris.py                          |  62 ++---------
 tests/unittests/test_allocate_symmetric.py | 120 +++++++++------------
 2 files changed, 58 insertions(+), 124 deletions(-)

diff --git a/iris/host/iris.py b/iris/host/iris.py
index 3badba7fb..6474be380 100644
--- a/iris/host/iris.py
+++ b/iris/host/iris.py
@@ -918,58 +918,13 @@ def get_heap_bases(self):
         """
         return self.heap_bases
 
-    def get_peer_bases(self, tensor: torch.Tensor) -> torch.Tensor:
-        """
-        Return the peer base-address table for a symmetric tensor.
-
-        This is the tensor-scoped form of :meth:`get_heap_bases`, and the host
-        half of the allocator-agnostic interface. Device translation needs one
-        thing -- a device-resident table of backing-allocation bases indexed by
-        rank -- and asking a tensor for its table rather than asking the context
-        for a global one is what lets a single kernel consume tensors from
-        different providers.
-
-        Iris maps every symmetric tensor through one context-wide heap today, so
-        every Iris tensor returns the same table. The call shape is what matters:
-        kernels written against it stop depending on the Iris context.
-
-        Args:
-            tensor (torch.Tensor): Tensor on the Iris symmetric heap.
-
-        Returns:
-            torch.Tensor: ``int64[world_size]`` device-resident base addresses,
-            indexed by rank. ``peer_bases[cur_rank]`` is this rank's base, which
-            is the value device-side translation subtracts.
-
-        Raises:
-            ValueError: If the tensor is not on the Iris symmetric heap.
-
-        Example:
-            >>> ctx = iris.iris(1 << 20)
-            >>> tensor = ctx.zeros(1024, dtype=torch.float32)
-            >>> peer_bases = ctx.get_peer_bases(tensor)
-        """
-        if not self.is_symmetric(tensor):
-            raise ValueError(
-                "tensor is not on the Iris symmetric heap; allocate it with an Iris "
-                "creation op or import it with as_symmetric()"
-            )
-        return self.heap_bases
-
     def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Tensor]:
         """
-        Allocate a symmetric tensor together with its peer base table.
-
-        This is the allocator-agnostic allocation entry point. Every provider
-        exposes the same call, so kernels written against the returned pair run
-        unchanged on tensors from any of them::
-
-            tensor, peer_bases = provider.allocate_symmetric(1024, dtype=torch.float32)
-            kernel[grid](tensor, peer_bases, target_rank, CUR_RANK=rank, ...)
+        Allocate a symmetric tensor and return it with its peer-base table.
 
-        The tensor and its table are passed to kernels as two separate arguments
-        -- a pointer and a tensor -- rather than bound into one struct, so this
-        works on the older Triton releases pinned in several environments.
+        Kernels take the pair as two ordinary arguments -- a pointer and a
+        tensor -- and inline the address translation, so the same device code
+        works for a tensor from any provider that returns this shape.
 
         Args:
             *size (int...): Shape of the tensor, as a sequence of integers or a
@@ -978,8 +933,9 @@ def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Ten
                 default dtype.
 
         Returns:
-            tuple[torch.Tensor, torch.Tensor]: The tensor, and the ``int64``
-            peer base table for its backing allocation.
+            tuple[torch.Tensor, torch.Tensor]: The tensor, uninitialized, and an
+            ``int64`` device-resident table of base addresses indexed by rank.
+            ``peer_bases[cur_rank]`` is the base translation subtracts.
 
         Note:
             Collective. All ranks must call this together, as with the other
@@ -988,10 +944,8 @@ def allocate_symmetric(self, *size, dtype=None) -> tuple[torch.Tensor, torch.Ten
         Example:
             >>> ctx = iris.iris(1 << 20)
             >>> tensor, peer_bases = ctx.allocate_symmetric(1024, dtype=torch.float32)
-            >>> tensor.shape  # torch.Size([1024])
         """
-        tensor = self.zeros(*size, dtype=dtype)
-        return tensor, self.get_peer_bases(tensor)
+        return self.empty(*size, dtype=dtype), self.heap_bases
 
     def _build_device_context(self):
         """
diff --git a/tests/unittests/test_allocate_symmetric.py b/tests/unittests/test_allocate_symmetric.py
index be4f03f06..a9f342858 100644
--- a/tests/unittests/test_allocate_symmetric.py
+++ b/tests/unittests/test_allocate_symmetric.py
@@ -2,14 +2,11 @@
 # Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
 
 """
-Test the allocator-agnostic allocate_symmetric() API.
-
-The kernel here deliberately does not call any Iris translation helper. It
-takes a pointer and a peer-base table as two ordinary kernel arguments and
-inlines the translation, which is the form the allocator-agnostic interface is
-specified against: identical device code has to work for a tensor from any
-provider, and hoisting the base loads out of the access path is only possible
-when the translation is written by hand.
+Test allocate_symmetric().
+
+The kernel takes a pointer and a peer-base table as two ordinary arguments and
+inlines the translation, so the same device code works for a tensor from any
+provider.
 """
 
 import gc
@@ -32,29 +29,24 @@ def _put_translated_kernel(
     CUR_RANK: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
-    """Copy src into dst on target_rank, translating dst by hand."""
+    """Copy src into dst on target_rank.
+
+    target_rank is runtime: a collective loops over peers, and specializing on
+    it would compile one kernel per destination.
+    """
     offsets = tl.arange(0, BLOCK_SIZE)
     mask = offsets < n_elements
 
-    # Hoisted: both base loads happen once, not once per access. target_rank is
-    # a runtime argument because a collective loops over peers; specializing on
-    # it would compile a separate kernel per destination.
+    # Loaded once, outside the access path.
     local_base = tl.load(dst_peer_bases + CUR_RANK)
     remote_base = tl.load(dst_peer_bases + target_rank)
 
-    # Translate the allocation, then index into it.
+    # Same offset within the allocation, resolved against the peer's base.
     offset = tl.cast(dst, tl.uint64) - local_base
     remote_base_byte = tl.cast(remote_base, tl.pointer_type(tl.int8))
     remote_dst = tl.cast(remote_base_byte + offset, dst.dtype)
 
-    # Re-apply the contiguity hint by hand. iris.load/store take a `hint` that
-    # does this inside __translate, and every production collective passes one;
-    # inlining the translation drops it unless the author puts it back. Manual
-    # translation buys control over hoisting, not a free win over the helper.
-    #
-    # It has to go here rather than on remote_dst: translating the base once and
-    # indexing after leaves the translated pointer scalar, and max_contiguous
-    # takes a block whose shape matches the hint.
+    # Hint goes on the indexed pointers; remote_dst is still scalar here.
     remote_ptrs = remote_dst + offsets
     remote_ptrs = tl.max_contiguous(tl.multiple_of(remote_ptrs, BLOCK_SIZE), BLOCK_SIZE)
 
@@ -64,62 +56,52 @@ def _put_translated_kernel(
 
 def test_allocate_symmetric_returns_peer_bases():
     """The table is device-resident, rank-indexed, and holds our own base."""
-    shmem = iris.iris(1 << 20)
+    ctx = iris.iris(1 << 20)
 
     try:
-        tensor, peer_bases = shmem.allocate_symmetric(1024, dtype=torch.float32)
+        tensor, peer_bases = ctx.allocate_symmetric(1024, dtype=torch.float32)
 
         assert tensor.shape == (1024,)
         assert tensor.dtype == torch.float32
-        assert shmem.is_symmetric(tensor)
+        assert ctx.is_symmetric(tensor)
 
-        assert peer_bases.numel() == shmem.get_num_ranks()
+        assert peer_bases.numel() == ctx.get_num_ranks()
         assert peer_bases.dtype in (torch.int64, torch.uint64)
         assert peer_bases.is_cuda
 
-        # peer_bases[cur_rank] is the base translation subtracts, so it has to
-        # be this rank's own heap base and the tensor has to sit inside it.
-        local_base = int(peer_bases[shmem.get_rank()].item())
-        assert local_base == int(shmem.get_heap_bases()[shmem.get_rank()].item())
-        assert tensor.data_ptr() >= local_base
-
-        # An external tensor has no table to hand out.
-        external = torch.zeros(1024, dtype=torch.float32, device=shmem.get_device())
-        with pytest.raises(ValueError, match="not on the Iris symmetric heap"):
-            shmem.get_peer_bases(external)
+        # peer_bases[cur_rank] is what translation subtracts, so the tensor has
+        # to sit inside it.
+        assert tensor.data_ptr() >= int(peer_bases[ctx.get_rank()].item())
     finally:
-        shmem.barrier()
-        del shmem
+        ctx.barrier()
+        del ctx
         gc.collect()
 
 
 def test_view_translates_against_allocation_root():
-    """A view translates against the allocation root, not its own pointer.
+    """A view keeps its offset within the allocation across translation.
 
-    The kernel subtracts peer_bases[local_rank] -- the allocation base -- so a
-    view's offset within the allocation survives translation. Subtracting the
-    view pointer instead would land at the start of the peer's allocation.
+    Subtracting the view pointer instead of the allocation base would land at
+    the start of the peer's allocation.
     """
-    shmem = iris.iris(1 << 24)
-    rank = shmem.get_rank()
-    world_size = shmem.get_num_ranks()
+    ctx = iris.iris(1 << 24)
+    rank = ctx.get_rank()
+    world_size = ctx.get_num_ranks()
     target_rank = (rank + 1) % world_size
 
     n_elements = 512
     offset = 128
 
     try:
-        src, _ = shmem.allocate_symmetric(n_elements, dtype=torch.float32)
-        dst, dst_peer_bases = shmem.allocate_symmetric(n_elements, dtype=torch.float32)
+        src, _ = ctx.allocate_symmetric(n_elements, dtype=torch.float32)
+        dst, dst_peer_bases = ctx.allocate_symmetric(n_elements, dtype=torch.float32)
 
-        # A view into the middle of the allocation. Its data_ptr is not the
-        # allocation base, but it translates against the same table.
         view = dst[offset:]
         assert view.data_ptr() != int(dst_peer_bases[rank].item())
 
         src.fill_(rank + 1)
         dst.fill_(-1)
-        shmem.barrier()
+        ctx.barrier()
 
         _put_translated_kernel[(1,)](
             src,
@@ -130,15 +112,15 @@ def test_view_translates_against_allocation_root():
             CUR_RANK=rank,
             BLOCK_SIZE=512,
         )
-        shmem.barrier()
+        ctx.barrier()
 
         source_rank = (rank - 1) % world_size
         # The view's region received the peer's data; everything before it did not.
         torch.testing.assert_close(dst[offset:], torch.full_like(dst[offset:], source_rank + 1))
         torch.testing.assert_close(dst[:offset], torch.full_like(dst[:offset], -1.0))
     finally:
-        shmem.barrier()
-        del shmem
+        ctx.barrier()
+        del ctx
         gc.collect()
 
 
@@ -146,23 +128,23 @@ def test_view_translates_against_allocation_root():
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
 def test_allocate_symmetric_remote_put(dtype, n_elements):
     """A kernel given (pointer, peer_bases) reaches the right peer allocation."""
-    shmem = iris.iris(1 << 24)
-    rank = shmem.get_rank()
-    world_size = shmem.get_num_ranks()
+    ctx = iris.iris(1 << 24)
+    rank = ctx.get_rank()
+    world_size = ctx.get_num_ranks()
     target_rank = (rank + 1) % world_size
 
-    # 200 is not a multiple of the block, so the mask is actually exercised.
+    # 200 is not a multiple of the block, so the mask is exercised.
     block_size = 256
 
     try:
-        src, _ = shmem.allocate_symmetric(n_elements, dtype=dtype)
-        dst, dst_peer_bases = shmem.allocate_symmetric(n_elements, dtype=dtype)
+        src, _ = ctx.allocate_symmetric(n_elements, dtype=dtype)
+        dst, dst_peer_bases = ctx.allocate_symmetric(n_elements, dtype=dtype)
 
-        # Each rank stamps a distinct value so a write landing on the wrong
-        # rank produces the wrong answer rather than a plausible one.
+        # Distinct per rank, so a write landing on the wrong rank gives a wrong
+        # answer rather than a plausible one.
         src.fill_(rank + 1)
         dst.fill_(-1)
-        shmem.barrier()
+        ctx.barrier()
 
         _put_translated_kernel[(1,)](
             src,
@@ -173,19 +155,17 @@ def test_allocate_symmetric_remote_put(dtype, n_elements):
             CUR_RANK=rank,
             BLOCK_SIZE=block_size,
         )
-        shmem.barrier()
+        ctx.barrier()
 
-        # We received from the rank that targets us, not from the rank we target.
+        # We received from the rank targeting us, not the one we target.
         source_rank = (rank - 1) % world_size
-        expected = torch.full_like(dst, source_rank + 1)
-        torch.testing.assert_close(dst, expected)
+        torch.testing.assert_close(dst, torch.full_like(dst, source_rank + 1))
 
         if world_size > 1:
-            # Pin the failure mode: a translation that silently resolved to the
-            # local allocation would leave this rank's own value here, and
-            # assert_close above would still pass at world_size 1.
+            # A translation resolving to the local allocation would leave our
+            # own value here, and the check above would still pass at 1 rank.
             assert dst[0].item() != rank + 1
     finally:
-        shmem.barrier()
-        del shmem
+        ctx.barrier()
+        del ctx
         gc.collect()

From 5406e27fb109c43171c24e28ecc0b8ca720ccbf5 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <muhaawad@amd.com>
Date: Mon, 31 Aug 2026 16:20:56 -0700
Subject: [PATCH 5/5] Tighten the tests and drop an unsound hint

Three fixes:

The view test passed trivially at world size 1, where the sender is the
receiver. Guard it the way the remote-put test already was.

The heap-membership assertion was one-sided, so every heap pointer cleared
it. Bound both ends.

Remove tl.multiple_of/tl.max_contiguous from the kernel. multiple_of
asserts the addresses are BLOCK_SIZE-divisible, which is false for element
i at base + i*itemsize, and a false assertion is a miscompile rather than
a lost optimization. Measured on gfx942: at BLOCK_SIZE 256 the hinted and
unhinted forms generate identical code, so it bought nothing here anyway.
It earns its place only above one element per lane, which is a question
for the P0 kernels, not for a correctness test.

Refs #546
---
 tests/unittests/test_allocate_symmetric.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/tests/unittests/test_allocate_symmetric.py b/tests/unittests/test_allocate_symmetric.py
index a9f342858..73ea8adce 100644
--- a/tests/unittests/test_allocate_symmetric.py
+++ b/tests/unittests/test_allocate_symmetric.py
@@ -46,12 +46,8 @@ def _put_translated_kernel(
     remote_base_byte = tl.cast(remote_base, tl.pointer_type(tl.int8))
     remote_dst = tl.cast(remote_base_byte + offset, dst.dtype)
 
-    # Hint goes on the indexed pointers; remote_dst is still scalar here.
-    remote_ptrs = remote_dst + offsets
-    remote_ptrs = tl.max_contiguous(tl.multiple_of(remote_ptrs, BLOCK_SIZE), BLOCK_SIZE)
-
     values = tl.load(src + offsets, mask=mask)
-    tl.store(remote_ptrs, values, mask=mask)
+    tl.store(remote_dst + offsets, values, mask=mask)
 
 
 def test_allocate_symmetric_returns_peer_bases():
@@ -70,8 +66,10 @@ def test_allocate_symmetric_returns_peer_bases():
         assert peer_bases.is_cuda
 
         # peer_bases[cur_rank] is what translation subtracts, so the tensor has
-        # to sit inside it.
-        assert tensor.data_ptr() >= int(peer_bases[ctx.get_rank()].item())
+        # to sit inside the heap it points at.
+        local_base = int(peer_bases[ctx.get_rank()].item())
+        assert local_base <= tensor.data_ptr()
+        assert tensor.data_ptr() + tensor.nbytes <= local_base + ctx.heap_size
     finally:
         ctx.barrier()
         del ctx
@@ -118,6 +116,11 @@ def test_view_translates_against_allocation_root():
         # The view's region received the peer's data; everything before it did not.
         torch.testing.assert_close(dst[offset:], torch.full_like(dst[offset:], source_rank + 1))
         torch.testing.assert_close(dst[:offset], torch.full_like(dst[:offset], -1.0))
+
+        if world_size > 1:
+            # At one rank the sender is the receiver, so the assertions above
+            # hold even if translation never left this rank.
+            assert dst[offset].item() != rank + 1
     finally:
         ctx.barrier()
         del ctx