SHMEM Host API

This module provides host-side ACLSHMEM operations for initializing the SHMEM runtime, managing symmetric memory, and performing host-initiated communication. These APIs are called from Python host code, not from within kernels.

The ACLSHMEM package must be installed separately. Install with:

pip install aclshmem

All host APIs are available through:

import shmem

Source: 3rdparty/shmem/src/python/shmem/__init__.py

For comprehensive ACLSHMEM documentation, refer to the ACLSHMEM docs.

Initialization and Finalization

aclshmem_init

aclshmem_init(attributes: InitAttr) → InitStatus

Initialize the ACLSHMEM runtime with specified attributes.

Parameters:: attributes – InitAttr object containing initialization configuration
Returns:: InitStatus — Status object containing initialization results

Must be called before any other SHMEM operations. Each PE must call this function.

Example:

import shmem

attr = shmem.InitAttr()
attr.mem_size = 1024 * 1024 * 1024  # 1GB symmetric heap
status = shmem.aclshmem_init(attr)

aclshmem_finalize

aclshmem_finalize() → None

Finalize and clean up the ACLSHMEM runtime.

Must be called by all PEs before program termination. Releases all symmetric memory and communication resources. After calling this, no other SHMEM operations can be performed.

Example:

# At program end
shmem.aclshmem_finalize()

aclshmem_init_using_unique_id

aclshmem_init_using_unique_id(rank: int, nranks: int, mem_size: int, uid: bytes) → InitStatus

Initialize ACLSHMEM using a unique ID for multi-process coordination.

Parameters:

rank – Rank number of this process (0 to nranks-1)
nranks – Total number of processes/ranks
mem_size – Size of symmetric heap in bytes
uid – Unique ID bytes obtained from aclshmem_get_unique_id()

Returns:

InitStatus — Initialization status

Usage Pattern:

import shmem

# On rank 0: generate and broadcast unique ID
if rank == 0:
    uid = shmem.aclshmem_get_unique_id()
    # ... broadcast uid to all ranks ...
else:
    # ... receive uid from rank 0 ...
    pass

# All ranks initialize with the same uid
status = shmem.aclshmem_init_using_unique_id(
    rank, nranks, mem_size=1024*1024*1024, uid=uid
)

aclshmem_get_unique_id

aclshmem_get_unique_id() → bytes

Generate a unique ID for coordinating multi-process SHMEM initialization.

Returns:: bytes — Unique identifier to be shared across all processes

Typically called by rank 0 and broadcast to other ranks. The same unique ID must be used by all participating processes.

Memory Management

aclshmem_malloc

aclshmem_malloc(nbytes: int) → int

Allocate symmetric memory from the SHMEM heap.

Parameters:: nbytes – Number of bytes to allocate
Returns:: int — Pointer (as integer) to the allocated memory, or 0 on failure

Memory is allocated at the same offset in the symmetric heap on all PEs. All PEs must call this with the same nbytes value. Use this for memory that will be accessed remotely via RMA operations.

Example:

# Allocate 4MB of symmetric memory
ptr = shmem.aclshmem_malloc(4 * 1024 * 1024)
if ptr == 0:
    raise RuntimeError("Symmetric memory allocation failed")

aclshmem_free

aclshmem_free(ptr: int) → None

Free symmetric memory previously allocated with aclshmem_malloc.

Parameters:: ptr – Pointer (as integer) to the memory to free

All PEs must free symmetric allocations collectively. Do not free memory that was not allocated with aclshmem_malloc.

aclshmem_calloc

aclshmem_calloc(count: int, size: int) → int

Allocate and zero-initialize symmetric memory.

Parameters:

count – Number of elements
size – Size of each element in bytes

Returns:

int — Pointer to allocated memory, or 0 on failure

Equivalent to aclshmem_malloc(count * size) followed by zeroing. All PEs must call with the same parameters.

aclshmem_align

aclshmem_align(alignment: int, size: int) → int

Allocate aligned symmetric memory.

Parameters:

alignment – Alignment requirement in bytes (must be power of 2)
size – Number of bytes to allocate

Returns:

int — Pointer to aligned memory, or 0 on failure

Example:

# Allocate 1MB aligned to 4KB boundary
ptr = shmem.aclshmem_align(4096, 1024 * 1024)

aclshmem_ptr

aclshmem_ptr(local_ptr: int, pe: int) → int

Get a pointer to a symmetric object on a remote PE (host-accessible).

Parameters:

local_ptr – Local pointer to symmetric memory
pe – Target PE number

Returns:

int — Pointer that can be used on the host to access the symmetric object on the specified PE, or 0 if not accessible

May return 0 for inter-node remote pointers if direct host access is not supported. Primarily useful for intra-node scenarios.

aclshmemx_get_heap_base

aclshmemx_get_heap_base() → int

Get the base address of the symmetric heap.

Returns:: int — Base address of the symmetric heap

PE Information

my_pe

my_pe() → int

Get the PE number of the calling process.

Returns:: int — PE number (0 to pe_count()-1)

pe_count

pe_count() → int

Get the total number of PEs.

Returns:: int — Total number of PEs in the system

Host RMA Operations

aclshmem_putmem

aclshmem_putmem(dest: int, source: int, bytes: int, pe: int) → None

Host-side blocking put (write) to remote symmetric memory.

Parameters:

dest – Destination pointer on remote PE
source – Source pointer on local host
bytes – Number of bytes to transfer
pe – Target PE number

Example:

import numpy as np
import shmem

# Allocate symmetric memory
remote_buf = shmem.aclshmem_malloc(1024)

# Create host data
data = np.arange(256, dtype=np.float32)

# Put data to PE 1
shmem.aclshmem_putmem(remote_buf, data.ctypes.data, data.nbytes, pe=1)

aclshmem_getmem

aclshmem_getmem(dest: int, source: int, bytes: int, pe: int) → None

Host-side blocking get (read) from remote symmetric memory.

Parameters:

dest – Destination pointer on local host
source – Source pointer on remote PE
bytes – Number of bytes to transfer
pe – Source PE number

aclshmem_putmem_nbi

aclshmem_putmem_nbi(dest: int, source: int, bytes: int, pe: int) → None

Host-side non-blocking put to remote symmetric memory.

Parameters:

dest – Destination pointer on remote PE
source – Source pointer on local host
bytes – Number of bytes to transfer
pe – Target PE number

Operation may complete asynchronously. Use appropriate synchronization before reusing source buffer or assuming remote visibility.

aclshmem_getmem_nbi

aclshmem_getmem_nbi(dest: int, source: int, bytes: int, pe: int) → None

Host-side non-blocking get from remote symmetric memory.

Parameters:

dest – Destination pointer on local host
source – Source pointer on remote PE
bytes – Number of bytes to transfer
pe – Source PE number

aclshmemx_putmem_signal

aclshmemx_putmem_signal(dest: int, source: int, nbytes: int, sig_addr: int, signal: int, sig_op: int, pe: int) → None

Host-side blocking put with atomic signal operation on completion.

Parameters:

dest – Destination pointer on remote PE
source – Source pointer on local host
nbytes – Number of bytes to transfer
sig_addr – Signal address on remote PE
signal – Signal value
sig_op – Signal operation (SET or ADD)
pe – Target PE number

aclshmemx_putmem_signal_nbi

aclshmemx_putmem_signal_nbi(dest: int, source: int, nbytes: int, sig_addr: int, signal: int, sig_op: int, pe: int) → None

Host-side non-blocking put with signal.

Parameters:

dest – Destination pointer on remote PE
source – Source pointer on local host
nbytes – Number of bytes to transfer
sig_addr – Signal address on remote PE
signal – Signal value
sig_op – Signal operation (SET or ADD)
pe – Target PE number

aclshmem_signal_wait_until

aclshmem_signal_wait_until(sig_addr: int, cmp: int, cmp_val: int) → None

Host-side wait on a signal location until a condition is met.

Parameters:

sig_addr – Pointer to signal location
cmp – Comparison operation (see device API for constants)
cmp_val – Comparison value

Team Management

team_split_strided

team_split_strided(start: int, stride: int, size: int, parent_team) → team

Create a new team by selecting PEs with a strided pattern from the parent team.

Parameters:

start – Starting PE in parent team
stride – Stride between selected PEs
size – Number of PEs in new team
parent_team – Parent team handle

Returns:

Team handle for the newly created team

Example:

# Create team with PEs 0, 2, 4, 6 from global team
team = shmem.team_split_strided(
    start=0, stride=2, size=4, parent_team=shmem.SHMEM_TEAM_WORLD
)

team_split_2d

team_split_2d(xdim: int, xaxis_teams, ydim: int, yaxis_teams) → None

Split PEs into 2D grid of teams.

Parameters:

xdim – Size of X dimension
xaxis_teams – Output array for X-axis teams
ydim – Size of Y dimension
yaxis_teams – Output array for Y-axis teams

team_translate_pe

team_translate_pe(src_team, src_pe: int, dest_team) → int

Translate a PE number from one team to another.

Parameters:

src_team – Source team handle
src_pe – PE number in source team
dest_team – Destination team handle

Returns:

int — Corresponding PE number in destination team

team_my_pe

team_my_pe(team) → int

Get the calling PE’s number within the specified team.

Parameters:: team – Team handle
Returns:: int — PE number within the team

team_n_pes

team_n_pes(team) → int

Get the number of PEs in the specified team.

Parameters:: team – Team handle
Returns:: int — Number of PEs in the team

team_destroy

team_destroy(team) → None

Destroy a team and release its resources.

Parameters:: team – Team handle to destroy

Do not destroy predefined teams. All PEs in the team must call this collectively.

Configuration and Info

InitAttr

class InitAttr

Initialization attributes class for configuring ACLSHMEM runtime.

Attributes:

mem_size: Size of symmetric heap in bytes

InitStatus

class InitStatus: Status object returned by initialization functions.

OpEngineType

class OpEngineType: Enum for operation engine types.

aclshmem_info_get_version

aclshmem_info_get_version() → str

Get the ACLSHMEM library version string.

Returns:: str — Version string

aclshmem_info_get_name

aclshmem_info_get_name() → str

Get the ACLSHMEM library name.

Returns:: str — Library name

set_log_level

set_log_level(level: int) → None

Configure ACLSHMEM logging level.

Parameters:: level – Logging level (use standard Python logging levels)

Utility Functions

aclshmem_global_exit

aclshmem_global_exit(status: int) → None

Perform a global exit of all PEs with the specified status code.

Parameters:: status – Exit status code

Terminates all PEs in the SHMEM job. Use for coordinated error handling.

aclshmem_create_tensor

aclshmem_create_tensor(shape: tuple, dtype: torch.dtype = torch.float32, device_id: int = 0) → torch.Tensor

Create a PyTorch tensor backed by symmetric memory.

Parameters:

shape – Tensor shape tuple
dtype – PyTorch data type (default: torch.float32)
device_id – NPU device ID (default: 0)

Returns:

torch.Tensor — Tensor backed by symmetric memory

Example:

import torch
import shmem

# Create symmetric tensor
tensor = shmem.aclshmem_create_tensor(
    (1024, 1024), dtype=torch.float32, device_id=0
)

# Use like any other PyTorch tensor
tensor.fill_(0.0)

aclshmem_free_tensor

aclshmem_free_tensor(tensor: torch.Tensor) → None

Free the symmetric memory backing a tensor created with aclshmem_create_tensor.

Parameters:: tensor – Tensor to free

Only use with tensors created via aclshmem_create_tensor. Tensor should not be used after calling this function.

Complete Example

import shmem
import torch
import numpy as np

# Initialize SHMEM
attr = shmem.InitAttr()
attr.mem_size = 1024 * 1024 * 1024  # 1GB
status = shmem.aclshmem_init(attr)

# Get PE info
my_rank = shmem.my_pe()
num_ranks = shmem.pe_count()

print(f"PE {my_rank} of {num_ranks}")

# Allocate symmetric memory
sym_ptr = shmem.aclshmem_malloc(1024 * 1024)  # 1MB
if sym_ptr == 0:
    raise RuntimeError("Failed to allocate symmetric memory")

# Or create symmetric tensor
sym_tensor = shmem.aclshmem_create_tensor((1024, 256), dtype=torch.float32)

# Perform communication
if my_rank == 0:
    # Rank 0 sends data to rank 1
    data = np.arange(256, dtype=np.float32)
    if num_ranks > 1:
        shmem.aclshmem_putmem(sym_ptr, data.ctypes.data, data.nbytes, pe=1)
elif my_rank == 1:
    # Rank 1 receives (data is already in symmetric memory)
    pass

# Clean up
shmem.aclshmem_free_tensor(sym_tensor)
shmem.aclshmem_free(sym_ptr)
shmem.aclshmem_finalize()

Installation Notes

The ACLSHMEM package is distributed separately and must be installed:

pip install aclshmem

Or build from source in the 3rdparty/shmem directory. See the Build from Source guide for details.

Notes

Symmetric Memory Model

SHMEM uses a symmetric memory model where allocations occur at the same virtual address offset on all PEs. This enables efficient remote memory access without explicit address translation.

Collective Operations

Many SHMEM operations are collective and must be called by all PEs in the team or globally. Examples include:

aclshmem_init / aclshmem_finalize
aclshmem_malloc / aclshmem_free
Barrier operations
Team creation and destruction

Thread Safety

SHMEM operations are generally not thread-safe. If using multiple threads per PE, appropriate synchronization must be added by the application.