TheRock/rocm-systems/projects/rocshmem/python

rocshmem4py: Python Bindings for rocSHMEM

rocshmem4py provides Python bindings for the ROCm OpenSHMEM (rocSHMEM) runtime library, enabling GPU-centric networking through an OpenSHMEM-like interface on AMD ROCm platforms.

Features

• Core rocSHMEM API: Memory management, data transfer, atomics, synchronization
• Teams: Team split/destroy/translate with tracked lifecycle and WORLD/INVALID sentinels
• Team-scoped collectives: Stream-ordered all-to-all and broadcast
• Framework-agnostic allocation: rocshmem_create_buffer / rocshmem_get_peer_buffer work without any ML framework
• Heap-base introspection: rocshmem_get_heap_bases for device-side pointer translation
• __cuda_array_interface__: Zero-copy interop between SymmetricBuffer and PyTorch tensors
• PyTorch interop submodule: rocshmem4py.interop.torch for tensor allocation, RMA, and collectives
• PyTorch coordination: init_with_torch() / finalize_with_torch() for torch.distributed-based init
• MPI Integration: init_with_mpi() / finalize_with_mpi() for existing MPI applications

API Coverage

rocshmem4py currently exposes a focused host-side subset of rocSHMEM APIs: initialization/finalization, PE/team queries, team management, memory allocation, put/get (blocking/non-blocking/stream variants), team-scoped collectives, the full host atomic-operation matrix, and key constants. Context APIs (rocshmem_ctx_*) are not yet exposed. It does not yet expose every rocSHMEM host API.

Host-facing symbols are exported directly from the compiled extension module (_rocshmem4py) to avoid drift between Python imports and C++ bindings.

Prerequisites

• AMD ROCm 6.0+ with HIP
• rocSHMEM library (built with RO, IPC, or GDA backend)
• Python 3.8+
• CMake 3.20+
• nanobind 2.12.0+ (binding backend)
• Distributed launcher: torchrun (for init_with_torch(), IPC/GDA backends) or mpirun (for init_with_mpi(), and required by the RO backend regardless of init path — see RO backend requires mpirun below).
• ROCm-aware Open MPI 5.0.x + UCX 1.17+ — RO backend only; any distro Open MPI works for IPC/GDA bootstrap. See ROCm-aware Open MPI required for RO below.
• mpi4py (only when using init_with_mpi()).

Installation

export ROCM_PATH=/opt/rocm
export ROCSHMEM_HOME=/path/to/rocshmem/build
export LD_LIBRARY_PATH=$ROCSHMEM_HOME/lib:$ROCM_PATH/lib:$LD_LIBRARY_PATH

pip install -e .

The build requirements (nanobind, cmake) are declared in pyproject.toml, so a standard pip install -e . pulls them into the build isolation environment automatically. If you build with --no-build-isolation, install them into your environment first: pip install nanobind cmake.

Note: this package is not published to PyPI yet; install from source only.

Binding backend

rocshmem4py uses nanobind to build the compiled extension module (_rocshmem4py). The public Python contract — module name, function names, argument behavior, and return types — is stable and independent of the binding framework. Framework-independent rocSHMEM glue lives in src/rocshmem4py_common.hpp; the thin m.def(...) registration layer lives in src/rocshmem4py.cc.

Quick Start

Backend launch matrix

Backend	Requires MPI runtime	Recommended launcher	Example
RO	Yes	mpirun	mpirun -n 2 python my_script.py
IPC/GDA	No	torchrun	torchrun --standalone --nproc_per_node=2 my_script.py

init_with_torch() works under both launchers. init_with_mpi() requires an mpirun-launched process environment.

With PyTorch coordination (recommended)

import torch
import rocshmem4py
from rocshmem4py.interop import torch as rshmem_torch

# RO backend:   mpirun -n 2 python my_script.py
# IPC/GDA:      torchrun --standalone --nproc_per_node=2 my_script.py
rocshmem4py.init_with_torch()

my_pe = rocshmem4py.rocshmem_my_pe()
n_pes = rocshmem4py.rocshmem_n_pes()

# Allocate a symmetric tensor (backed by rocshmem_malloc)
src = rshmem_torch.create_tensor((64,), torch.float32)
dst = rshmem_torch.create_tensor((64,), torch.float32)
src.fill_(float(my_pe))
dst.fill_(-1.0)
torch.cuda.synchronize()

# Stream-ordered transfer to the next PE (portable across all backends)
peer = (my_pe + 1) % n_pes
rshmem_torch.barrier_all()
rshmem_torch.put(dst, src, peer)
rshmem_torch.barrier_all()
torch.cuda.synchronize()

rocshmem4py.finalize_with_torch()

With MPI coordination

from mpi4py import MPI
import rocshmem4py

rocshmem4py.init_with_mpi(MPI.COMM_WORLD)

my_pe = rocshmem4py.rocshmem_my_pe()
buf = rocshmem4py.SymmetricBuffer(1024)
rocshmem4py.rocshmem_barrier_all()

buf.free()
rocshmem4py.finalize_with_mpi()

API Reference

Initialization / Finalization

Function	Description
init_with_torch(group=None)	Init rocSHMEM via torch.distributed (recommended)
finalize_with_torch()	Synchronized teardown of rocSHMEM + torch.distributed
init_with_mpi(comm)	Init rocSHMEM via mpi4py
finalize_with_mpi()	Synchronized teardown for init_with_mpi() sessions
init_rocshmem_by_uniqueid(group)	Low-level init with a torch process group
set_hip_device_from_env()	Pin HIP device from LOCAL_RANK before raw init
rocshmem_init()	Raw rocSHMEM init (rarely needed directly)
rocshmem_finalize()	Raw rocSHMEM finalize
rocshmem_init_attr(rank, nranks, uid)	Init with unique ID
rocshmem_get_uniqueid()	Get a unique ID for init_attr

PE Queries

Function	Description
rocshmem_my_pe()	PE number of the calling process
rocshmem_n_pes()	Total number of PEs
rocshmem_team_my_pe(team)	PE number within a team
rocshmem_team_n_pes(team)	Number of PEs in a team

Teams

Team handles are passed as Python int values (raw intptr_t from the C library). Use the sentinel constants below for the special teams; all other handles come from rocshmem_team_split_strided.

Function / type	Description
TeamConfig()	Configuration record for rocshmem_team_split_strided (num_contexts)
rocshmem_team_split_strided(parent, start, stride, size, config=None, mask=0)	Split a parent team into a strided sub-team; returns (status, team_handle). Non-members receive ROCSHMEM_TEAM_INVALID (-1), not 0
rocshmem_team_destroy(team)	Destroy a team (no-op for ROCSHMEM_TEAM_WORLD / ROCSHMEM_TEAM_INVALID)
rocshmem_team_translate_pe(src_team, src_pe, dest_team)	Map a PE index between teams; returns -1 if unmappable

Pass ROCSHMEM_TEAM_WORLD (0) for world-scope team operations. The C binding's resolve_team_handle() translates this sentinel to the runtime handle; use the constant rather than a raw runtime pointer so rocshmem_team_destroy and the tracked split/destroy wrappers stay safe.

finalize_with_torch() and finalize_with_mpi() automatically destroy any teams created via the tracked rocshmem_team_split_strided wrapper before calling rocshmem_finalize().

Memory Management

Framework-agnostic (rocshmem4py)

Function	Description
rocshmem_malloc(size)	Allocate symmetric memory (returns raw pointer)
rocshmem_free(ptr)	Free symmetric memory
rocshmem_ptr(dest, pe)	Get remote symmetric pointer (IPC backends only)
SymmetricBuffer(size)	RAII wrapper; exposes __cuda_array_interface__ for zero-copy torch interop
rocshmem_create_buffer(nbytes)	Collective allocation returning a SymmetricBuffer
rocshmem_get_peer_buffer(buf, peer)	Non-owning SymmetricBuffer view of a peer's buffer (IPC only)
rocshmem_get_heap_bases(ptr)	Per-PE base addresses for a symmetric allocation (for device-side pointer translation)

PyTorch interop (rocshmem4py.interop.torch)

Function	Description
create_tensor(shape, dtype)	Collective allocation returning a symmetric torch.Tensor
get_peer_tensor(tensor, peer)	Zero-copy tensor view of a peer's symmetric tensor (IPC only)
free_tensor(tensor)	Explicit collective-safe deallocation
put(dst, src, peer, stream=None)	Stream-ordered put (all backends)
get(dst, src, peer, stream=None)	Stream-ordered get (all backends)
barrier_all(stream=None)	Stream-ordered collective barrier
get_heap_bases(tensor)	(n_pes,) int64 GPU tensor of per-PE heap bases for a symmetric tensor
alltoall(team, dst, src, stream=None)	Stream-ordered all-to-all over a team
broadcast(team, dst, src, pe_root, stream=None)	Stream-ordered broadcast over a team

Collectives

Host (rocshmem4py)

Function	Description
rocshmem_alltoallmem_on_stream(team, dest, source, bytes_per_pe, stream)	Stream-ordered all-to-all; bytes_per_pe is bytes sent to each PE in the team
rocshmem_broadcastmem_on_stream(team, dest, source, nbytes, pe_root, stream)	Stream-ordered broadcast; nbytes is total bytes; pe_root is in the team's PE space
rocshmem_sync_all()	Lighter-weight sync (local-store visibility)
rocshmem_sync_all_on_stream(stream)	Stream-ordered sync_all

Pass ROCSHMEM_TEAM_WORLD (0) as the team handle for world-scope collectives.

Data Transfer

Function	Description
rocshmem_putmem(dest, src, nbytes, pe)	Blocking put
rocshmem_getmem(dest, src, nbytes, pe)	Blocking get
rocshmem_putmem_nbi(dest, src, nbytes, pe)	Non-blocking put
rocshmem_getmem_nbi(dest, src, nbytes, pe)	Non-blocking get
rocshmem_putmem_on_stream(dest, src, nbytes, pe, stream)	Stream-ordered put
rocshmem_getmem_on_stream(dest, src, nbytes, pe, stream)	Stream-ordered get
rocshmem_putmem_signal_on_stream(...)	Stream-ordered put with signal
rocshmem_signal_wait_until_on_stream(...)	Stream-ordered signal wait

Synchronization

Function	Description
rocshmem_barrier_all()	Barrier across all PEs
rocshmem_barrier_all_on_stream(stream)	Stream-ordered barrier across all PEs
rocshmem_barrier(team)	Blocking barrier across all PEs in team
rocshmem_barrier_on_stream(team, stream)	Stream-ordered team barrier; ROCSHMEM_TEAM_INVALID is a no-op
rocshmem_team_sync(team)	Team member rendezvous plus local-store visibility (lighter than barrier)
rocshmem_team_sync_on_stream(team, stream)	Stream-ordered team sync; ROCSHMEM_TEAM_INVALID is a no-op
rocshmem_sync_all()	World-scope sync (local-store visibility)
rocshmem_sync_all_on_stream(stream)	Stream-ordered world sync
rocshmem_fence()	Ordering fence
rocshmem_quiet()	Wait for all outstanding operations
hip_device_synchronize()	Synchronize the current HIP device

Team-scoped calls use the same team handle conventions as collectives (ROCSHMEM_TEAM_WORLD or a handle from rocshmem_team_split_strided). Non-members of a child team should pass ROCSHMEM_TEAM_INVALID (-1); the runtime returns without enqueueing work. If RMA was issued on a HIP stream, synchronize that stream before a blocking rocshmem_barrier(team); use rocshmem_barrier_on_stream on the same stream as the RMA for overlap.

Atomic Operations

Host atomic operations are auto-exported from the compiled extension by naming convention (rocshmem_{type}_atomic_{op}). The matrix covers integer, unsigned, and floating types with fetch/set/swap/CAS, fetch-add/add, fetch-inc/inc, and bitwise ops where the rocSHMEM header provides them.

Examples:

Function	Description
rocshmem_int_atomic_fetch_add(dest, value, pe)	Atomic int fetch-and-add
rocshmem_long_atomic_fetch_add(dest, value, pe)	Atomic long fetch-and-add
rocshmem_int_atomic_compare_swap(dest, cond, value, pe)	Atomic int CAS
rocshmem_uint64_atomic_or(dest, value, pe)	Atomic uint64 OR

Note: Host AMO bindings are present in all builds, but behavioral correctness on the IPC no-MPI backend depends on runtime support that is still being extended.

Constants

Constant	Value	Description
ROCSHMEM_TEAM_WORLD	0	Team containing all PEs
ROCSHMEM_TEAM_INVALID	-1	Invalid team identifier
ROCSHMEM_SUCCESS	0	Success status code
ROCSHMEM_SIGNAL_SET	impl-defined	Signal set op enum
ROCSHMEM_SIGNAL_ADD	impl-defined	Signal add op enum
ROCSHMEM_CMP_EQ/NE/GT/GE/LT/LE	impl-defined	Signal wait compare enums

Running the tests

# IPC / GDA backend (any AMD multi-GPU node)
torchrun --standalone --nproc_per_node=2 -m pytest tests/ -v

# RO backend (must use mpirun)
./launch_test.sh -n 2 -c "pytest tests/ -v"

Test source layout, the full launcher × backend × init-tier matrix used by conftest.py, and CI-author guidance live in the test suite's own README: https://github.com/ROCm/rocm-systems/blob/develop/projects/rocshmem/python/tests/README.md.

Troubleshooting

Symptom	Cause	Fix
ImportError: _rocshmem4py	rocSHMEM not on the loader path	export LD_LIBRARY_PATH=$ROCSHMEM_HOME/lib:$ROCM_PATH/lib:$LD_LIBRARY_PATH
CMake cannot find rocSHMEM at build time	ROCSHMEM_HOME unset	export ROCSHMEM_HOME=/path/to/rocshmem/build before pip install -e .
Link error: recompile with -fPIC	librocshmem.a built without PIC	Rebuild rocSHMEM with -DCMAKE_POSITION_INDEPENDENT_CODE=ON
RO backend hangs in mca_btl_vader.so (Wrote -1, errno = 14), aborts with MPI_ERR_WIN: invalid window, or hangs at exit in __run_exit_handlers → libamdhip64	Non-ROCm-aware Open MPI / UCX cannot handle GPU buffers in the data plane	Use ROCm-aware Open MPI 5.0.x + UCX 1.17+ — see below
Unsupported configuration to initialize rocSHMEM. Please initialize the MPI library using MPI_Init first	RO backend launched under torchrun — see "RO backend requires mpirun" below	Launch with mpirun (you can keep init_with_torch()), or build an IPC/GDA-only rocSHMEM if you don't need inter-node RDMA
Rendezvous / port conflict under torchrun	Default MASTER_PORT=29500 already taken	Set MASTER_PORT (or ROCSHMEM_MASTER_PORT) to a free port

RO backend requires mpirun

This is a structural property of rocSHMEM's C library, not a packaging gap in rocshmem4py: library_init_subcomm (in src/rocshmem.cpp) requires either MPI_Initialized() to be true or the OpenMPI launcher env var OMPI_COMM_WORLD_SIZE to be set. mpirun/prterun exports those env vars; torchrun does not. Pre-importing mpi4py is not a workaround — it makes MPI_Initialized() return true, which then routes the C library through a subgroup-creation path (MPI_Group_incl + MPI_Comm_create_group) across processes that live in disjoint singleton MPI universes, and that crashes inside OMPI's PMIx wireup. Use mpirun for RO; init_with_torch() itself still works under mpirun, so you keep the torch.distributed unique-id exchange and only the launcher changes.

ROCm-aware Open MPI required for RO

Only RO backend builds need ROCm-aware Open MPI 5.0.x + UCX 1.17+ built with --with-rocm for data-plane RMA. IPC and GDA backends use MPI only for bootstrap and work with any distro Open MPI (see prereqs and troubleshooting above).

Verify on the launch host:

mpirun --version           # must say "Open MPI v5.0.x"
ucx_info -d | grep rocm    # must list rocm_copy and rocm_ipc transports

See the rocSHMEM install docs for build instructions.

Diagnosing which rocSHMEM backend you actually linked

A pre-built _rocshmem4py.so (or any wheel you might pick up later) statically links one rocSHMEM backend. When initialization fails in ways that look backend-specific, two checks pin down which one:

# What backend does the rocSHMEM install at $ROCSHMEM_HOME advertise?
"${ROCSHMEM_HOME}/bin/rocshmem_info" | grep "Vendor String"

# Which backend's code is actually linked into the loaded extension?
nm -D --defined-only "$(python -c 'import _rocshmem4py; print(_rocshmem4py.__file__)')" \
  | grep -E " T _ZN8rocshmem(9RO|10IPC|10GDA)Backend" | head

The auto-detect order is IPC → GDA → RO; force a specific one with ROCSHMEM_BACKEND=ipc if the build supports it. The same _rocshmem4py source builds against any backend — there are no backend #ifdefs in the binding layer — so two wheels with identical Python APIs can behave differently at runtime depending on how the C library was configured.

License

MIT License. See LICENSE for details.