SHMEM Device API¶

This module provides device-side ACLSHMEM operations that can be called from within Triton kernels running on Ascend NPUs. These APIs enable fine-grained control over remote memory access, synchronization, and inter-PE communication.

All device APIs are available through:

from triton_dist.language.extra import libshmem_device

For more information about the underlying ACLSHMEM library, refer to the ACLSHMEM documentation.

Topology and Team Information¶

my_pe() → int32¶

Get the processing element (PE) ID of the current PE.

Returns:: int32 — Current PE number (0 to n_pes - 1)

Example:

pe_id = libshmem_device.my_pe()

n_pes() → int32¶

Get the total number of PEs in the system.

Returns:: int32 — Total number of processing elements

Example:

total_pes = libshmem_device.n_pes()

team_my_pe(team: int32) → int32¶

Get the PE number within a specific team.

Parameters:: team – Team handle (int32)
Returns:: int32 — PE number within the specified team

team_n_pes(team: int32) → int32¶

Get the number of PEs in a specific team.

Parameters:: team – Team handle (int32)
Returns:: int32 — Number of PEs in the team

team_translate_pe(src_team: int32, pe_in_src_team: int32, dest_team: int32) → int32¶

Translate a PE number from one team’s local numbering to another team’s local numbering. The function identifies the global PE represented by pe_in_src_team in src_team, then returns that PE’s local number in dest_team.

Parameters:

src_team – Source team handle
pe_in_src_team – PE number in the source team
dest_team – Destination team handle

Returns:

int32 — Corresponding PE number in the destination team

Example:

world_team = libshmem_device.ACLSHMEM_TEAM_WORLD
team_pe = libshmem_device.team_my_pe(world_team)
team_size = libshmem_device.team_n_pes(world_team)
world_pe = libshmem_device.team_translate_pe(
    world_team, team_pe, libshmem_device.ACLSHMEM_TEAM_WORLD)

team_pe_mapping(team: int32, pe: int32) → int32¶

Convert a PE number local to team into its global PE number in ACLSHMEM_TEAM_WORLD. For valid inputs, this is equivalent to team_translate_pe(team, pe, libshmem_device.ACLSHMEM_TEAM_WORLD). For example, in a team containing global PEs [1, 3, 5], local PE 1 maps to global PE 3. For the world team, the PE number is unchanged.

Parameters:

team – Team handle available on the calling PE
pe – Nonnegative PE number within the specified team

Returns:

int32 — Global PE number, or -1 if team is ACLSHMEM_TEAM_INVALID or pe is at least the team size

Both arguments are cast to int32. Use a valid team handle and a PE number in the range 0 <= pe < team_n_pes(team). Negative PE numbers and other invalid team handles are not checked by the underlying operation.

Example (team is a valid team handle passed to the kernel):

# team is the valid team handle passed to the kernel.
team_pe = libshmem_device.team_my_pe(team)
# global_pe equals libshmem_device.my_pe()
global_pe = libshmem_device.team_pe_mapping(team, team_pe)

Remote Memory Access (RMA)¶

remote_ptr(local_ptr, pe) → pointer¶

Get a pointer to a symmetric memory object on a remote PE.

Parameters:

local_ptr – Pointer to local symmetric memory
pe – Target PE number (int32 or uint32)

Returns:

Pointer of the same type as local_ptr, addressing the symmetric object on the remote PE

Works with all supported data types. The local pointer must reference symmetric memory.

int_p(dest, value: int32, pe: int32) → void¶

Put a single int32 value to a remote PE (blocking operation).

Parameters:

dest – Pointer to destination on remote PE (must be int32*)
value – int32 value to write
pe – Target PE number

Example:

libshmem_device.int_p(remote_flag, 1, target_pe)

getmem(dest, source, bytes: uint32, pe: int32) → void¶

Blocking get (read) operation from remote memory.

Parameters:

dest – Pointer to local destination buffer
source – Pointer to source on remote PE
bytes – Number of bytes to transfer (cast to uint32)
pe – Source PE number

Supported Types: fp16, fp32, int8, int16, int32, int64, uint8, uint16, uint32, uint64, bf16

putmem(dest, source, bytes: uint32, pe: int32) → void¶

Blocking put (write) operation to remote memory.

Parameters:

dest – Pointer to destination on remote PE
source – Pointer to local source buffer
bytes – Number of bytes to transfer (cast to uint32)
pe – Destination PE number

Supported Types: fp16, fp32, int8, int16, int32, int64, uint8, uint16, uint32, uint64, bf16

getmem_nbi(dest, source, bytes: uint32, pe: int32, qp_id: int = 0) → void¶

Non-blocking get (read) operation from remote memory.

Parameters:

dest – Pointer to local destination buffer
source – Pointer to source on remote PE
bytes – Number of bytes to transfer
pe – Source PE number
qp_id – Queue pair ID (default: 0, currently unused)

Operation may complete asynchronously. Use quiet() or fence() to ensure completion before accessing result.

putmem_nbi(dest, source, bytes: uint32, pe: int32, qp_id: int = 0) → void¶

Non-blocking put (write) operation to remote memory.

Parameters:

dest – Pointer to destination on remote PE
source – Pointer to local source buffer
bytes – Number of bytes to transfer
pe – Destination PE number
qp_id – Queue pair ID (default: 0, currently unused)

Operation may complete asynchronously. Use quiet() or fence() to ensure completion.

Scalar Get and Atomics¶

g(source, pe: int32)¶: Read one scalar from symmetric memory on pe. The result has the source pointer’s element type. Supports the RMA data types listed below.

The scalar atomic operations take a symmetric destination pointer and a remote PE. Values and comparison operands are cast to the destination pointer’s element type; PE arguments are cast to int32.

API	Result
`atomic_op(dest, value, pe, op)`	Old value for `op=shmem.ACLSHMEM_AMO_SWAP`; void otherwise
`atomic_fetch_op(dest, value, pe, op)`	Old value in the destination element type
`atomic_compare_swap(dest, cond, value, pe)`	Old value; replacement occurs only when it equals `cond`

op must be a compile-time ACLSHMEMAtomicOp enum member, exposed as shmem.ACLSHMEM_AMO_*. Plain integers, booleans, strings and members of other enum classes are rejected. atomic_op accepts shmem.ACLSHMEM_AMO_ADD, shmem.ACLSHMEM_AMO_AND, shmem.ACLSHMEM_AMO_OR, shmem.ACLSHMEM_AMO_XOR, shmem.ACLSHMEM_AMO_INC, shmem.ACLSHMEM_AMO_SET and shmem.ACLSHMEM_AMO_SWAP. atomic_fetch_op accepts shmem.ACLSHMEM_AMO_ADD, shmem.ACLSHMEM_AMO_AND, shmem.ACLSHMEM_AMO_OR, shmem.ACLSHMEM_AMO_XOR, shmem.ACLSHMEM_AMO_INC and shmem.ACLSHMEM_AMO_NONE. Pass None as value for shmem.ACLSHMEM_AMO_INC and shmem.ACLSHMEM_AMO_NONE; no placeholder operand is emitted. shmem.ACLSHMEM_AMO_NONE performs an atomic read without updating the destination. Runtime selectors and unsupported selections are rejected before emitting an external call.

All selections support int32, uint32, int64 and uint64. atomic_op with op=shmem.ACLSHMEM_AMO_ADD additionally supports int8, int16, bf16, fp16 and fp32; op=shmem.ACLSHMEM_AMO_SET and op=shmem.ACLSHMEM_AMO_SWAP additionally support fp32. Fetch-add does not inherit add’s small or floating types. These are scalar operations; tensor operands and masks are not supported. Hardware support must match the underlying ACLSHMEM wrapper.

Atomic operations use generic SHMEM runtime routing. For example:

shmem.atomic_op(dest, 1, pe, shmem.ACLSHMEM_AMO_ADD)
old = shmem.atomic_fetch_op(dest, None, pe, shmem.ACLSHMEM_AMO_INC)
current = shmem.atomic_fetch_op(dest, None, pe, shmem.ACLSHMEM_AMO_NONE)
old = shmem.atomic_op(dest, value, pe, shmem.ACLSHMEM_AMO_SWAP)

Fetching operations and swap return synchronously. Void updates require completion before dependent reads or buffer reuse, according to the selected runtime transport’s completion contract. In the current ACLSHMEM runtime, quiet and fence complete operations issued by the calling rank; they do not synchronize ranks. In a multi-rank scenario, when ranks depend on updates from other ranks, each participating rank must call quiet or fence followed by a matching sync before consuming those updates. sync alone does not complete outstanding remote memory operations.

Signaling Operations¶

putmem_signal(dest, source, nbytes: uint32, sig_addr, signal: int32, sig_op: int32, pe: int32) → void¶

Blocking put operation with atomic signal update on completion.

Parameters:

dest – Pointer to destination on remote PE
source – Pointer to local source buffer
nbytes – Number of bytes to transfer
sig_addr – Pointer to signal location on remote PE (must be int32*)
signal – Signal value to apply
sig_op – Signal operation (see Constants)
pe – Target PE number

Signal Operations:

Use constants from libshmem_device:

ACLSHMEM_SIGNAL_SET — Atomically set signal location to signal value
ACLSHMEM_SIGNAL_ADD — Atomically add signal value to existing value

Example:

libshmem_device.putmem_signal(remote_buf, local_buf, nbytes, signal_addr,
                    1, libshmem_device.ACLSHMEM_SIGNAL_SET, target_pe)

putmem_signal_nbi(dest, source, nbytes: uint32, sig_addr, signal: int32, sig_op: int32, pe: int32) → void¶

Non-blocking put operation with atomic signal update on completion.

Parameters:

dest – Pointer to destination on remote PE
source – Pointer to local source buffer
nbytes – Number of bytes to transfer
sig_addr – Pointer to signal location on remote PE (must be int32*)
signal – Signal value to apply
sig_op – Signal operation (see Constants)
pe – Target PE number

The signal operation occurs when the data transfer completes. Use quiet() to ensure all non-blocking operations have completed.

signal_op(sig_addr, signal: int32, sig_op: int32, pe: int32) → void¶

Perform an atomic signal operation on a remote PE without data transfer.

Parameters:

sig_addr – Pointer to signal location on remote PE (must be int32*)
signal – Signal value to apply
sig_op – Signal operation (see Constants)
pe – Target PE number

signal_wait_until(sig_addr, cmp_: int32, cmp_val: int32) → int32¶

Wait until the value at a signal location satisfies a comparison condition.

Parameters:

sig_addr – Pointer to local signal location (must be int32*)
cmp_ – Comparison operation (see Constants)
cmp_val – Value to compare against

Returns:

int32 — The observed value that satisfied the condition

Comparison Operations:

Use constants from libshmem_device:

ACLSHMEM_CMP_EQ — Equal to
ACLSHMEM_CMP_NE — Not equal to
ACLSHMEM_CMP_GT — Greater than
ACLSHMEM_CMP_GE — Greater than or equal to
ACLSHMEM_CMP_LT — Less than
ACLSHMEM_CMP_LE — Less than or equal to

Example:

# Wait until signal >= 1
value = libshmem_device.signal_wait_until(
    signal_ptr, libshmem_device.ACLSHMEM_CMP_GE, 1)

Synchronization¶

barrier_all() → void¶

Global barrier across all PEs in the system.

All PEs must call this operation. Execution blocks until all PEs reach the barrier.

barrier_all_vec() → void¶

Vectorized global barrier across all PEs (optimized implementation).

Ascend-specific optimized barrier implementation. Semantically equivalent to barrier_all but may have better performance.

barrier(team: int32) → void¶

Team-specific barrier synchronization.

Parameters:: team – Valid team handle, such as libshmem_device.ACLSHMEM_TEAM_WORLD or a subteam created by the host API

All PEs in the team must call this operation. Execution blocks until all team members reach the barrier.

barrier_vec(team: int32) → void¶

Vectorized team barrier (optimized implementation).

Parameters:: team – Valid team handle, such as libshmem_device.ACLSHMEM_TEAM_WORLD or a subteam created by the host API

sync(team: int32) → void¶
sync_all() → void¶
sync_vec(team: int32) → void¶
sync_vec_all() → void¶: Synchronize the PEs in a team or in the world team (*_all). The sync_vec variants involve vector cores only; sync and sync_all involve both cube and vector cores. Every participating core must call the operation. Unlike barriers, these operations do not complete outstanding remote memory operations; use quiet() when needed.

quiet() → void¶

Wait for completion of all outstanding RMA operations initiated by the local PE.

Ensures all preceding non-blocking put and get operations have completed. Does not guarantee remote visibility (use fence() for ordering).

fence() → void¶

Memory fence ensuring ordering of RMA operations.

Ensures all preceding RMA operations are ordered before subsequent RMA operations. Provides memory ordering guarantees without waiting for completion.

Supported Data Types¶

The RMA operations support the following data types:

Triton Type	C/Kernel Suffix
`fp16`	`half`
`fp32`	`float`
`bf16`	`bfloat16`
`int8`	`int8`
`int16`	`int16`
`int32`	`int32`
`int64`	`int64`
`uint8`	`uint8`
`uint16`	`uint16`
`uint32`	`uint32`
`uint64`	`uint64`

Constants¶

Note

The constants in libshmem_device should be used within kernels; using them on the host side may introduce risks.

Signal operation types:

libshmem_device.ACLSHMEM_SIGNAL_SET = 0¶: Set signal to specified value.

libshmem_device.ACLSHMEM_SIGNAL_ADD = 1¶: Add to signal value.

Comparison operations for signal_wait_until:

libshmem_device.ACLSHMEM_CMP_EQ = 0¶: Equal to.

libshmem_device.ACLSHMEM_CMP_NE = 1¶: Not equal to.

libshmem_device.ACLSHMEM_CMP_GT = 2¶: Greater than.

libshmem_device.ACLSHMEM_CMP_GE = 3¶: Greater than or equal to.

libshmem_device.ACLSHMEM_CMP_LT = 4¶: Less than.

libshmem_device.ACLSHMEM_CMP_LE = 5¶: Less than or equal to.

Team identifiers:

libshmem_device.ACLSHMEM_TEAM_INVALID = -1¶: Invalid team.

libshmem_device.ACLSHMEM_TEAM_WORLD = 0¶: World team (all PEs).

Usage Example¶

The following kernel sends one float32 block from every nonzero PE to its corresponding slot in PE 0’s symmetric result buffer. Each sender sets a dedicated signal after its transfer, and PE 0 waits for all sender signals.

import triton
import triton.language as tl
from triton.language.extra.cann.extension import sub_vec_id
from triton_dist.language.extra import libshmem_device


@triton.jit
def distributed_kernel(
    data_ptr,
    signal_ptr,
    result_ptr,
    BLOCK_SIZE: tl.constexpr,
    my_rank: tl.constexpr,
):
   pe = libshmem_device.my_pe()
   n_pe = libshmem_device.n_pes()
   subblock_idx = sub_vec_id()

   # Synchronize all PEs at start
   libshmem_device.barrier_all_vec()

   # Run the transfer or wait logic once per vector core group
   if subblock_idx == 0:
      if my_rank > 0:
         source_ptr = data_ptr + pe * BLOCK_SIZE

         # Send this PE's local data block to its slot in PE 0's symmetric buffer
         libshmem_device.putmem_signal(
            result_ptr + pe * BLOCK_SIZE,
            source_ptr,
            BLOCK_SIZE * 4,  # 4 bytes per float32
            signal_ptr + pe,
            1,
            libshmem_device.ACLSHMEM_SIGNAL_SET,
            0,
         )
      else:
         # PE 0 waits for all others
         for sender in range(1, n_pe):
            libshmem_device.signal_wait_until(
                signal_ptr + sender,
                libshmem_device.ACLSHMEM_CMP_EQ,
                1,
            )

   # Final barrier
   libshmem_device.barrier_all_vec()

Notes¶

Symmetric Memory Requirement¶

Most RMA operations require symmetric memory — memory allocated at the same virtual address offset on all PEs. Use the host-side aclshmem_malloc() to allocate symmetric memory (see SHMEM Host API).

Error Handling¶

Invalid enum values for sig_op or cmp_ will raise ValueError with a detailed message listing valid options.

Performance Considerations¶

Non-blocking operations (*_nbi) can improve performance by overlapping communication with computation
Use barrier_vec and barrier_all_vec for potentially better performance than standard barriers
Batch multiple small transfers into larger ones when possible to amortize communication overhead
Use extern_call infrastructure to invoke these operations from Triton kernels