SHMEM 主机端 API

本模块提供主机端 ACLSHMEM 操作，用于初始化 SHMEM 运行时、管理对称内存以及执行主机发起的通信。这些 API 从 Python 主机代码调用，而非从 kernel 内部调用。

ACLSHMEM 包必须单独安装。使用以下命令安装：

pip install aclshmem

所有主机端 API 可通过以下方式访问：

import shmem

源文件： 3rdparty/shmem/src/python/shmem/__init__.py

关于完整的 ACLSHMEM 文档，请参考 ACLSHMEM 文档。

初始化与终止

aclshmem_init

aclshmem_init(attributes: InitAttr) → InitStatus

使用指定属性初始化 ACLSHMEM 运行时。

参数:: attributes -- 包含初始化配置的 InitAttr 对象
返回:: InitStatus —— 包含初始化结果的状态对象

必须在任何其他 SHMEM 操作之前调用。每个 PE 必须调用此函数。

示例：

import shmem

attr = shmem.InitAttr()
attr.mem_size = 1024 * 1024 * 1024  # 1GB 对称堆
status = shmem.aclshmem_init(attr)

aclshmem_finalize

aclshmem_finalize() → None

终止并清理 ACLSHMEM 运行时。

必须在程序终止前由所有 PE 调用。释放所有对称内存和通信资源。调用此函数后，无法执行其他 SHMEM 操作。

示例：

# 程序结束时
shmem.aclshmem_finalize()

aclshmem_init_using_unique_id

aclshmem_init_using_unique_id(rank: int, nranks: int, mem_size: int, uid: bytes) → InitStatus

使用唯一 ID 初始化 ACLSHMEM 以进行多进程协调。

参数:

rank -- 此进程的 rank 编号（0 到 nranks-1）
nranks -- 进程/rank 总数
mem_size -- 对称堆大小（字节）
uid -- 从 aclshmem_get_unique_id() 获取的唯一 ID 字节

返回:

InitStatus —— 初始化状态

使用模式：

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

aclshmem_get_unique_id() → bytes

生成用于协调多进程 SHMEM 初始化的唯一 ID。

返回:: bytes —— 要在所有进程间共享的唯一标识符

通常由 rank 0 调用并广播到其他 rank。所有参与进程必须使用相同的唯一 ID。

内存管理

aclshmem_malloc

aclshmem_malloc(nbytes: int) → int

从 SHMEM 堆分配对称内存。

参数:: nbytes -- 要分配的字节数
返回:: int —— 指向已分配内存的指针（整数形式），失败时返回 0

内存在所有 PE 的对称堆中以相同偏移量分配。所有 PE 必须使用相同的 nbytes 值调用此函数。用于将通过 RMA 操作远程访问的内存。

示例：

# 分配 4MB 对称内存
ptr = shmem.aclshmem_malloc(4 * 1024 * 1024)
if ptr == 0:
    raise RuntimeError("对称内存分配失败")

aclshmem_free

aclshmem_free(ptr: int) → None

释放先前使用 aclshmem_malloc 分配的对称内存。

参数:: ptr -- 要释放的内存指针（整数形式）

所有 PE 必须集体释放对称分配。不要释放未使用 aclshmem_malloc 分配的内存。

aclshmem_calloc

aclshmem_calloc(count: int, size: int) → int

分配并零初始化对称内存。

参数:

count -- 元素数量
size -- 每个元素的字节大小

返回:

int —— 指向已分配内存的指针，失败时返回 0

等价于 aclshmem_malloc(count * size) 后进行清零。所有 PE 必须使用相同参数调用。

aclshmem_align

aclshmem_align(alignment: int, size: int) → int

分配对齐的对称内存。

参数:

alignment -- 对齐要求（字节，必须是 2 的幂）
size -- 要分配的字节数

返回:

int —— 指向对齐内存的指针，失败时返回 0

示例：

# 分配 1MB 内存，对齐到 4KB 边界
ptr = shmem.aclshmem_align(4096, 1024 * 1024)

aclshmem_ptr

aclshmem_ptr(local_ptr: int, pe: int) → int

获取远程 PE 上对称对象的指针（主机可访问）。

参数:

local_ptr -- 本地对称内存指针
pe -- 目标 PE 编号

返回:

int —— 可在主机上用于访问指定 PE 上对称对象的指针，如果不可访问则返回 0

对于节点间远程指针，如果不支持直接主机访问，可能返回 0。主要用于节点内场景。

aclshmemx_get_heap_base

aclshmemx_get_heap_base() → int

获取对称堆的基地址。

返回:: int —— 对称堆的基地址

PE 信息

my_pe

my_pe() → int

获取调用进程的 PE 编号。

返回:: int —— PE 编号（0 到 pe_count()-1）

pe_count

pe_count() → int

获取 PE 总数。

返回:: int —— 系统中的 PE 总数

主机 RMA 操作

aclshmem_putmem

aclshmem_putmem(dest: int, source: int, bytes: int, pe: int) → None

主机端阻塞写入（put）到远程对称内存。

参数:

dest -- 远程 PE 上的目标指针
source -- 本地主机上的源指针
bytes -- 要传输的字节数
pe -- 目标 PE 编号

示例：

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

aclshmem_getmem(dest: int, source: int, bytes: int, pe: int) → None

主机端阻塞读取（get）远程对称内存。

参数:

dest -- 本地主机上的目标指针
source -- 远程 PE 上的源指针
bytes -- 要传输的字节数
pe -- 源 PE 编号

aclshmem_putmem_nbi

aclshmem_putmem_nbi(dest: int, source: int, bytes: int, pe: int) → None

主机端非阻塞写入到远程对称内存。

参数:

dest -- 远程 PE 上的目标指针
source -- 本地主机上的源指针
bytes -- 要传输的字节数
pe -- 目标 PE 编号

操作可能异步完成。在重用源缓冲区或假设远程可见性之前，使用适当的同步。

aclshmem_getmem_nbi

aclshmem_getmem_nbi(dest: int, source: int, bytes: int, pe: int) → None

主机端非阻塞读取远程对称内存。

参数:

dest -- 本地主机上的目标指针
source -- 远程 PE 上的源指针
bytes -- 要传输的字节数
pe -- 源 PE 编号

aclshmemx_putmem_signal

aclshmemx_putmem_signal(dest: int, source: int, nbytes: int, sig_addr: int, signal: int, sig_op: int, pe: int) → None

主机端阻塞写入，完成时原子性信号操作。

参数:

dest -- 远程 PE 上的目标指针
source -- 本地主机上的源指针
nbytes -- 要传输的字节数
sig_addr -- 远程 PE 上的信号地址
signal -- 信号值
sig_op -- 信号操作（SET 或 ADD）
pe -- 目标 PE 编号

aclshmemx_putmem_signal_nbi

aclshmemx_putmem_signal_nbi(dest: int, source: int, nbytes: int, sig_addr: int, signal: int, sig_op: int, pe: int) → None

主机端非阻塞写入，带信号。

参数:

dest -- 远程 PE 上的目标指针
source -- 本地主机上的源指针
nbytes -- 要传输的字节数
sig_addr -- 远程 PE 上的信号地址
signal -- 信号值
sig_op -- 信号操作（SET 或 ADD）
pe -- 目标 PE 编号

aclshmem_signal_wait_until

aclshmem_signal_wait_until(sig_addr: int, cmp: int, cmp_val: int) → None

主机端等待信号位置直到满足条件。

参数:

sig_addr -- 信号位置指针
cmp -- 比较操作（参见设备 API 的常量）
cmp_val -- 比较值

团队管理

team_split_strided

team_split_strided(start: int, stride: int, size: int, parent_team) → team

通过从父团队中选择具有跨步模式的 PE 创建新团队。

参数:

start -- 父团队中的起始 PE
stride -- 所选 PE 之间的跨步
size -- 新团队中的 PE 数量
parent_team -- 父团队句柄

返回:

新创建团队的团队句柄

示例：

# 从全局团队创建包含 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

team_split_2d(xdim: int, xaxis_teams, ydim: int, yaxis_teams) → None

将 PE 拆分为团队的 2D 网格。

参数:

xdim -- X 维度大小
xaxis_teams -- X 轴团队的输出数组
ydim -- Y 维度大小
yaxis_teams -- Y 轴团队的输出数组

team_translate_pe

team_translate_pe(src_team, src_pe: int, dest_team) → int

将 PE 编号从一个团队转换到另一个团队。

参数:

src_team -- 源团队句柄
src_pe -- 源团队中的 PE 编号
dest_team -- 目标团队句柄

返回:

int —— 目标团队中对应的 PE 编号

team_my_pe

team_my_pe(team) → int

获取调用 PE 在指定团队内的编号。

参数:: team -- 团队句柄
返回:: int —— 团队内的 PE 编号

team_n_pes

team_n_pes(team) → int

获取指定团队中的 PE 数量。

参数:: team -- 团队句柄
返回:: int —— 团队中的 PE 数量

team_destroy

team_destroy(team) → None

销毁团队并释放其资源。

参数:: team -- 要销毁的团队句柄

不要销毁预定义的团队。团队中的所有 PE 必须集体调用此函数。

配置与信息

InitAttr

class InitAttr

用于配置 ACLSHMEM 运行时的初始化属性类。

属性：

mem_size: 对称堆大小（字节）

InitStatus

class InitStatus: 初始化函数返回的状态对象。

OpEngineType

class OpEngineType: 操作引擎类型枚举。

aclshmem_info_get_version

aclshmem_info_get_version() → str

获取 ACLSHMEM 库版本字符串。

返回:: str —— 版本字符串

aclshmem_info_get_name

aclshmem_info_get_name() → str

获取 ACLSHMEM 库名称。

返回:: str —— 库名称

set_log_level

set_log_level(level: int) → None

配置 ACLSHMEM 日志级别。

参数:: level -- 日志级别（使用标准 Python 日志级别）

实用函数

aclshmem_global_exit

aclshmem_global_exit(status: int) → None

使用指定状态代码执行所有 PE 的全局退出。

参数:: status -- 退出状态代码

终止 SHMEM 作业中的所有 PE。用于协调的错误处理。

aclshmem_create_tensor

aclshmem_create_tensor(shape: tuple, dtype: torch.dtype = torch.float32, device_id: int = 0) → torch.Tensor

创建由对称内存支持的 PyTorch 张量。

参数:

shape -- 张量形状元组
dtype -- PyTorch 数据类型（默认：torch.float32）
device_id -- NPU 设备 ID（默认：0）

返回:

torch.Tensor —— 由对称内存支持的张量

示例：

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

aclshmem_free_tensor(tensor: torch.Tensor) → None

释放使用 aclshmem_create_tensor 创建的张量背后的对称内存。

参数:: tensor -- 要释放的张量

仅用于通过 aclshmem_create_tensor 创建的张量。调用此函数后不应使用该张量。

完整示例

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 包单独分发，必须安装：

pip install aclshmem

或从 3rdparty/shmem 目录中的源码构建。详见从源码构建指南。

注意事项

对称内存模型

SHMEM 使用对称内存模型，其中分配在所有 PE 上以相同虚拟地址偏移量发生。这使得无需显式地址转换即可实现高效的远程内存访问。

集合操作

许多 SHMEM 操作是集合操作，必须由团队或全局中的所有 PE 调用。示例包括：

aclshmem_init / aclshmem_finalize
aclshmem_malloc / aclshmem_free
屏障操作
团队创建和销毁

线程安全

SHMEM 操作通常不是线程安全的。如果每个 PE 使用多个线程，应用程序必须添加适当的同步。