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: .. code-block:: bash pip install aclshmem All host APIs are available through: .. code-block:: python import shmem **Source:** ``3rdparty/shmem/src/python/shmem/__init__.py`` For comprehensive ACLSHMEM documentation, refer to the `ACLSHMEM docs <../../../3rdparty/shmem/docs/index.rst>`_. Initialization and Finalization -------------------------------- aclshmem_init ^^^^^^^^^^^^^ .. py:function:: aclshmem_init(attributes: InitAttr) -> InitStatus Initialize the ACLSHMEM runtime with specified attributes. :param 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: .. code-block:: python import shmem attr = shmem.InitAttr() attr.mem_size = 1024 * 1024 * 1024 # 1GB symmetric heap status = shmem.aclshmem_init(attr) aclshmem_finalize ^^^^^^^^^^^^^^^^^ .. py:function:: 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: .. code-block:: python # At program end shmem.aclshmem_finalize() aclshmem_init_using_unique_id ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: 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. :param rank: Rank number of this process (0 to nranks-1) :param nranks: Total number of processes/ranks :param mem_size: Size of symmetric heap in bytes :param uid: Unique ID bytes obtained from aclshmem_get_unique_id() :returns: InitStatus — Initialization status Usage Pattern: .. code-block:: python 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 ^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: 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 ^^^^^^^^^^^^^^^ .. py:function:: aclshmem_malloc(nbytes: int) -> int Allocate symmetric memory from the SHMEM heap. :param 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: .. code-block:: python # Allocate 4MB of symmetric memory ptr = shmem.aclshmem_malloc(4 * 1024 * 1024) if ptr == 0: raise RuntimeError("Symmetric memory allocation failed") aclshmem_free ^^^^^^^^^^^^^ .. py:function:: aclshmem_free(ptr: int) -> None Free symmetric memory previously allocated with ``aclshmem_malloc``. :param 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 ^^^^^^^^^^^^^^^ .. py:function:: aclshmem_calloc(count: int, size: int) -> int Allocate and zero-initialize symmetric memory. :param count: Number of elements :param 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 ^^^^^^^^^^^^^^ .. py:function:: aclshmem_align(alignment: int, size: int) -> int Allocate aligned symmetric memory. :param alignment: Alignment requirement in bytes (must be power of 2) :param size: Number of bytes to allocate :returns: int — Pointer to aligned memory, or 0 on failure Example: .. code-block:: python # Allocate 1MB aligned to 4KB boundary ptr = shmem.aclshmem_align(4096, 1024 * 1024) aclshmem_ptr ^^^^^^^^^^^^ .. py:function:: aclshmem_ptr(local_ptr: int, pe: int) -> int Get a pointer to a symmetric object on a remote PE (host-accessible). :param local_ptr: Local pointer to symmetric memory :param 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 ^^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: 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 ^^^^^ .. py:function:: my_pe() -> int :no-index: Get the PE number of the calling process. :returns: int — PE number (0 to pe_count()-1) pe_count ^^^^^^^^ .. py:function:: pe_count() -> int Get the total number of PEs. :returns: int — Total number of PEs in the system Host RMA Operations ------------------- aclshmem_putmem ^^^^^^^^^^^^^^^ .. py:function:: aclshmem_putmem(dest: int, source: int, bytes: int, pe: int) -> None Host-side blocking put (write) to remote symmetric memory. :param dest: Destination pointer on remote PE :param source: Source pointer on local host :param bytes: Number of bytes to transfer :param pe: Target PE number Example: .. code-block:: python 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 ^^^^^^^^^^^^^^^ .. py:function:: aclshmem_getmem(dest: int, source: int, bytes: int, pe: int) -> None Host-side blocking get (read) from remote symmetric memory. :param dest: Destination pointer on local host :param source: Source pointer on remote PE :param bytes: Number of bytes to transfer :param pe: Source PE number aclshmem_putmem_nbi ^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_putmem_nbi(dest: int, source: int, bytes: int, pe: int) -> None Host-side non-blocking put to remote symmetric memory. :param dest: Destination pointer on remote PE :param source: Source pointer on local host :param bytes: Number of bytes to transfer :param pe: Target PE number Operation may complete asynchronously. Use appropriate synchronization before reusing source buffer or assuming remote visibility. aclshmem_getmem_nbi ^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_getmem_nbi(dest: int, source: int, bytes: int, pe: int) -> None Host-side non-blocking get from remote symmetric memory. :param dest: Destination pointer on local host :param source: Source pointer on remote PE :param bytes: Number of bytes to transfer :param pe: Source PE number aclshmemx_putmem_signal ^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: 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. :param dest: Destination pointer on remote PE :param source: Source pointer on local host :param nbytes: Number of bytes to transfer :param sig_addr: Signal address on remote PE :param signal: Signal value :param sig_op: Signal operation (SET or ADD) :param pe: Target PE number aclshmemx_putmem_signal_nbi ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: 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. :param dest: Destination pointer on remote PE :param source: Source pointer on local host :param nbytes: Number of bytes to transfer :param sig_addr: Signal address on remote PE :param signal: Signal value :param sig_op: Signal operation (SET or ADD) :param pe: Target PE number aclshmem_signal_wait_until ^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: 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. :param sig_addr: Pointer to signal location :param cmp: Comparison operation (see device API for constants) :param cmp_val: Comparison value Team Management --------------- team_split_strided ^^^^^^^^^^^^^^^^^^ .. py:function:: 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. :param start: Starting PE in parent team :param stride: Stride between selected PEs :param size: Number of PEs in new team :param parent_team: Parent team handle :returns: Team handle for the newly created team Example: .. code-block:: python # 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 ^^^^^^^^^^^^^ .. py:function:: team_split_2d(xdim: int, xaxis_teams, ydim: int, yaxis_teams) -> None Split PEs into 2D grid of teams. :param xdim: Size of X dimension :param xaxis_teams: Output array for X-axis teams :param ydim: Size of Y dimension :param yaxis_teams: Output array for Y-axis teams team_translate_pe ^^^^^^^^^^^^^^^^^ .. py:function:: team_translate_pe(src_team, src_pe: int, dest_team) -> int :no-index: Translate a PE number from one team to another. :param src_team: Source team handle :param src_pe: PE number in source team :param dest_team: Destination team handle :returns: int — Corresponding PE number in destination team team_my_pe ^^^^^^^^^^ .. py:function:: team_my_pe(team) -> int :no-index: Get the calling PE's number within the specified team. :param team: Team handle :returns: int — PE number within the team team_n_pes ^^^^^^^^^^ .. py:function:: team_n_pes(team) -> int :no-index: Get the number of PEs in the specified team. :param team: Team handle :returns: int — Number of PEs in the team team_destroy ^^^^^^^^^^^^ .. py:function:: team_destroy(team) -> None Destroy a team and release its resources. :param team: Team handle to destroy Do not destroy predefined teams. All PEs in the team must call this collectively. Configuration and Info ---------------------- InitAttr ^^^^^^^^ .. py:class:: InitAttr Initialization attributes class for configuring ACLSHMEM runtime. **Attributes:** .. py:attribute:: mem_size Size of symmetric heap in bytes InitStatus ^^^^^^^^^^ .. py:class:: InitStatus Status object returned by initialization functions. OpEngineType ^^^^^^^^^^^^ .. py:class:: OpEngineType Enum for operation engine types. aclshmem_info_get_version ^^^^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_info_get_version() -> str Get the ACLSHMEM library version string. :returns: str — Version string aclshmem_info_get_name ^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_info_get_name() -> str Get the ACLSHMEM library name. :returns: str — Library name set_log_level ^^^^^^^^^^^^^ .. py:function:: set_log_level(level: int) -> None Configure ACLSHMEM logging level. :param level: Logging level (use standard Python logging levels) Utility Functions ----------------- aclshmem_global_exit ^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_global_exit(status: int) -> None Perform a global exit of all PEs with the specified status code. :param status: Exit status code Terminates all PEs in the SHMEM job. Use for coordinated error handling. aclshmem_create_tensor ^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_create_tensor(shape: tuple, dtype: torch.dtype = torch.float32, device_id: int = 0) -> torch.Tensor Create a PyTorch tensor backed by symmetric memory. :param shape: Tensor shape tuple :param dtype: PyTorch data type (default: torch.float32) :param device_id: NPU device ID (default: 0) :returns: torch.Tensor — Tensor backed by symmetric memory Example: .. code-block:: python 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 ^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_free_tensor(tensor: torch.Tensor) -> None Free the symmetric memory backing a tensor created with ``aclshmem_create_tensor``. :param tensor: Tensor to free Only use with tensors created via ``aclshmem_create_tensor``. Tensor should not be used after calling this function. Complete Example ---------------- .. code-block:: python 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: .. code-block:: bash pip install aclshmem Or build from source in the ``3rdparty/shmem`` directory. See the :doc:`../getting-started/build` 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.