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.ascend import libaclshmem_device as shmem

Source: python/triton_dist/language/extra/ascend/libaclshmem_device.py

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

Topology and Team Information

my_pe

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 = shmem.my_pe()

n_pes

n_pes() → int32

Get the total number of PEs in the system.

Returns:: int32 — Total number of processing elements

Example:

total_pes = shmem.n_pes()

team_my_pe

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_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

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

Translate a PE number from one team’s context to another team’s context.

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

Remote Memory Access (RMA)

remote_ptr

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

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:

shmem.int_p(remote_flag, tl.int32(1), target_pe)

getmem

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

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

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

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.

Signaling Operations

putmem_signal

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 ACLSHMEMSignalOp enum)
pe – Target PE number

Signal Operations:

Use constants from triton_dist.language.extra.ascend.aclshmem_constants.ACLSHMEMSignalOp:

SET — Atomically set signal location to signal value
ADD — Atomically add signal value to existing value

Example:

from triton_dist.language.extra.ascend.aclshmem_constants import ACLSHMEMSignalOp

shmem.putmem_signal(remote_buf, local_buf, nbytes, signal_addr,
                    tl.int32(1), ACLSHMEMSignalOp.SET, target_pe)

putmem_signal_nbi

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 (SET or ADD)
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

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 (SET or ADD)
pe – Target PE number

signal_wait_until

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

Wait until a signal location satisfies a comparison condition.

Parameters:

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

Returns:

int32 — The observed value that satisfied the condition

Comparison Operations:

Use constants from triton_dist.language.extra.ascend.aclshmem_constants.ACLSHMEMCmpOp:

EQ — Equal to
NE — Not equal to
GT — Greater than
GE — Greater than or equal to
LT — Less than
LE — Less than or equal to

Example:

from triton_dist.language.extra.ascend.aclshmem_constants import ACLSHMEMCmpOp

# Wait until signal >= 1
value = shmem.signal_wait_until(signal_ptr, ACLSHMEMCmpOp.GE, tl.int32(1))

Synchronization

barrier_all

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

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

barrier(team: int32) → void

Team-specific barrier synchronization.

Parameters:: team – Team handle

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

barrier_vec

barrier_vec(team: int32) → void

Vectorized team barrier (optimized implementation).

Parameters:: team – Team handle

quiet

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

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

ACLSHMEMSignalOp

Signal operation types for signaling operations:

class ACLSHMEMSignalOp

SET = 0: Set signal to specified value

ADD = 1: Add to signal value

ACLSHMEMCmpOp

Comparison operations for signal_wait_until:

class ACLSHMEMCmpOp

EQ = 0: Equal to

NE = 1: Not equal to

GT = 2: Greater than

GE = 3: Greater than or equal to

LT = 4: Less than

LE = 5: Less than or equal to

ACLSHMEMTeam

Team identifiers:

class ACLSHMEMTeam

INVALID = -1: Invalid team

WORLD = 0: World team (all PEs)

Usage Example

import triton
import triton.language as tl
from triton_dist.language.extra.ascend import libaclshmem_device as shmem
from triton_dist.language.extra.ascend.aclshmem_constants import ACLSHMEMSignalOp, ACLSHMEMCmpOp

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

    # Synchronize all PEs at start
    shmem.barrier_all()

    # Each PE processes local data
    local_data = tl.load(data_ptr + pe * BLOCK_SIZE)
    processed = local_data * 2

    # Send result to PE 0
    if pe > 0:
        remote_ptr = shmem.remote_ptr(result_ptr, tl.int32(0))
        shmem.putmem_signal(
            remote_ptr + pe * BLOCK_SIZE,
            tl.pointer(processed),
            BLOCK_SIZE * 4,  # 4 bytes per float32
            signal_ptr,
            tl.int32(1),
            ACLSHMEMSignalOp.ADD,
            tl.int32(0)
        )
    else:
        # PE 0 waits for all others
        shmem.signal_wait_until(signal_ptr, ACLSHMEMCmpOp.EQ, tl.int32(n_pe - 1))

    # Final barrier
    shmem.barrier_all()

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