CANN SHMEM Primitives

导言

我是零基础的算子开发初学者。第一次面对混合了 C++、Ascend C、SIMT 和 SHMEM 的通信算子,难点并不只是某个 API 不认识,而是同一行里往往同时出现模板参数、地址空间、数据搬运和同步语义。我想迁移这样的文件,就需要先回答四个问题:数据在哪里、谁在执行、这一句改变什么、下一步何时能使用结果

这篇文章从实际遇到的代码片段出发,把接口拆成尽量小的学习单元,再逐项解释模板参数、普通参数、返回值、地址偏移和同步范围。内容核对了 CANN 官方文档及 CANN/asc-devkitCANN/shmem 公开源码。尚未取得原算子完整文件,因此不虚构原文件行号,也不把业务封装的推测写成 API 定义。 文中的示意代码用于理解语义,未在 NPU 上编译运行。

阅读地图

这些名称属于不同层次,不能全部当作同一种“基础 API”。

层次 要学习什么 本文中的例子
C++ 语法 如何定义、取值、转换、赋值 constexprstatic_castreinterpret_castsizeof[]return
编译器扩展 函数在哪执行、指针指向哪种内存 __aicore____simt_vf____gm____ubuf__、launch bounds
Ascend C 内存与计算 建立内存视图、分配片上缓冲、搬运与填充 InitBufferSetGlobalBufferDataCopyPadDataCopyDuplicate
执行与可见性 哪些工作必须先完成,哪些线程必须等待 SetFlag/WaitFlagThreadBarrierasc_threadfenceDataCacheCleanAndInvalid
SIMT 并行原语 启动线程、获得编号、并发修改同一位置 Simt::VF_CALLGetThreadIdxatomicAddatomicOr
SHMEM 通信 识别参与者、翻译对端地址、读取通信上下文 aclshmem_my_peaclshmem_n_pesaclshmem_ptraclshmemi_udma_qp_info_fetch
算子业务封装 选择通信窗口、组织状态区和数据区 GetBaseWindOutAddrByRankIdGetBaseWindAddrByRankIdSyncFunc

输入/输出有两个不同概念:函数返回值,以及函数通过指针或 Tensor 修改的数据。void 只说明没有返回值,不表示没有输出。例如 DataCopy(dst, src, count) 没有返回值,但会写入 dst 指向的内存。

为避免混淆,下文的“读写”表示数据的实际作用方向。官方文档有时把用于读改写的地址参数标成“输出”,把需要初始化的对象标成“输入”;理解代码时还要看它指向的数据或对象状态是否被修改。

1
2
3
4
5
flowchart LR
GM["GM 中的数据"] -->|"DataCopyPad / MTE2"| UB["本核 UB"]
UB -->|"Duplicate / SIMD / SIMT"| RESULT["UB 中的结果"]
RESULT -->|"DataCopy / MTE3"| OUT["GM 中的结果"]
PTR["本地对称地址 + 目标 PE"] -->|"aclshmem_ptr:地址翻译"| REMOTE["目标 PE 对应地址"]

图中最后一条路径只计算地址,没有搬运数据。GM 与本地内存之间的搬入、计算、搬出模型可参考 CANN Add 入门教程;对端地址翻译参见 SHMEM 接口定义

最小 C++ 语法

看懂一行代码

语法 最小含义 你的代码中的读法
T name; 声明类型为 T 的对象 DataCopyExtParams elasticInfoParams:声明搬运配置对象
T name = expression; 用右侧表达式的值初始化左侧新变量 int64_t pos = ScalarGetSFFValue<1>(mask)
name = expression; 将右侧结果赋给已有左侧对象 wqeWords[0] = make_ulonglong4(...)
T name{a, b, c}; 列表初始化;对这里的配置结构体按成员声明顺序填字段 {false, 0U, 0U, 0U} 对应四个 padding 字段
T* p p 是指向 T 的指针,保存地址 __gm__ float* 是 GM 中 float 数据的地址
*p 访问 p 指向的对象 *counter 是计数值;counter 是计数器地址
&x 取得变量 x 的地址 可用于将计数器地址传给原子操作
T& x 引用,是已有对象的别名 InitBuffer(TBuf<pos>& buf, ...) 可以修改传入缓冲对象的管理状态
p[i] 对原生指针,访问第 i 个元素,下标从 0 开始 dw[0]wqeWords[0]
tensor[i] 对 Ascend C Tensor,是重载操作,通常取得偏移后的 Tensor 视图 不能直接套用“返回一个数”的原生数组含义
object.Method(...) 在对象上调用成员函数 scalesGMTensor_.SetGlobalBuffer(...)
p->field 访问指针所指结构体的成员 qp_ctx_entry->db_addr 取 QP 上下文中的门铃地址
Namespace::Name 指定名称所属命名空间或类型 AscendC::HardEvent::MTE2_V 是枚举值
Fn<T>(...) 用模板参数 T 选择函数版本 Duplicate<float>:选择处理 float 的版本
Fn<1>(...) 模板参数也可以是编译期数值 ScalarGetSFFValue<1>:查找第一个 1
sizeof(T) 一个 T 占多少字节 在这里 sizeof(int32_t) 为 4,sizeof 表达式的类型为 size_t
static_cast<T>(x) 按 C++ 转换规则将值转为 T 将搬运字节数转成 uint32_t,不会做范围校验
reinterpret_cast<T*>(x) 将地址或指针重新解释为指定指针类型 把 UB 地址解释为 int32_t 指针;不会复制数据或转换元素数值
(T*)x C 风格强制转换 (__gm__ float*)scales;阅读时要同时检查元素类型和地址空间
constexpr 变量初始化必须能在编译期求值 适合线程数上限、固定长度等
const 不允许通过该限定对象进行相应修改 不必然是编译期常量,也不必然保证指针指向的数据不变
inline C++ 函数声明属性,允许满足规则的多处定义 不表示强制内联,更不表示同步执行
void 函数不通过 return value 返回值 输出可以写到参数指向的内存
return expr; 把表达式值返回给调用者 返回窗口地址,本身不发送数据
0U / 1U 无符号整数常量 配置计数、步长字段
0.0F float 类型的零 Duplicate<float>(..., 0.0F, ...)
args... 展开参数包 把 VF 的各个参数按顺序传递进去
& / | / ~ / << 在整数表达式中分别表示按位与、或、非、左移 位图、标志位和 WQE 字段编码;与取地址符 &x 要按上下文区分

变量名尾部的 _ 通常是成员变量命名习惯,没有额外语言语义scalesGMTensor_elasticInfoGMTensor_ 都是具体对象名,不是独立 API。

地址偏移的单位

1
2
3
4
5
__gm__ uint8_t* bytes = /* 已存在的合法地址 */;
__gm__ int32_t* ints = /* 已存在的合法地址 */;

// bytes + 3:向后移动 3 字节。
// ints + 3:向后移动 3 个 int32_t,即 12 字节。

偏移单位由参与运算的指针类型决定,不由变量名中的 Offset 决定。 对整数形式的地址进行加法时,则是普通整数运算,需要调用方自己约定单位。

void* 没有元素大小,标准 C++ 不允许直接对其做 + offset。部分设备编译器可能提供扩展,但迁移时应确认规则,或先明确转换成字节指针。SHMEM 的公开示例也采用先转 GM_ADDR、再加偏移的写法。源码示例

内存对象与数据搬运

地址空间与 Tensor

语法 / 类型 作用 需要区分的事情
__gm__ T* 指向设备 Global Memory 的 T 类型数据 修饰符不分配内存,也不自动建立跨卡映射
__ubuf__ T* 指向本核 Unified Buffer 中的数据 同一个 SIMT 线程块可通过共享 UB 协作;不表示所有核共用同一 UB
GM_ADDR Ascend C 中常见的 GM 地址宏,通常为 __gm__ uint8_t* 属于宏,不是普通 C++ 内置类型;具体展开以所用头文件为准
GlobalTensor<T> 描述 GM 中一段数据的视图 绑定地址不等于申请显存,不等于搬运
LocalTensor<T> 描述片上本地内存中一段数据的视图 本地位置不只 UB;不能将任意 LocalTensor 都强转为 UB 指针
TBuf<pos> 管理一段本地缓冲空间 pos 指定逻辑位置,VECCALC 等位置对应 UB
TQue<pos, depth> 管理流水阶段之间传递的 Tensor 队列深度和物理 buffer 数是不同参数,不要混为一谈
TPipe 管理片上缓冲和相关流水资源 与主机侧显存分配职责不同

地址限定符见 官方扩展语法;GM 地址和 Tensor 使用方式见 Add 教程。跨卡访问还依赖 SHMEM 的初始化、内存映射、传输引擎与拓扑。SHMEM 地址接口

InitBuffer:先准备缓冲空间

两个最常见的调用形式:

1
2
pipe.InitBuffer(buf, lenBytes);
pipe.InitBuffer(queue, numBuffers, lenBytesPerBuffer);
部分 输入 / 输出 含义与单位
pipe 管理状态被修改 TPipe 对象,负责分配和登记本地缓冲
buf 传入并初始化 TBuf<pos>&;为这个缓冲对象建立内存信息
queue 传入并初始化 队列对象;为其准备物理 buffer
numBuffers 输入 uint8_t,物理内存块个数;常用 1 或 2
lenBytes / lenBytesPerBuffer 输入 uint32_t字节数;队列版本表示每块大小
返回值 返回 bool 所核对实现正常路径返回 true;它不是申请到的地址,也不是 Tensor

示例:

1
2
3
AscendC::TBuf<AscendC::TPosition::VECCALC> buffer;
pipe.InitBuffer(buffer, 128 * sizeof(int32_t));
AscendC::LocalTensor<int32_t> local = buffer.Get<int32_t>();

这里申请 512 字节,再取一个按 int32_t 访问的视图。Get<int32_t>() 的模板参数选择元素类型,无普通参数;返回 LocalTensor<int32_t>不会再次申请一份内存

InitBuffer 的长度不足 32 字节对齐时会向上补齐。numBuffers=2 为 double buffer 准备两块存储,但实际流水重叠还需要队列和正确的生产消费流程。InitBuffer 不负责把数据清零。 如果需要零值,应另做填充。

依据:TPipe::InitBuffer 文档实现。部分官方页面虽在原型中写 bool,却在“返回值”处写“无”;这里按源码解释,不能将 false 当作一套已承诺的完整错误处理协议。

SetGlobalBuffer:给已有 GM 地址建立视图

原片段:

1
scalesGMTensor_.SetGlobalBuffer((__gm__ float*)scales);

从内向外读:

  1. scales:已有地址,通常来自核函数参数。
  2. (__gm__ float*)scales:按 GM 中的 float 数组地址解释;不将其他类型的数据逐元素转成 float
  3. scalesGMTensor_:预计是 GlobalTensor<float> 对象,最终仍应以声明为准。
  4. .SetGlobalBuffer(...):把地址绑定到这个对象,后续通过它做访问或搬运。
参数 / 返回 方向 说明
buffer 输入 指向 GM 数据的指针,基础元素类型与 Tensor 匹配
bufferSize,仅双参数版本存在 输入 uint64_t元素个数,不是字节数
Tensor 对象状态 被修改 保存数据地址,以及显式传入时的长度信息
返回值 void

双参数示例:scalesGMTensor_.SetGlobalBuffer((__gm__ float*)scales, 128) 表示绑定 128 个 float。单参数调用并没有提供真实数组长度,不能据此判断可访问范围;所引较新文档说明此时 GetSize() 为 0。官方文档

DataCopyExtParams:描述“搬多少、怎么跨块”

原片段:

1
2
3
4
5
6
7
8
9
DataCopyExtParams elasticInfoParams = {
1U,
static_cast<uint32_t>(
(ELASTIC_INFO_OFFSET + RANK_LIST_NUM * epWorldSizeOriginal_)
* sizeof(int32_t)),
0U,
0U,
0U
};

这是一条结构体初始化语句,还没有开始搬数据。 左侧声明配置对象,右侧按成员顺序提供五个值。

顺序 字段与类型 本例输入 作用
1 blockCount: uint16_t 1U 搬 1 个连续块;不是 1 个线程块,也不是 1 个元素
2 blockLen: uint32_t 计算后的字节数 每个连续块要搬多少字节
3 srcStride: uint32_t 0U 相邻源块之间额外留出的间隔;本例 GM 侧单位是字节
4 dstStride: uint32_t 0U 相邻目的块之间额外间隔;本例 UB 侧按 32 字节块
5 rsv: uint32_t 0U 保留字段,置零;不是偏移地址或事件编号

本例只搬一块,因此不存在“下一块”的跨块步长效果。对多块情况,stride 描述前块结束到后块开始的额外间隔,不能直接当成两块起点之间的距离。

右侧长度表达式分解为:

1
2
3
元素数量 E = ELASTIC_INFO_OFFSET + RANK_LIST_NUM × epWorldSizeOriginal_
字节数量 B = E × sizeof(int32_t) = E × 4
blockLen = 将 B 转为 uint32_t

由表达式可以确定:括号整体在这里被当作 int32 元素数量使用。 各个常量具体表示多少个头字段、几份 rank 列表,必须结合定义确认。

为了演示,假设三个数依次为 8、2、4,则 E=16B=64,本次配置搬运 64 字节。这只是演算例子,不是原算子的常量。

static_cast<uint32_t> 不会防止溢出;如果前面的乘法已溢出,最后转换不能补救。迁移时应检查中间类型、API 的实际长度上限和 UB 容量,必要时先提升为较宽整数再运算。字段和单位依据 DataCopyPad 文档

DataCopyPadExtParams:描述“如何填充”

1
2
3
DataCopyPadExtParams<int32_t> elasticInfoCopyPadParams{
false, 0U, 0U, 0U
};
部分 字段 / 类型 本例含义
<int32_t> 模板元素类型 填充值按 int32_t 表示
false isPad: bool 不采用结构体中 paddingValue 指定的填充值
第一个 0U leftPadding: uint8_t 左侧显式填充 0 个元素
第二个 0U rightPadding: uint8_t 右侧显式填充 0 个元素
第三个 0U paddingValue: int32_t 字段数值为 0,但本例不能据此推出硬件执行补零

“没有要求自定义填充”与“没有对齐填充”是两回事。 非对齐搬运仍可能占用向上对齐后的 UB 空间;尾部 dummy 的规则随配置和架构变化。所核对较新源码文档还允许 isPad=false 时使用外部 SetPadValue 配置。不要把未声明为有效数据的尾部参与后续计算。官方字段说明较新 GM→UB 文档

DataCopyPad:真正发起 GM → UB 搬运

1
2
3
4
5
6
DataCopyPad(
elasticInfoTensor_,
elasticInfoGMTensor_,
elasticInfoParams,
elasticInfoCopyPadParams
);
位置 参数 输入 / 输出 含义
1 elasticInfoTensor_ 输出数据 目的 LocalTensor<int32_t>,接收搬入数据
2 elasticInfoGMTensor_ 输入数据 GlobalTensor<int32_t>,指向 GM 中原数据
3 elasticInfoParams 输入配置 块数、每块字节数、源目的间隔
4 elasticInfoCopyPadParams 输入配置 左右填充和填充值规则
返回值 void 结果写入第一个参数所指内存

阅读为:按指定搬运和填充规则,将源 GM 中一块 int32 数据搬到目标本地缓冲。 Tensor 类型以匹配此重载为前提,原声明仍需核对。

常见原型中的 const LocalTensor<T>& dst 仍可以是数据输出。这里 const 约束的是传入的 Tensor 视图对象,不能据此推出它描述的存储永远只读。

搬运属于 MTE2 流水;后续由 Vector 或 Scalar 消费时,需要建立相应依赖,见下一节。接口及方向依据 DataCopyPad 官方说明

DataCopy 与 Duplicate

1
2
DataCopy(stageDoneTensor, cleanStageDoneTensor, UB_ALIGN_DATA_COUNT);
Duplicate<float>(cleanStageDoneTensor, 0.0F, UB_ALIGN_DATA_COUNT);
接口部分 方向 含义
DataCopy 第 1 参数 stageDoneTensor 输出 目的 Tensor,接收复制的数据
第 2 参数 cleanStageDoneTensor 输入 源 Tensor,复制其现有内容
第 3 参数 UB_ALIGN_DATA_COUNT 输入 连续搬运的元素个数,原型类型为 uint32_t
Duplicate<float><float> 模板输入 操作的数据类型
第 1 参数 cleanStageDoneTensor 输出 被填充的本地 Tensor
第 2 参数 0.0F 输入 每个元素要写入的 float 值
第 3 参数 UB_ALIGN_DATA_COUNT 输入 要填充的元素数量,常见原型为 const int32_t&
两个接口的返回值 均为 void

若两个 Tensor 都在 UB,第一句是 UB→UB;如果 stageDoneTensor 是 GlobalTensor,则是 UB→GM。不能只靠名字猜地址空间。如果源数据已填成零,再执行复制,才会把零复制到目的地。

连续版 DataCopy 要检查 count * sizeof(T) 的 32 字节对齐。所核对文档指出未对齐的搬运量会向下取整,尾部可能不被搬走。UB→UB 文档GM↔UB 文档

Duplicate 是批量填值,不复制另一个 Tensor。假设 UB_ALIGN_DATA_COUNT=8,float 情况下会写 8 个零,即 32 字节;常量实际值需看原文件。Duplicate 官方说明

GetPhyAddr 与 reinterpret_cast

1
2
3
reinterpret_cast<__ubuf__ int32_t*>(
validExpertIdsTensor_.GetPhyAddr()
)
部分 输入 / 输出 解释
validExpertIdsTensor_ 输入对象 已经建立的 LocalTensor 视图
.GetPhyAddr() 返回地址 无普通参数;NPU 侧常见返回类型为 uint64_t,CPU 调试模式可返回指针
<__ubuf__ int32_t*> 转换目标类型 告诉编译器将地址当作 UB 中 int32 数据的指针
整个表达式 返回指针值 可作为 VF 的地址参数传入,不复制内存

它不提取“专家 ID 的数值”,而是提取“存放专家 ID 的地址”。如果原 Tensor 不是 UB、生命周期已结束、元素布局不匹配或对齐不满足,强制转换不会让这些问题消失。GetPhyAddr(offset) 还有带元素偏移的重载;这里调用的是无参数版本。官方说明

同步与缓存

SetFlag / WaitFlag:让后续流水等待前序流水

1
2
AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
部分 输入 / 输出 含义
<HardEvent::MTE2_V> 模板参数 同步方向为 MTE2 → Vector,Vector 等待 MTE2
EVENT_ID0 普通输入参数 事件槽位标识;不是线程 ID、字节数或地址
SetFlag 建立事件 源流水在前序工作完成后发出相应事件
WaitFlag 等待事件 目标流水等待匹配事件,再执行依赖它的后续工作
返回值 两者均为 void

将异步流水想成两名协作者:搬运单元搬完,计算单元才能读。C++ 代码行的先后顺序本身不够表达不同硬件流水之间的完成关系。

事件类型 常见用途
MTE2_V GM→UB 搬完,Vector 再计算
MTE2_S GM→UB 搬完,Scalar 再读取本地数据
V_MTE3 Vector 算完,MTE3 再把结果搬到 GM
V_S Vector 算完,Scalar 再消费结果
MTE3_MTE2 搬出完成后,允许下一轮搬入复用相关缓冲

事件方向和编号需要配对;在 TPipe/TQue 管理的代码中应考虑框架已使用的事件资源。官方建议通过分配或获取事件 ID 的接口管理资源,不能随意将全部同步都写成 EVENT_ID0SetFlag / WaitFlag 官方文档

SyncFunc:需要识别封装内容

1
SyncFunc<AscendC::HardEvent::MTE2_S>();
部分 能确定的含义
<MTE2_S> 传入编译期枚举值,表达 MTE2→Scalar 这一同步方向
() 没有显式普通参数;配置可能来自模板、局部常量或外部状态
SyncFunc 不能仅凭名字确定返回类型、事件 ID 和完整实现

未在此次核对的两个公开仓快照中找到这个精确函数定义,但 SHMEM 示例有 SyncFuncStatic<event, eventId>(),内部就是配对的 SetFlagWaitFlag实际封装源码

因此,原片段很可能是同类封装;这是推断。只有确认原定义后,才能断言它确实等待了什么、是否还包含事件分配或释放。不要把未知函数自动等同于“全设备同步”。

DataCacheCleanAndInvalid:处理缓存可见性

常见形式:

1
2
3
4
5
AscendC::DataCacheCleanAndInvalid<
uint64_t,
AscendC::CacheLine::SINGLE_CACHE_LINE,
AscendC::DcciDst::CACHELINE_OUT
>(global);
参数 / 返回 位置 含义
uint64_t 模板参数 T Tensor 元素类型;这里不表示“只处理 8 字节”
SINGLE_CACHE_LINE 模板参数 entireType 针对传入地址所在缓存行操作
ENTIRE_DATA_CACHE 上一参数的另一种取值 操作本核整个 Data Cache,传入地址不再选择范围
CACHELINE_OUT 模板参数 dcciDst GM 相关一致性模式;不同模式和产品的适用性需查文档
global 普通参数 提供地址的 GlobalTensor;单行模式下用于确定操作地址
返回值 void,效果是缓存处理,不返回读到的数据

Scalar 访问 GM 可能经过本核 DCache。其他核或 DMA 更新 GM 后,缓存里可能仍是旧内容;Scalar 写入也可能需要刷新才能让外部看到。该接口处理这类缓存一致性问题。不是清零函数,也不是让所有 PE 到齐的 barrier。

所核对当前文档特别区分:单行模式下地址生效、dcciDst 不生效;全缓存模式下地址不生效、dcciDst 生效。还存在两模板参数重载和特定产品支持的 LocalTensor 版本;不能凭一个函数名泛化全部行为。官方文档当前源码文档

四种“同步”不能互换

机制 解决的问题 不应从它推出什么
SetFlag/WaitFlag 本核不同流水之间的生产消费依赖 所有线程或所有卡已经到齐
ThreadBarrier / asc_syncthreads 当前线程块内线程到齐和共享数据可见性 所有核或所有 PE 已完成
asc_threadfence 调用线程内存写入的可见性与顺序 其他线程已经执行到此处,或远端网络传输已经完成
DataCacheCleanAndInvalid 缓存中的数据与目标存储的一致性 DMA、网络队列和所有线程都已结束

依据:流水事件线程同步内存栅栏缓存处理。跨 PE 的通信完成条件还要依据具体 SHMEM 操作的完成语义。

SIMT 线程与原子操作

函数修饰、线程上限与启动

1
2
3
4
__simt_vf__ __aicore__ LAUNCH_BOUND(PREPARE_THREADS_NUM)
inline void Prepare(__ubuf__ int32_t* counters, int32_t count);

Simt::VF_CALL<Prepare>(Simt::Dim3{N}, counters, count);

这是沿用问题中名称的结构示意,不是已验证的可独立编译示例

语法 / API 参数或输出 作用及注意点
__aicore__ 无参数 表示函数运行在 AI Core 设备侧;不是开启多核并行的指令
__simt_vf__ 无参数 标记 SIMT VF 入口;函数体描述每个线程执行的工作
LAUNCH_BOUND(N) N 为编译期上限 这是片段中的宏拼写;精确定义还需配套编译器头文件。官方扩展文档描述的是 __launch_bounds__(N)
__launch_bounds__(N) 输入最大线程数 给编译器资源分配提供上限;声明上限不等于真正启动 N 个线程
Simt::Dim3{x,y,z} 三个维度 实际线程总数为 x*y*zDim3 在所查封装中是 cce::dim3 的别名
Simt::VF_CALL<Fn>(dims, args...) Fn 是模板参数;dims 和 VF 参数是普通参数 提交指定 VF 子任务,返回类型为 void

当前扩展文档说明 dim3 未指定的维度默认为 1,因而一维 Dim3{N} 的意图是启动 N×1×1 个线程。上限必须覆盖实际线程总数;所核对混合编程文档的上限范围是 1~2048,这个范围不应推广到所有芯片和所有 CANN 版本。线程配置dim3 类型

Simt::VF_CALL<Fn>(Dim3{N}, a, b),参数逐项为:

  1. **Fn**:编译期选定的 VF 入口函数,可以是模板实例;不是运行时字符串。
  2. **Dim3{N}**:线程块形状,决定这次启动的线程数量;不是启用多少个核。
  3. **a, b**:按 Fn 声明顺序传进去。调用层传递的是值或地址,但指针所指数据是输入、输出还是读写,必须看 Fn 的函数体。
  4. **返回 void**:VF 结果通常通过传入地址写出。

VF 调用不等于通用阻塞屏障

在此次核对的 Simt::VF_CALL 源码中,内部调用 cce::async_invoke<funcPtr>(threadNums, args...),封装里没有显式等待。这个证据不足以支持“调用线程阻塞等待所有线程完成”的通用描述,也不能只根据 async 这个名字推断所有硬件依赖细节。迁移时需要结合目标编译器和后续消费者的流水依赖确认完成语义。

也不能把 SIMD 寄存器 VF 和 SIMT VF 的同名调用规则混用。

依据:Simt::VF_CALL 实现。本例 PREPARE_THREADS_NUM 是否等于 1024、原文件哪些位置启动 VF,仍需原文件定义;用户提供的行号不在本文中当作已核实定位。

GetThreadIdx:我是哪一个线程

1
uint32_t tid = Simt::GetThreadIdx<0>();
部分 输入 / 输出 说明
<0> 编译期模板参数 dim 取 x 维编号;1 对应 y,2 对应 z
() 无普通参数 编号来自运行环境
右侧返回值 uint32_t x 维线程编号,即 threadIdx.x
左侧 tid 新变量 保存当前线程自己的编号

一维启动 N 个线程时,编号范围为 0~N−1。若启用多个核,每个核的线程块可能都有自己的 tid=0tid 不是全设备唯一编号,更不是 SHMEM rank。三维情况下仅取 x 维也不是整个线程块的线性编号。接口实现

atomicAdd:计数器与取号

1
int32_t ticket = atomicAdd(counter, 1);
部分 输入 / 输出 含义
counter 地址输入;所指数据读写 int32_t*,本例需指向此重载支持的 UB 或 GM
1 输入 本次增加的值,类型与计数器匹配
返回值 输出 加法之前的旧值
左侧 ticket 接收返回值 当前线程拿到的编号
*counter 内存输出 原值加 1

假设初值为 0,三个线程分别原子加 1,则返回值形成 0、1、2,最终计数为 3;哪个线程拿到哪个编号由执行顺序决定。这是“计数器 + 取号机”的准确含义。

当前 Ascend C 的 asc_atomic_add 整数实现直接调用 atomicAdd,可据此核对片段中该 intrinsic 的语义。封装实现官方原子加说明

原子性只覆盖同一地址的读改写,不能把整个复杂代码段变成原子事务。尤其要区分“已经领取多少个任务”和“已经完成多少个任务”:拿到 ticket 后再写数据,不代表增加计数器时数据就准备好了。

本例以 32 位整数讲解。所核对接口中 64 位整数原子加只支持 GM;half/bfloat16 的返回值还有专门限制,不能把 32 位整数的取号模式照搬到所有数据类型。当前数据类型约束

atomicOr:并发设置位图

1
2
uint32_t bit = 1U << k;      // 前提:0 <= k < 32
uint32_t oldMask = atomicOr(maskPtr, bit);
参数 / 输出 含义
maskPtr 指向共享位图的地址,所指数据会被读取和修改
bit 本次需要置 1 的位;也可以同时包含多个位
返回值 oldMask 更新前的位图
更新后 *maskPtr 旧位图 | bit

例如原值为二进制 0010、本次输入 0100,更新后为 0110,返回值仍为 0010。它可以安全合并多个线程的置位请求,避免普通 *maskPtr |= bit 在并发读改写时丢失更新。

其公开包装 asc_atomic_or 调用 atomicOr;32 位整数支持 UB/GM,所查 64 位整数重载只支持 GM。实现原子或文档

ThreadBarrier 与 asc_threadfence

API 输入 / 返回 核心含义
Simt::ThreadBarrier() 无普通参数,无返回值 当前线程块内所有线程到达同步点后才能继续
asc_syncthreads() 无普通参数,无返回值 当前公开接口的线程块同步原语,提供块内共享数据可见性
asc_threadfence() 无普通参数,无返回值 约束调用线程对全局/共享内存写入的可见性顺序;不等待其他线程到齐

ThreadBarrier 的所查设备实现调用 __sync_workitems()实现。要理解其用途,可以类比 CUDA 的线程块屏障,但不应把不同平台的全部细节当作相同。官方线程同步说明

1
2
3
4
5
6
// 结构示意:所有线程都必须走到下面的 barrier。
if (tid < validCount) {
shared[tid] = input[tid];
}
Simt::ThreadBarrier();
// 后续线程才可按算法读取其他线程写入的 shared 数据。

如果把 barrier 放进只有部分线程会进入的分支,或者部分线程提前退出,就可能死锁。

在“先写数据、再发布就绪标志”的模式中,内存 fence 用来建立顺序;它不等于网络传输完成确认,也不能代替接收方必要的观察和同步操作。asc_threadfence 官方文档

ScalarGetSFFValue:从低位扫描 0 或 1

1
int64_t pos = AscendC::ScalarGetSFFValue<1>(mask);
部分 输入 / 输出 含义
<1> / <0> 模板参数 countValue 1 查找第一个 1;0 查找第一个 0,只允许这两种取值
mask 输入 uint64_t 被扫描的 64 位无符号整数
返回值 pos 输出 int64_t 从最低有效位开始计数,范围 0~63;找不到为 −1

28 的低位二进制为 11100,因此 <1>(28) 返回 2,<0>(28) 返回 0。<1>(0) 返回 −1;<0>(UINT64_MAX) 返回 −1。用这个返回值做下标或移位前,必须先处理 −1。官方定义与示例

它只是位扫描。将结果解释成“有效专家编号”“连续到达范围”还是“空闲槽位”,取决于原算子的位图编码。

make_ulonglong4:把四个数放进短向量

1
wqeWords[0] = make_ulonglong4(dw[0], dw[1], dw[2], dw[3]);
部分 输入 / 输出 含义
dw[0] 输入参数 x 转为第一个 unsigned long long 分量
dw[1] 输入参数 y 第二个分量
dw[2] 输入参数 z 第三个分量
dw[3] 输入参数 w 第四个分量
函数返回值 ulonglong4 .x/.y/.z/.w 四个 64 位分量的短向量
左侧 wqeWords[0] 被赋值的第 0 个元素 把整个返回对象保存到这里

所查实现就是依次设置四个成员后返回。ulonglong4 的对象大小和对齐要求在所查文档中均为 32 字节。实现类型说明

这不代表一次 32 字节原子写,也不会自动发送通信请求。 如果 wqeWords 是指向 GM WQE 区域的指针,赋值会写那个区域;如果它是局部数组,作用位置就不同。若 dw 实际是 32 位元素,传参会先扩展成 64 位数,并非把四个 32 位字原样打包成 16 字节。必须核对原声明。

SHMEM 地址与通信状态

PE、rank、线程编号

PE(Processing Element)是 SHMEM 通信参与者。 在常见一进程一设备配置中,它对应一个进程所管理的设备,但不能把 PE 编号直接当作物理卡号。程序中的专家并行 rank、子 team rank 和全局 SHMEM PE 编号也可能需要映射。

API 显式输入 返回类型与含义 示例
aclshmem_my_pe() int,当前调用 PE 在全局参与者中的编号,范围 0~npes−1 int me = aclshmem_my_pe();
aclshmem_n_pes() int,当前 SHMEM 程序的 PE 总数 int peers = aclshmem_n_pes();

这两个查询依赖已经初始化的 SHMEM 状态,调用本身不创建参与者。对于 8 个 PE,可以返回 peers=8me=3me=3 并不表示当前 SIMT 线程 tid=3官方接口头文件

aclshmem_ptr:求同一对称位置在对端的地址

公开设备侧原型:

1
__gm__ void* aclshmem_ptr(__gm__ void* ptr, int pe);
参数 / 输出 含义
ptr,输入 本 PE 上的对称内存地址;不是任意显存地址都能传入
pe,输入 要访问的目标全局 PE 编号
返回值 指定 PE 上对应对称位置的地址
内存输出 此查询不复制负载,也不改写目标数据

此次核对的设备实现采用如下地址换算:

1
2
offset        = ptr - 本地对称堆基址
remoteAddress = 目标 PE 的映射堆基址 + offset

例如本地基址为 0x1000ptr=0x1120,偏移为 0x120。若目标 PE 的映射基址是 0x8000,对应地址就是 0x8120。这表示相同的对称位置,不表示两端当前数据内容相同。实现

返回地址允许哪种方式访问,取决于传输引擎和拓扑,不能保证任何机器间都能直接用普通 load/store 访问;也不能把 aclshmem_ptr 当作“远程 memcpy”。接口契约

你的 SHMEM 地址表达式

1
2
3
aclshmem_ptr((GM_ADDR)shmemBuffer_, rankId)
+ dataBaseOffset
+ shmemDataSizeOffset_
部分 能确定的作用 还需核对的内容
(GM_ADDR)shmemBuffer_ 将保存的地址转成常见的 GM 字节指针形式 必须对应合法的 SHMEM 对称地址
rankId 选择目标 PE 是否已经从 EP/team rank 转成 SHMEM 全局编号
aclshmem_ptr(...) 返回对端对应地址 返回地址的可访问方式由引擎和拓扑决定
+ dataBaseOffset 继续移动到某个子区域 具体单位和区域定义
+ shmemDataSizeOffset_ 再叠加一个偏移 是否用于数据/状态区或缓冲轮转,不能只凭名字确定
整个表达式 计算地址值 没有函数调用来发送这些数据

本次核对的 aclshmem_ptr 返回 __gm__ void*,所以原写法的指针加法依赖编译器扩展或另一版本的签名。 如果两个偏移确实都是字节数,更明确的写法是:

1
2
3
4
auto remoteBase = reinterpret_cast<__gm__ uint8_t*>(
aclshmem_ptr(reinterpret_cast<GM_ADDR>(shmemBuffer_), rankId)
);
auto remoteData = remoteBase + dataBaseOffset + shmemDataSizeOffset_;

这里左侧 remoteBaseremoteData 都接收地址值,不接收远端内容。SHMEM 的公开 GetRemoteAddress 示例同样先转成 GM_ADDR 再加 offset,可作为阅读参照。示例源码

GetBaseWindOutAddrByRankId 与 GetBaseWindAddrByRankId

原片段:

1
2
3
4
5
return GetBaseWindAddrByRankId(
winContext_[COMM_EP_IDX],
rankId,
epRankIdOriginal_
) + winDataSizeOffset_;

没有在此次核对的两个公开仓快照中找到这两个名称的定义。 它们可能来自原算子、公共通信头文件或其他依赖。检索未找到不等于证明它们不是公开代码。

可以确定的 C++ 语义如下:

  1. winContext_[COMM_EP_IDX]:从上下文容器中选择一个元素,作为第 1 个实参。COMM_EP_IDX 是下标,具体值未知。
  2. rankId:第 2 个实参;按名称推测是目标 rank,但需定义确认。
  3. epRankIdOriginal_:第 3 个实参;按名称推测是原始 EP rank 编号,不能确定是否参与本地/远端分支或映射。
  4. GetBaseWindAddrByRankId(...):调用函数得到一个可参与加法的值,按名称推测为窗口基地址。
  5. + winDataSizeOffset_:在返回值上增加偏移。若返回值是指针,单位由元素类型决定;若是整数地址,单位由业务约定决定。
  6. return:把最终结果返回外层调用者。语句中没有显式数据搬运。

要完成原算子级解释,还缺函数声明、函数体、返回类型以及四个成员/常量的定义。当前可标记为“地址计算封装”,不能将其写成已确认的 CANN 或 SHMEM 标准 API。

aclshmemi_udma_qp_info_fetch:取通信上下文指针

所查函数的核心签名为:

1
__gm__ aclshmemi_aiv_udma_info_t* aclshmemi_udma_qp_info_fetch();
参数 / 输出 说明
显式普通参数 无;调用者没有传 rank、QP 编号或长度
隐式输入 已初始化的 SHMEM UDMA 运行时状态
返回类型 __gm__ aclshmemi_aiv_udma_info_t*
返回值 GM 中 UDMA 信息结构的指针;不是一个 QP 编号,也不是新建出的 QP
负载传输 这个函数本身不发送业务负载

当前实现取得 aclshmemi_get_udma_info_address(0) 的返回地址,转换为上述结构指针后返回。返回上下文可用于后续选择队列和对端内存信息;具体结构布局绑定到 SHMEM 版本。

函数位于 src/device/gm2gm/engine/ 内部实现,名称中的 aclshmemi_ 也与这一内部用途一致。代码公开不等于接口承诺稳定。 在基础学习清单中应把它放在公开 RMA 接口之后,迁移调用它的代码时应固定 SHMEM commit。实现与声明

st_dev:底层存储与门铃

原片段:

1
st_dev(val, addr, 0);
位置 输入 / 输出 本例的可确认解释
val 输入 待写入的标量;SHMEM 门铃示例中是 uint32_t 队列索引
addr 地址输入、目标存储输出 所查门铃示例为 __gm__ uint32_t*
0 输入 底层存储偏移参数取零;封装源码命名为 ASC_C_API_DEFAULT_OFFSET。非零值的单位、编码和范围需编译器 intrinsic 说明,不能在这里推定
返回 原样例不使用返回值 关注的是写入目标地址的副作用

st_dev 并不是专门表示“触发 URMA”的通用函数名。CANN 公开 asc_store_dev(addr, value) 的实现用到了 st_dev(value, addr, 0);其说明是绕过 DCache 向 GM 地址写数据。公开 store 接口说明底层包装实现

在 SHMEM 中,下面两个上下文把目标地址设置成门铃地址:

调用场景 写什么 作用
提交 SQ 工作 新的生产者 head 通知硬件有新的队列工作
消费 CQ 完成项 更新后的消费者 tail 通知硬件完成队列的消费进度

例如 aclshmemi_udma_post_send_update_info 把 QP 的 db_addr 转成 __gm__ uint32_t*,通过 st_dev(cur_head, door_bell_addr, 0) 写门铃,再更新软件上下文中的 head。SHMEM 实现

所以准确读法是:底层 store + 指向门铃的地址 + 对应硬件协议,共同构成通知操作。 不能用普通指针赋值随意替换,因为可能改变缓存路径和可见性;写了门铃也不代表传输已经完成。

把基础单元连起来

一条搬入与计算链

以下是逻辑顺序,不是可直接运行的完整算子:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
1. InitBuffer
给当前核准备 UB 空间。

2. SetGlobalBuffer
把输入 GM 地址绑定到 GlobalTensor。

3. DataCopyExtParams / DataCopyPadExtParams
描述真实数据长度、跨块间隔、填充策略。

4. DataCopyPad(local, global, copy, pad)
发起 GM → UB 搬运。

5. MTE2_V 或 MTE2_S 对应的同步
根据下一个消费者是 Vector 还是 Scalar 建立依赖。

6. Duplicate / SIMD 运算 / SIMT VF
处理 UB 或相关 GM 数据;按执行模式处理线程和流水依赖。

7. 对应的计算 → 搬出依赖,再 DataCopy
将结果写回目标存储。

源 Tensor、目的 Tensor 和有效长度应贯穿整个过程。例如实际搬入 13 个 int32,UB 因对齐占用更大空间,并不意味着后续算法可以把尾部都当作第 14~16 个有效元素。

一条通信提交链

按照 SHMEM 当前 UDMA 代码可概括为:

1
2
3
4
5
6
7
8
9
10
11
12
13
本 PE / 目标 PE 编号

找到对称位置及对端地址 / 内存描述

取得 UDMA 上下文,选择具体队列

构造 WQE(Work Queue Element,工作队列项)

保证描述符写入的可见性和所需依赖

写 SQ 门铃,发布队列工作

依据协议观察 CQE(Completion Queue Entry,完成队列项)

make_ulonglong4 只能完成其中局部的数据组装;st_dev 可以承担门铃写入;aclshmemi_udma_qp_info_fetch 提供上下文指针。三者都不能单独等同于一次已完成的远端数据传输。 这条链是对 SHMEM UDMA 实现 的概括,不表示已确认原算子的完整提交协议。

迁移前优先检查的差异

检查项 容易产生的错误 本文中的定位
元素数与字节数 长度多乘或少乘 sizeof(T) InitBufferSetGlobalBufferDataCopyExtParamsDataCopy
地址空间与映射 将本地地址当成远端可访问地址 __gm____ubuf__aclshmem_ptr
原子返回值 把旧值当新值,或把取号当完成 atomicAddatomicOr
编译期与运行期 将启动上限误当实际启动数量 launch bounds、Dim3VF_CALL
完成与可见性 只改同步函数名,却改变依赖范围 event、barrier、fence、cache clean
结构体与短向量布局 WQE 宽度、对齐或字段编码改变 reinterpret_castmake_ulonglong4
版本与内部 ABI 使用另一版本的 QP 布局解释指针 aclshmemi_udma_qp_info_fetchst_dev

推荐学习顺序:C++ 类型与指针 → GM/UB 与 Tensor → 搬运和填充 → 流水同步 → SIMT 线程与原子操作 → SHMEM 地址 → WQE 和门铃。 先能准确解释数据在哪里、多少字节、谁负责写,再进入通信队列细节。

回到迁移现场

再看到一条陌生语句时,我不必先记住整套 CANN。可以先把它拆成四层:

  1. C++ 层:左侧接收什么,右侧返回什么,转换只改类型还是也改变数据。
  2. 内存层:地址指向 GM、UB 还是 SHMEM 对称内存;长度的单位是元素还是字节。
  3. 执行层:当前是 Scalar、Vector、MTE 还是 SIMT 线程在工作。
  4. 完成层:当前操作只发起任务、保证可见性,还是确实等到了所需范围内的完成。

这套拆法已经能解释本文列出的公开基础单元,也暴露了仍然不能从片段确定的边界:GetBaseWindAddrByRankIdGetBaseWindOutAddrByRankId 和原算子的 SyncFunc 还缺定义;st_dev 的第三个参数在非零时如何解释,也需要对应编译器 intrinsic 的资料;完整迁移仍要回到目标芯片、CANN/SHMEM 版本和原算子声明。

所以下一步不是继续背函数名,而是拿到原文件后,从地址类型、长度单位和同步依赖开始逐段标注。只要这三件事没有猜,后续的接口替换才有可靠基础。

核对范围

核对日期:2026-09-03。所用公开源码固定为:

仓库 Commit 用途
CANN/asc-devkit 2ccf1031ba8085c0a6e44c2f271ba24f3363648f Ascend C 文档、SIMT 包装、原子操作与 store 实现
CANN/shmem a759c51376b6e0b4b8fa7d967bf859121346d6cb PE 查询、对称地址翻译、UDMA 上下文与门铃用法

CANN 网页引用包含不同版本,用来核对具体接口;不是声明这些 API 在所有版本和所有 Ascend 设备上都同时可用。原算子的目标芯片、CANN/SHMEM 版本、完整声明与真实行号尚未核实;本文完成的是问题中所列公开基础单元的解释。

Author

Shaojie Tan

Posted on

2026-09-03

Updated on

2026-09-04

Licensed under