Add allocate_symmetric() for allocator-agnostic kernels - #549
Open
mawad-amd wants to merge 5 commits into
Open
Conversation
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.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
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
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
github-actions
Bot
added
in-progress
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
nirvedhmeshram
left a comment
nirvedhmeshram
left a comment
There was a problem hiding this comment.
The API looks good to me, I am not certain its strictly necessary for rocSHMEM as a consumer, iris.store/load/copy already take heap_bases as a plain pointer, and __translate just does offset = ptr - bases[from]; bases[to] + offset, so any table satisfying the invariant drives them. I got rocSHMEM-allocated tensors working through unmodified Iris device code, with no Iris context and no Iris heap anywhere in the process. One table covers every allocation, too, since rocshmem_ptr is a linear translation over the whole symmetric heap., here is a branch I experimented with main...nirvedhmeshram:iris:nmeshram/rocshmem-provider#diff-bd62caa6315bd8129f290bb8e246aa5cfa4e6173ba7b5a5bb814935a0c233ab5R63-R83
nirvedhmeshram
added a commit
to nirvedhmeshram/iris
that referenced
this pull request
Sep 2, 2026
Independent check of the alignment finding in mawad-amd's variant analysis: Triton specializes pointer ARGUMENTS on 16-byte divisibility, so an address computed from a value loaded out of memory loses that information and the store narrows. Confirmed on gfx950 / Triton 3.8.0. Three kernels differing only in where the destination address comes from: direct (pointer argument) buffer_store_dwordx4 x1 unhinted (PR ROCm#549's kernel verbatim) buffer_store_dword x4 hinted (+ tl.multiple_of(.,16)) buffer_store_dwordx4 x1 The control is the point: without it, "four narrow stores" could just mean this shape never vectorizes -- masked store, block size, dtype -- rather than anything about pointer provenance. All three keep an identical mask so the only variable is the alignment information available to the compiler. Two things the run rules out. The destination was 16B-aligned, so this is not real misalignment, only unprovable alignment. And the compiler did not fold the identity translation (peer_bases[0] == dst.data_ptr()), which would have made the comparison meaningless. Practical consequence: iris.store/load already accept a `hint` that applies tl.multiple_of/max_contiguous, and all_gather passes one. Kernels that inline the translation instead -- as PR ROCm#549's example does -- get the narrow pattern unless they re-assert it themselves. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Iris device kernels are already independent of the host allocator -- given a
local pointer and a peer-base table, they can translate to a peer's address.
The host API does not reflect that: it exposes a context-wide heap, which
assumes every symmetric tensor comes from the Iris heap.
This adds the allocation entry point that other providers (rocSHMEM, Torch
Symmetric Memory) can expose with the same shape, so one set of device kernels
serves all of them. First step toward #546.
Technical Details
Adds
Iris.allocate_symmetric(*size, dtype=None), returning the tensor and theint64device-resident peer-base table indexed by rank.Kernels take the pair as two ordinary arguments -- a pointer and a tensor --
and inline the address translation. Two arguments rather than one struct: the
named-tuple flattening needed for the struct form requires a recent Triton, and
several environments are pinned to older ROCm/Triton stacks.
Purely additive.
get_heap_bases()and all existing RMA APIs are unchanged.Test Plan
New
tests/unittests/test_allocate_symmetric.py, with a kernel that calls noIris translation helper.
Test Result
Submission Checklist