[experimental] rocSHMEM as an allocation provider - #550
Conversation
Lets Iris device code operate on tensors allocated by rocSHMEM instead of from
Iris's own symmetric heap. No Iris device code changes are required: store, load
and copy take heap_bases as a plain pointer argument, so any table satisfying
peer_bases[local_rank] == local allocation base drives them.
iris/experimental/rocshmem_provider.py builds that table from
rocshmem_ptr(base, peer), which returns an address in this process's own space
for a peer's counterpart of a symmetric object, or NULL when that peer is not
reachable by direct load/store.
allocate_symmetric(*size, dtype) -> (tensor, peer_bases)
allocate_symmetric_map(*size, dtype) -> (tensor, SymmetricAddressMap)
symmetric_address_map(tensor) -> SymmetricAddressMap
The first matches Iris.allocate_symmetric's shape so the same kernels drive
either provider. The descriptor form adds local_rank, allocation_base,
allocation_bytes and a per-peer `direct` mask; callers check that mask before
launching, since a peer that is not directly addressable has a base of 0 and
would translate to a wild pointer rather than an error.
One table serves every allocation. rocSHMEM's peer mapping is a linear
translation of the whole symmetric heap, so any symmetric address anchors a
table valid for all allocations, and rocSHMEM's heap base -- which it does not
expose publicly -- is never needed. That also keeps iris.copy usable, since it
translates two pointers against a single heap_bases.
Scope is intra-node. Inter-node peers are reported as unreachable rather than
driven; they need a transport this module does not provide.
Tests:
tests/unittests/test_rocshmem_provider.py pytest under the repo launcher,
skipping when rocshmem4py is absent, when fewer than 2 ranks are present,
or when peers are not directly addressable
tests/manual_rocshmem_provider.py multi-node script, including the
non-addressable-peer path via EXPECT_INDIRECT=1
The provider module is not imported by iris/experimental/__init__.py, so
`import iris` does not require rocshmem4py.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
5abb6e6 to
a28ffe6
Compare
There was a problem hiding this comment.
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.pywithRocshmemProvider,allocate_symmetric(*), andSymmetricAddressMap. - Adds distributed pytest coverage that drives unmodified
iris.storeover rocSHMEM memory (skipping cleanly when unavailable). - Adds a manual
torchrunscript 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.
allocate_symmetric() cached the table built from the first allocation and returned it for every later one. Translation still worked, because the per-peer offset is constant across the heap, but peer_bases[local_rank] was the first allocation's base rather than the current tensor's -- contradicting the stated invariant and breaking the assertion in tests/manual_rocshmem_provider.py, which checks it on a second allocation. Cache the per-peer offsets instead and build each allocation's table from its own base. peer_bases[local_rank] is now that allocation's base for every allocation, while the shared offsets keep a table from one allocation able to translate another's pointers, which iris.copy relies on. This also removes a duplicated rocshmem_ptr sweep: the first allocation previously queried every peer twice, once to build the map and once to seed the cache. It is now queried once per process. peer_bases is created on tensor.device rather than a device captured when the provider was constructed, so the table cannot end up on a different device than the memory it describes. test_table_is_context_wide asserted the two tables were equal, which no longer holds and was the weaker property anyway. It is now test_peer_offsets_are_shared and checks what actually matters: each table's local entry is its own allocation's base, and the per-peer offsets agree. Unreachable peers are excluded from that comparison, since their entry is 0 rather than base + offset. Verified on 2 ranks: pytest 4 passed, manual test PASS including the cross-allocation check. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
allocate_symmetric() now returns a table whose entry r is the address of this tensor on rank r, so peer_bases[cur_rank] == tensor.data_ptr(). It previously returned the heap base table, which named the heap rather than the allocation. This matches what the rocSHMEM provider in #550 returns, so the two line up on the same integers and not merely the same signature. Device-side translation is unchanged either way -- it subtracts peer_bases[cur_rank] and adds peer_bases[to], and any consistent anchor works. Computed on device: heap_bases[cur_rank] stays a tensor rather than going through .item(), so there is no device-to-host sync on the allocation path. The test's heap-membership assertion becomes an equality against data_ptr(), which is the invariant the design actually rests on; the old one-sided bound was cleared by any heap pointer.
|
|
||
| torchrun --nproc_per_node=2 tests/manual_rocshmem_provider.py | ||
|
|
||
| Set EXPECT_INDIRECT=1 and run across 2 nodes to check that peers which are not |
There was a problem hiding this comment.
I don't follow the difference between manual_rocshmem_provider.py and test_rocshmem_provider.py. Do we need both?
There was a problem hiding this comment.
yes the manual one is a multi-node test that is not picked up by the CI, I think its useful when locally testing in that setup. Let me know if you would rather not have it though.
| provider.free(data) | ||
| provider.free(results) | ||
| dist.destroy_process_group() | ||
| return 0 |
There was a problem hiding this comment.
Did this test run in CI or do we need to update the environment with the rocSHMEM dependency?
There was a problem hiding this comment.
In CI env we dont have rocshmem4py so it skips, I have tested it locally.
There was a problem hiding this comment.
I wasnt sure if its ok to add the dependency in the env, if thats ok, I can add it and make sure the test passes in CI
There was a problem hiding this comment.
Let's add it! If you have trouble updating the CI files, just let me know what you need installed.
Actually, a Readme next to the provider file describing how to install the dependancies would be great and we can copy that into the CI scripts.
There was a problem hiding this comment.
if you edit this file, you should be able to get it in the CI container env:
https://github.com/ROCm/iris/blob/main/apptainer/iris.def#L35
There was a problem hiding this comment.
Added iris/experimental/README.md, covering install, verification and the run-time contract.
Good news for the CI side: the bindings are a one-line pip install, which I tested in a fresh venv with nothing else set —
CMAKE_PREFIX_PATH=<rocshmem-install> pip install \
"rocshmem4py @ git+https://github.com/ROCm/rocm-systems.git#subdirectory=python/rocshmem"The part that isn't free is what it builds against: there's no prebuilt wheel anywhere — not PyPI, not the ROCm nightly wheel indexes — and the bindings don't build rocSHMEM themselves (find_package(rocshmem 3.5.0 CONFIG REQUIRED), no FetchContent), so the image needs a rocSHMEM install first: a cmake/ninja build for the target arch with USE_IPC=ON, which this provider requires since rocshmem_ptr returns NULL unconditionally without it.
So: one pip line, plus a rocSHMEM build. I've gone ahead and made that change rather than just describing it — .github/scripts/install_rocshmem.sh, called from both docker/Dockerfile and apptainer/iris.def so there's one implementation rather than two copies. It builds for gfx942 to match the MI325X runners, IPC only, which is all a single-node runner can exercise anyway.
One difference from the command above: the script installs from the checkout it already made rather than from the git+ URL. Since it has to build the core from that repo anyway, taking both from one checkout keeps the core and the bindings at the same revision — a git+ URL would have pip clone independently at whatever develop is at by then, and find_package wouldn't catch the skew because it only compares versions while the bindings statically link the core.
Two incidental changes it needed. The Docker build context was docker/, so the Dockerfile couldn't COPY anything from .github/ — it now builds from the repo root with -f, which also matches how the Apptainer %files paths resolve, and a .dockerignore keeps the wider context from shipping .git.
One caveat worth knowing, noted in the def file: container_build.sh caches the Apptainer image on the checksum of iris.def alone, so editing install_rocshmem.sh on its own will silently reuse a stale image — touch the def file too.
CI is still churning through the first rebuild (every image recompiles rocSHMEM from source), so I can't claim the provider tests are green yet — only that they should now run instead of skipping. I'll follow up here once it reports.
There was a problem hiding this comment.
Could you update the def file I linked too please? We actually are running Apptainer these days because of runner/docker issue.
There was a problem hiding this comment.
Sorry, ignore that. Looks like you already did and image was built.
Review follow-ups, all documentation: - allocate_symmetric: say it is collective, describe the returned table's shape, dtype, device and invariant, and point at allocate_symmetric_map for the direct mask. - symmetric_address_map: note it builds a fresh table per call, why that is once per allocation rather than per launch, and why it is deliberately not memoised by data_ptr. - free: explain it cannot be automated. rocshmem_free is documented as collective and must be called by all PEs, so a __del__ or weakref finalizer would let ranks diverge on GC timing and hang. - module: record that rocshmem4py is a standalone package that statically links rocSHMEM rather than linking it at run time, so its USE_IPC setting is fixed at its build time.
The provider tests have only ever skipped in CI, because rocshmem4py was not installed. Add it, so they actually execute. Installing it is a source build but a short one. There is no prebuilt rocshmem4py wheel on any index -- not PyPI, not the ROCm nightly indexes -- and the bindings do not build rocSHMEM themselves (find_package(rocshmem 3.5.0 CONFIG REQUIRED), no FetchContent), so rocSHMEM is built first and the bindings are pip-installed against it with CMAKE_PREFIX_PATH. Scope is IPC only, which is what the provider uses and all a single-node runner can exercise. Upstream already defaults USE_IPC=ON and USE_GDA=OFF, so no conduit flags are passed, which keeps MPI and the RDMA provider libraries out of it. Built for gfx942 to match the MI325X runners. One installer shared by both image definitions rather than two copies. That required widening the Docker build context from docker/ to the repo root so the Dockerfile can COPY it, hence the .dockerignore; the Apptainer def pulls the same file in via %files. Also adds iris/experimental/README.md covering install, verification and the collective-call contract, as requested in review.
Both are pip installs from source, but the README documented the git+ form while the CI script used a local path, with nothing explaining the difference. One checkout serves the core build and the bindings, so they are guaranteed to be the same revision. A git+ URL would have pip clone the monorepo again at whatever develop is at by then, and find_package would not catch the skew because it only compares versions while the bindings statically link the core.
The docstring claimed the same signature and return shape as Iris.allocate_symmetric. Iris has no such method -- its allocation API is zeros/ones/full/uniform/as_symmetric, with a context-wide table from get_heap_bases(). The (tensor, peer_bases) shape comes from the provider interface being proposed, not from an existing method.
Comment only. The CI bases are ROCm 7.2.1 (apptainer) and 7.1 (docker), both older than the 7.14 artifacts that ship rocSHMEM's static library and headers, so the source build is still required. Records where to cut it when a base image bumps, and that rocshmem4py stays a source build either way until its TheRock packaging lands.
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/copyalready takeheap_basesas a plain pointer argument, and__translatecomputesoffset = ptr - bases[from]; bases[to] + offset. Any table satisfyingpeer_bases[local_rank] == local allocation basedrives them. The test kernel calls unmodifiediris.storeon 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
int64table.Technical details
iris/experimental/rocshmem_provider.pyadds:RocshmemProvider.allocate_symmetric(*size, dtype=None) -> (tensor, peer_bases)— the same signature and return shape as Addallocate_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, carryinglocal_rank,allocation_base,allocation_bytes, and a per-peerdirectmask.symmetric_address_map(tensor)for describing an already-allocated tensor.The table is built from
rocshmem_ptr(base, peer), which is OpenSHMEM'sshmem_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_ptris a single linear translation of the whole symmetric heap —GDAHostContext::shmem_ptrcomputesipc_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.copytakes a singleheap_basesand 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
mainand was validated at6432c101with #549 not applied — it needs nothing from that PR, and touches no file it touches.allocate_symmetricmirrors 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:Skips when rocshmem4py is absent, when fewer than 2 ranks are present, or when peers are not directly addressable.
tests/manual_rocshmem_provider.pycovers the multi-node case the unit test skips.CI does not install rocSHMEM, so these skip there: the
Test unittestsjobs reportcollected 4 items/4 skippedat 1, 2, 4 and 8 ranks and pass. The skip is applied in a fixture rather than at module scope on purpose — a module-levelimportorskipcollects zero items, and pytest then returns exit code 5 (NO_TESTS_COLLECTED), whichrun_tests_distributed.pypropagates 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__.pydoes not import the provider module, soimport irisdoes not require rocshmem4py.Test results
On MI355X (gfx950, ROCm 7.14, rocSHMEM 3.7.0 GDA/IONIC+IPC):
test_rocshmem_provider.py, 2 ranks under the repo launcheriris.storeover rocSHMEM memoryNot covered
iris.storeis exercised.load,put,getand the atomics are single-translation and should behave identically, but are untested here.