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: .. code-block:: python 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 <../../../3rdparty/shmem/docs/index.rst>`_. Topology and Team Information ------------------------------ my_pe ^^^^^ .. py:function:: my_pe() -> int32 Get the processing element (PE) ID of the current PE. :returns: int32 — Current PE number (0 to n_pes - 1) Example: .. code-block:: python pe_id = shmem.my_pe() n_pes ^^^^^ .. py:function:: n_pes() -> int32 Get the total number of PEs in the system. :returns: int32 — Total number of processing elements Example: .. code-block:: python total_pes = shmem.n_pes() team_my_pe ^^^^^^^^^^ .. py:function:: team_my_pe(team: int32) -> int32 Get the PE number within a specific team. :param team: Team handle (int32) :returns: int32 — PE number within the specified team team_n_pes ^^^^^^^^^^ .. py:function:: team_n_pes(team: int32) -> int32 Get the number of PEs in a specific team. :param team: Team handle (int32) :returns: int32 — Number of PEs in the team team_translate_pe ^^^^^^^^^^^^^^^^^ .. py:function:: 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. :param src_team: Source team handle :param pe_in_src_team: PE number in the source team :param dest_team: Destination team handle :returns: int32 — Corresponding PE number in the destination team Remote Memory Access (RMA) --------------------------- remote_ptr ^^^^^^^^^^ .. py:function:: remote_ptr(local_ptr, pe) -> pointer Get a pointer to a symmetric memory object on a remote PE. :param local_ptr: Pointer to local symmetric memory :param 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 ^^^^^ .. py:function:: int_p(dest, value: int32, pe: int32) -> void Put a single int32 value to a remote PE (blocking operation). :param dest: Pointer to destination on remote PE (must be int32\*) :param value: int32 value to write :param pe: Target PE number Example: .. code-block:: python shmem.int_p(remote_flag, tl.int32(1), target_pe) getmem ^^^^^^ .. py:function:: getmem(dest, source, bytes: uint32, pe: int32) -> void Blocking get (read) operation from remote memory. :param dest: Pointer to local destination buffer :param source: Pointer to source on remote PE :param bytes: Number of bytes to transfer (cast to uint32) :param pe: Source PE number **Supported Types:** fp16, fp32, int8, int16, int32, int64, uint8, uint16, uint32, uint64, bf16 putmem ^^^^^^ .. py:function:: putmem(dest, source, bytes: uint32, pe: int32) -> void Blocking put (write) operation to remote memory. :param dest: Pointer to destination on remote PE :param source: Pointer to local source buffer :param bytes: Number of bytes to transfer (cast to uint32) :param pe: Destination PE number **Supported Types:** fp16, fp32, int8, int16, int32, int64, uint8, uint16, uint32, uint64, bf16 getmem_nbi ^^^^^^^^^^ .. py:function:: getmem_nbi(dest, source, bytes: uint32, pe: int32, qp_id: int = 0) -> void Non-blocking get (read) operation from remote memory. :param dest: Pointer to local destination buffer :param source: Pointer to source on remote PE :param bytes: Number of bytes to transfer :param pe: Source PE number :param 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 ^^^^^^^^^^ .. py:function:: putmem_nbi(dest, source, bytes: uint32, pe: int32, qp_id: int = 0) -> void Non-blocking put (write) operation to remote memory. :param dest: Pointer to destination on remote PE :param source: Pointer to local source buffer :param bytes: Number of bytes to transfer :param pe: Destination PE number :param qp_id: Queue pair ID (default: 0, currently unused) Operation may complete asynchronously. Use ``quiet()`` or ``fence()`` to ensure completion. Signaling Operations -------------------- putmem_signal ^^^^^^^^^^^^^ .. py:function:: 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. :param dest: Pointer to destination on remote PE :param source: Pointer to local source buffer :param nbytes: Number of bytes to transfer :param sig_addr: Pointer to signal location on remote PE (must be int32\*) :param signal: Signal value to apply :param sig_op: Signal operation (see ACLSHMEMSignalOp enum) :param 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: .. code-block:: python 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 ^^^^^^^^^^^^^^^^^ .. py:function:: 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. :param dest: Pointer to destination on remote PE :param source: Pointer to local source buffer :param nbytes: Number of bytes to transfer :param sig_addr: Pointer to signal location on remote PE (must be int32\*) :param signal: Signal value to apply :param sig_op: Signal operation (SET or ADD) :param 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 ^^^^^^^^^ .. py:function:: signal_op(sig_addr, signal: int32, sig_op: int32, pe: int32) -> void Perform an atomic signal operation on a remote PE without data transfer. :param sig_addr: Pointer to signal location on remote PE (must be int32\*) :param signal: Signal value to apply :param sig_op: Signal operation (SET or ADD) :param pe: Target PE number signal_wait_until ^^^^^^^^^^^^^^^^^ .. py:function:: signal_wait_until(sig_addr, cmp_: int32, cmp_val: int32) -> int32 Wait until a signal location satisfies a comparison condition. :param sig_addr: Pointer to local signal location (must be int32\*) :param cmp_: Comparison operation (see ACLSHMEMCmpOp enum) :param 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: .. code-block:: python 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 ^^^^^^^^^^^ .. py:function:: 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 ^^^^^^^^^^^^^^^ .. py:function:: 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 ^^^^^^^ .. py:function:: barrier(team: int32) -> void Team-specific barrier synchronization. :param team: Team handle All PEs in the team must call this operation. Execution blocks until all team members reach the barrier. barrier_vec ^^^^^^^^^^^ .. py:function:: barrier_vec(team: int32) -> void Vectorized team barrier (optimized implementation). :param team: Team handle quiet ^^^^^ .. py:function:: 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 ^^^^^ .. py:function:: 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: .. list-table:: :header-rows: 1 :widths: 30 70 * - 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: .. py:class:: ACLSHMEMSignalOp .. py:attribute:: SET :value: 0 Set signal to specified value .. py:attribute:: ADD :value: 1 Add to signal value ACLSHMEMCmpOp ^^^^^^^^^^^^^ Comparison operations for ``signal_wait_until``: .. py:class:: ACLSHMEMCmpOp .. py:attribute:: EQ :value: 0 Equal to .. py:attribute:: NE :value: 1 Not equal to .. py:attribute:: GT :value: 2 Greater than .. py:attribute:: GE :value: 3 Greater than or equal to .. py:attribute:: LT :value: 4 Less than .. py:attribute:: LE :value: 5 Less than or equal to ACLSHMEMTeam ^^^^^^^^^^^^ Team identifiers: .. py:class:: ACLSHMEMTeam .. py:attribute:: INVALID :value: -1 Invalid team .. py:attribute:: WORLD :value: 0 World team (all PEs) Usage Example ------------- .. code-block:: python 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 :doc:`shmem_host`). 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