[experimental] rocSHMEM as an allocation provider

[nirvedhmeshram]

Summary (human)

Uses rocSHMEM API to do allocations and form map suggested in #546 , no changes to iris kernels are needed.

Motivation

Issue #546 proposes an allocator-agnostic boundary so Iris device kernels can operate on tensors from providers Iris does not own. This is a second implementation of that shape, against rocSHMEM, to test whether it holds.

The main result is stronger than expected: Iris device code needs no changes at all. iris.store/load/copy already take heap_bases as a plain pointer argument, and __translate computes offset = ptr - bases[from]; bases[to] + offset. Any table satisfying peer_bases[local_rank] == local allocation base drives them. The test kernel calls unmodified iris.store on memory allocated entirely by rocSHMEM, with no Iris context and no Iris heap anywhere in the process.

So the entire integration surface is host-side, and what a provider owes Iris is exactly one int64 table.

Technical details

iris/experimental/rocshmem_provider.py adds:

• RocshmemProvider.allocate_symmetric(*size, dtype=None) -> (tensor, peer_bases) — the same signature and return shape as Add allocate_symmetric() for allocator-agnostic kernels #549, so the same device kernels drive either provider.
• RocshmemProvider.allocate_symmetric_map(...) -> (tensor, SymmetricAddressMap) — the richer descriptor from [Feature]: Formalize allocator-agnostic symmetric tensor address translation #546, carrying local_rank, allocation_base, allocation_bytes, and a per-peer direct mask.
• symmetric_address_map(tensor) for describing an already-allocated tensor.

The table is built from rocshmem_ptr(base, peer), which is OpenSHMEM's shmem_ptr: an address in our own address space for the peer's counterpart, or NULL when that peer is not reachable by direct load/store.

One context-wide table serves every allocation. rocshmem_ptr is a single linear translation of the whole symmetric heap — GDAHostContext::shmem_ptr computes ipc_bases[peer] + (p - ipc_bases[me]) — so the peer delta is constant for every heap address regardless of allocation. Any symmetric anchor yields a table valid for all of them, which also means the provider never needs rocSHMEM's heap base (not exposed publicly). This is verified rather than assumed: a table built from one allocation is used to translate pointers belonging to another, and the data lands correctly.

That property matters for Iris specifically, because iris.copy takes a single heap_bases and translates two pointers against it. A provider handing out genuinely per-allocation tables could not drive it.

Relationship to #549

Compatible, not dependent. This targets main and was validated at 6432c101 with #549 not applied — it needs nothing from that PR, and touches no file it touches. allocate_symmetric mirrors its signature so the two line up when it lands, and will follow whatever shape it settles on.

Test plan

tests/unittests/test_rocshmem_provider.py, in the repo's pytest-under-torchrun convention:

python tests/run_tests_distributed.py tests/unittests/test_rocshmem_provider.py --num_ranks 2

Skips when rocshmem4py is absent, when fewer than 2 ranks are present, or when peers are not directly addressable. tests/manual_rocshmem_provider.py covers the multi-node case the unit test skips.

CI does not install rocSHMEM, so these skip there: the Test unittests jobs report collected 4 items / 4 skipped at 1, 2, 4 and 8 ranks and pass. The skip is applied in a fixture rather than at module scope on purpose — a module-level importorskip collects zero items, and pytest then returns exit code 5 (NO_TESTS_COLLECTED), which run_tests_distributed.py propagates and torchrun reports as a child failure, failing the whole job.

Run against real hardware with rocSHMEM installed (2+ ranks, one node, rocSHMEM built with USE_IPC=ON) the four tests execute and pass.

Note iris/experimental/__init__.py does not import the provider module, so import iris does not require rocshmem4py.

Test results

On MI355X (gfx950, ROCm 7.14, rocSHMEM 3.7.0 GDA/IONIC+IPC):

Run	Result
test_rocshmem_provider.py, 2 ranks under the repo launcher	4 passed on both ranks, no skips
manual, 2 ranks 1 node	PASS — unmodified iris.store over rocSHMEM memory
manual, cross-allocation	PASS — a table from one allocation translates another's pointers
manual, 4 ranks across 2 nodes	PASS — non-addressable peers detected on all 4 ranks

Not covered

• Inter-node peers are detected and refused, not driven. Intra-node (IPC) only.
• Only iris.store is exercised. load, put, get and the atomics are single-translation and should behave identically, but are untested here.
• No views, no mixing two providers in one kernel, no imported tensors — three cases [Feature]: Formalize allocator-agnostic symmetric tensor address translation #546 calls out that this does not reach.

[Copilot]

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Pull request overview

Adds an experimental rocSHMEM-backed “allocation provider” so existing Iris Triton device kernels can operate on rocSHMEM-allocated symmetric memory by supplying a compatible peer_bases table / address map.

Changes:

• Introduces iris/experimental/rocshmem_provider.py with RocshmemProvider, allocate_symmetric(*), and SymmetricAddressMap.
• Adds distributed pytest coverage that drives unmodified iris.store over rocSHMEM memory (skipping cleanly when unavailable).
• Adds a manual torchrun script to exercise intra-node (IPC) and multi-node indirect-peer detection.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 10 comments.

File	Description
iris/experimental/rocshmem_provider.py	Implements rocSHMEM allocations and produces Iris-compatible peer-base tables / richer address descriptor.
tests/unittests/test_rocshmem_provider.py	Adds distributed unit tests validating translation tables and iris.store over rocSHMEM memory.
tests/manual_rocshmem_provider.py	Adds a manual launcher script to validate IPC and indirect-peer reporting outside CI.

---

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.