SHMEM 主机端 API ================ 本模块提供主机端 ACLSHMEM 操作，用于初始化 SHMEM 运行时、管理对称内存以及执行主机发起的通信。 这些 API 从 Python 主机代码调用，而非从 kernel 内部调用。 ACLSHMEM 包必须单独安装。使用以下命令安装： .. code-block:: bash pip install aclshmem 所有主机端 API 可通过以下方式访问： .. code-block:: python import shmem **源文件：** ``3rdparty/shmem/src/python/shmem/__init__.py`` 关于完整的 ACLSHMEM 文档，请参考 `ACLSHMEM 文档 <../../../3rdparty/shmem/docs/index.rst>`_。 初始化与终止 ------------ aclshmem_init ^^^^^^^^^^^^^ .. py:function:: aclshmem_init(attributes: InitAttr) -> InitStatus 使用指定属性初始化 ACLSHMEM 运行时。 :param attributes: 包含初始化配置的 InitAttr 对象 :returns: InitStatus —— 包含初始化结果的状态对象 必须在任何其他 SHMEM 操作之前调用。每个 PE 必须调用此函数。 示例： .. code-block:: python import shmem attr = shmem.InitAttr() attr.mem_size = 1024 * 1024 * 1024 # 1GB 对称堆 status = shmem.aclshmem_init(attr) aclshmem_finalize ^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_finalize() -> None 终止并清理 ACLSHMEM 运行时。 必须在程序终止前由所有 PE 调用。释放所有对称内存和通信资源。调用此函数后， 无法执行其他 SHMEM 操作。 示例： .. code-block:: python # 程序结束时 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 使用唯一 ID 初始化 ACLSHMEM 以进行多进程协调。 :param rank: 此进程的 rank 编号（0 到 nranks-1） :param nranks: 进程/rank 总数 :param mem_size: 对称堆大小（字节） :param uid: 从 aclshmem_get_unique_id() 获取的唯一 ID 字节 :returns: InitStatus —— 初始化状态 使用模式： .. code-block:: python import shmem # 在 rank 0 上：生成并广播唯一 ID if rank == 0: uid = shmem.aclshmem_get_unique_id() # ... 将 uid 广播到所有 rank ... else: # ... 从 rank 0 接收 uid ... pass # 所有 rank 使用相同的 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 生成用于协调多进程 SHMEM 初始化的唯一 ID。 :returns: bytes —— 要在所有进程间共享的唯一标识符 通常由 rank 0 调用并广播到其他 rank。所有参与进程必须使用相同的唯一 ID。 内存管理 -------- aclshmem_malloc ^^^^^^^^^^^^^^^ .. py:function:: aclshmem_malloc(nbytes: int) -> int 从 SHMEM 堆分配对称内存。 :param nbytes: 要分配的字节数 :returns: int —— 指向已分配内存的指针（整数形式），失败时返回 0 内存在所有 PE 的对称堆中以相同偏移量分配。所有 PE 必须使用相同的 ``nbytes`` 值调用此函数。 用于将通过 RMA 操作远程访问的内存。 示例： .. code-block:: python # 分配 4MB 对称内存 ptr = shmem.aclshmem_malloc(4 * 1024 * 1024) if ptr == 0: raise RuntimeError("对称内存分配失败") aclshmem_free ^^^^^^^^^^^^^ .. py:function:: aclshmem_free(ptr: int) -> None 释放先前使用 ``aclshmem_malloc`` 分配的对称内存。 :param ptr: 要释放的内存指针（整数形式） 所有 PE 必须集体释放对称分配。不要释放未使用 ``aclshmem_malloc`` 分配的内存。 aclshmem_calloc ^^^^^^^^^^^^^^^ .. py:function:: aclshmem_calloc(count: int, size: int) -> int 分配并零初始化对称内存。 :param count: 元素数量 :param size: 每个元素的字节大小 :returns: int —— 指向已分配内存的指针，失败时返回 0 等价于 ``aclshmem_malloc(count * size)`` 后进行清零。所有 PE 必须使用相同参数调用。 aclshmem_align ^^^^^^^^^^^^^^ .. py:function:: aclshmem_align(alignment: int, size: int) -> int 分配对齐的对称内存。 :param alignment: 对齐要求（字节，必须是 2 的幂） :param size: 要分配的字节数 :returns: int —— 指向对齐内存的指针，失败时返回 0 示例： .. code-block:: python # 分配 1MB 内存，对齐到 4KB 边界 ptr = shmem.aclshmem_align(4096, 1024 * 1024) aclshmem_ptr ^^^^^^^^^^^^ .. py:function:: aclshmem_ptr(local_ptr: int, pe: int) -> int 获取远程 PE 上对称对象的指针（主机可访问）。 :param local_ptr: 本地对称内存指针 :param pe: 目标 PE 编号 :returns: int —— 可在主机上用于访问指定 PE 上对称对象的指针，如果不可访问则返回 0 对于节点间远程指针，如果不支持直接主机访问，可能返回 0。主要用于节点内场景。 aclshmemx_get_heap_base ^^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmemx_get_heap_base() -> int 获取对称堆的基地址。 :returns: int —— 对称堆的基地址 PE 信息 ------- my_pe ^^^^^ .. py:function:: my_pe() -> int :no-index: 获取调用进程的 PE 编号。 :returns: int —— PE 编号（0 到 pe_count()-1） pe_count ^^^^^^^^ .. py:function:: pe_count() -> int 获取 PE 总数。 :returns: int —— 系统中的 PE 总数 主机 RMA 操作 ------------- aclshmem_putmem ^^^^^^^^^^^^^^^ .. py:function:: aclshmem_putmem(dest: int, source: int, bytes: int, pe: int) -> None 主机端阻塞写入（put）到远程对称内存。 :param dest: 远程 PE 上的目标指针 :param source: 本地主机上的源指针 :param bytes: 要传输的字节数 :param pe: 目标 PE 编号 示例： .. code-block:: python import numpy as np import shmem # 分配对称内存 remote_buf = shmem.aclshmem_malloc(1024) # 创建主机数据 data = np.arange(256, dtype=np.float32) # 将数据写入 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 主机端阻塞读取（get）远程对称内存。 :param dest: 本地主机上的目标指针 :param source: 远程 PE 上的源指针 :param bytes: 要传输的字节数 :param pe: 源 PE 编号 aclshmem_putmem_nbi ^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_putmem_nbi(dest: int, source: int, bytes: int, pe: int) -> None 主机端非阻塞写入到远程对称内存。 :param dest: 远程 PE 上的目标指针 :param source: 本地主机上的源指针 :param bytes: 要传输的字节数 :param pe: 目标 PE 编号 操作可能异步完成。在重用源缓冲区或假设远程可见性之前，使用适当的同步。 aclshmem_getmem_nbi ^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_getmem_nbi(dest: int, source: int, bytes: int, pe: int) -> None 主机端非阻塞读取远程对称内存。 :param dest: 本地主机上的目标指针 :param source: 远程 PE 上的源指针 :param bytes: 要传输的字节数 :param pe: 源 PE 编号 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 主机端阻塞写入，完成时原子性信号操作。 :param dest: 远程 PE 上的目标指针 :param source: 本地主机上的源指针 :param nbytes: 要传输的字节数 :param sig_addr: 远程 PE 上的信号地址 :param signal: 信号值 :param sig_op: 信号操作（SET 或 ADD） :param pe: 目标 PE 编号 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 主机端非阻塞写入，带信号。 :param dest: 远程 PE 上的目标指针 :param source: 本地主机上的源指针 :param nbytes: 要传输的字节数 :param sig_addr: 远程 PE 上的信号地址 :param signal: 信号值 :param sig_op: 信号操作（SET 或 ADD） :param pe: 目标 PE 编号 aclshmem_signal_wait_until ^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_signal_wait_until(sig_addr: int, cmp: int, cmp_val: int) -> None 主机端等待信号位置直到满足条件。 :param sig_addr: 信号位置指针 :param cmp: 比较操作（参见设备 API 的常量） :param cmp_val: 比较值 团队管理 -------- team_split_strided ^^^^^^^^^^^^^^^^^^ .. py:function:: team_split_strided(start: int, stride: int, size: int, parent_team) -> team 通过从父团队中选择具有跨步模式的 PE 创建新团队。 :param start: 父团队中的起始 PE :param stride: 所选 PE 之间的跨步 :param size: 新团队中的 PE 数量 :param parent_team: 父团队句柄 :returns: 新创建团队的团队句柄 示例： .. code-block:: python # 从全局团队创建包含 PE 0、2、4、6 的团队 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 将 PE 拆分为团队的 2D 网格。 :param xdim: X 维度大小 :param xaxis_teams: X 轴团队的输出数组 :param ydim: Y 维度大小 :param yaxis_teams: Y 轴团队的输出数组 team_translate_pe ^^^^^^^^^^^^^^^^^ .. py:function:: team_translate_pe(src_team, src_pe: int, dest_team) -> int :no-index: 将 PE 编号从一个团队转换到另一个团队。 :param src_team: 源团队句柄 :param src_pe: 源团队中的 PE 编号 :param dest_team: 目标团队句柄 :returns: int —— 目标团队中对应的 PE 编号 team_my_pe ^^^^^^^^^^ .. py:function:: team_my_pe(team) -> int :no-index: 获取调用 PE 在指定团队内的编号。 :param team: 团队句柄 :returns: int —— 团队内的 PE 编号 team_n_pes ^^^^^^^^^^ .. py:function:: team_n_pes(team) -> int :no-index: 获取指定团队中的 PE 数量。 :param team: 团队句柄 :returns: int —— 团队中的 PE 数量 team_destroy ^^^^^^^^^^^^ .. py:function:: team_destroy(team) -> None 销毁团队并释放其资源。 :param team: 要销毁的团队句柄 不要销毁预定义的团队。团队中的所有 PE 必须集体调用此函数。 配置与信息 ---------- InitAttr ^^^^^^^^ .. py:class:: InitAttr 用于配置 ACLSHMEM 运行时的初始化属性类。 **属性：** .. py:attribute:: mem_size 对称堆大小（字节） InitStatus ^^^^^^^^^^ .. py:class:: InitStatus 初始化函数返回的状态对象。 OpEngineType ^^^^^^^^^^^^ .. py:class:: OpEngineType 操作引擎类型枚举。 aclshmem_info_get_version ^^^^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_info_get_version() -> str 获取 ACLSHMEM 库版本字符串。 :returns: str —— 版本字符串 aclshmem_info_get_name ^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_info_get_name() -> str 获取 ACLSHMEM 库名称。 :returns: str —— 库名称 set_log_level ^^^^^^^^^^^^^ .. py:function:: set_log_level(level: int) -> None 配置 ACLSHMEM 日志级别。 :param level: 日志级别（使用标准 Python 日志级别） 实用函数 -------- aclshmem_global_exit ^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_global_exit(status: int) -> None 使用指定状态代码执行所有 PE 的全局退出。 :param status: 退出状态代码 终止 SHMEM 作业中的所有 PE。用于协调的错误处理。 aclshmem_create_tensor ^^^^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_create_tensor(shape: tuple, dtype: torch.dtype = torch.float32, device_id: int = 0) -> torch.Tensor 创建由对称内存支持的 PyTorch 张量。 :param shape: 张量形状元组 :param dtype: PyTorch 数据类型（默认：torch.float32） :param device_id: NPU 设备 ID（默认：0） :returns: torch.Tensor —— 由对称内存支持的张量 示例： .. code-block:: python import torch import shmem # 创建对称张量 tensor = shmem.aclshmem_create_tensor( (1024, 1024), dtype=torch.float32, device_id=0 ) # 像使用任何其他 PyTorch 张量一样使用 tensor.fill_(0.0) aclshmem_free_tensor ^^^^^^^^^^^^^^^^^^^^ .. py:function:: aclshmem_free_tensor(tensor: torch.Tensor) -> None 释放使用 ``aclshmem_create_tensor`` 创建的张量背后的对称内存。 :param tensor: 要释放的张量 仅用于通过 ``aclshmem_create_tensor`` 创建的张量。调用此函数后不应使用该张量。 完整示例 -------- .. code-block:: python import shmem import torch import numpy as np # 初始化 SHMEM attr = shmem.InitAttr() attr.mem_size = 1024 * 1024 * 1024 # 1GB status = shmem.aclshmem_init(attr) # 获取 PE 信息 my_rank = shmem.my_pe() num_ranks = shmem.pe_count() print(f"PE {my_rank} of {num_ranks}") # 分配对称内存 sym_ptr = shmem.aclshmem_malloc(1024 * 1024) # 1MB if sym_ptr == 0: raise RuntimeError("对称内存分配失败") # 或创建对称张量 sym_tensor = shmem.aclshmem_create_tensor((1024, 256), dtype=torch.float32) # 执行通信 if my_rank == 0: # Rank 0 向 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 接收（数据已在对称内存中） pass # 清理 shmem.aclshmem_free_tensor(sym_tensor) shmem.aclshmem_free(sym_ptr) shmem.aclshmem_finalize() 安装说明 -------- ACLSHMEM 包单独分发，必须安装： .. code-block:: bash pip install aclshmem 或从 ``3rdparty/shmem`` 目录中的源码构建。详见 :doc:`../getting-started/build` 指南。 注意事项 -------- 对称内存模型 ^^^^^^^^^^^^ SHMEM 使用对称内存模型，其中分配在所有 PE 上以相同虚拟地址偏移量发生。这使得无需显式地址 转换即可实现高效的远程内存访问。 集合操作 ^^^^^^^^ 许多 SHMEM 操作是集合操作，必须由团队或全局中的所有 PE 调用。示例包括： - ``aclshmem_init`` / ``aclshmem_finalize`` - ``aclshmem_malloc`` / ``aclshmem_free`` - 屏障操作 - 团队创建和销毁 线程安全 ^^^^^^^^ SHMEM 操作通常不是线程安全的。如果每个 PE 使用多个线程，应用程序必须添加适当的同步。