[AI][shmem] Aiv 直驱RDMA

CANN aclshmem 之 AIV 直驱 RDMA 源码分析报告

文档说明

分析对象：CANN 开源仓库 cann/aclshmem（昇腾 SHMEM 对称内存通信库）

分析内容：AIV 直驱 RDMA 的功能实现与方案设计

重要提示 ：本文基于公开版本的架构与实现逻辑整理。文中代码为按实现逻辑重构的示意代码，请以aclshmem 仓为准。

1. 概述

复制代码

AIV kernel 内部
  → 在 HBM 中直接构造 NIC 认识的 WQE
  → AIV 直接向 doorbell 寄存器写值（MMIO store）
  → NIC 取描述符、DMA 读数据、发包

核心特征：

Host CPU 完全不参与数据面，只做初始化（建路）；
通信可在 kernel 运行中任意时刻发起，与计算与通信做 block 级重叠；
一次提交开销为几十至一百多 ns（几条指令 + 一次 MMIO）；
小包端到端（对端 flag 可见）延迟约 1.5 ~ 3 μs（RoCE 同机房）。

为什么由 AIV 执行

昇腾 AICORE 分为 AIC（Cube 核）与 AIV（向量核），都支持：

通信描述符构造、搬数、flag 轮询属于标量/访存密集工作，适合向量核；
AIC 可继续承担 Matmul，实现通算融合（如 MoE dispatch 与专家计算重叠）。

硬件前提（A2/A3 代际满足）

NIC 队列内存（SQ/CQ）可放在 NPU HBM，且 NIC 可直接 DMA 访问（SoC 内一致性或 P2P）；
doorbell 寄存器被映射进 AICORE 可寻址空间，初始化时由驱动交出 device 侧 DB 地址；
说明. AIV + URMA 的 Doorbell 在 A5 上，有其他实现方案，见源代码。
AICORE 写 HBM 与 NIC DMA 读之间具有缓存一致性；
通信 buffer 位于已注册 MR 内。

2. 术语与缩略语

术语	全称 / 含义
SHMEM	Symmetric Hierarchical Memory，对称内存 PGAS 编程模型
PE	Processing Element，SHMEM 中的 rank 概念
AIV / AIC	Ascend AI Core 的向量单元 / 矩阵（Cube）单元
WQE	Work Queue Element，提交给 NIC 的工作描述符
SQ / CQ	Send Queue / Completion Queue，NIC 队列
QP	Queue Pair，RDMA 连接端点（本文均为 RC 类型）
MR / lkey / rkey	Memory Region 及其本地/远端访问密钥
doorbell	NIC 门铃寄存器，写它即通知 NIC 消费 SQ
MMIO	Memory-Mapped I/O，用普通 store 指令访问设备寄存器
PGAS	Partitioned Global Address Space，分区全局地址空间
RoCE	RDMA over Converged Ethernet（本文使用 RoCE v2）
GM / UB	NPU 全局内存（HBM）/ 统一缓冲（片上 SRAM）
IBGDA	InfiniBand GPUDirect Async，NVSHMEM 的同构机制

3. 总体架构

3.1 分层视图

复制代码

┌────────────────────────────────────────────────────────┐
│ 用户算子 (AscendC kernel, AIV)                          │
│   shmem_putmem_nbi / shmem_putmem_signal_nbi / wait... │
├────────────────────────────────────────────────────────┤
│ aclshmem Device 库 (编译进 kernel)                      │
│   ├─ 地址翻译 (heap_base 表)                            │
│   ├─ SQ 槽位管理 (原子分配)                             │
│   ├─ WQE 构造                                          │
│   ├─ 缓存一致性 (clean/barrier)                         │
│   └─ doorbell MMIO 写入        ★ 直驱本体               │
├────────────────────────────────────────────────────────┤
│ Device Context (GM/HBM, init 时由 host 填充)            │
│   QP 设备视图 / SQ ring / CQ ring / signal 池 / 锁       │
├────────────────────────────────────────────────────────┤
│ On-chip RoCE NIC                                        │
│   DMA 取 WQE → DMA 读数据 → 组包发送                     │
└────────────────────────────────────────────────────────┘

┌────────────────────────────────────────────────────────┐
│ aclshmem Host 库 (libshmem.so, 仅初始化/控制面)          │
│   bootstrap / QP 建链 / MR 注册 / 对称堆 / context 下发  │
└────────────────────────────────────────────────────────┘

3.2 一句话概括

Host 面负责"建路"（QP、MR、对称堆、上下文下发），Device 面负责"走路"（AIV 造 WQE + 敲门铃 + 轮询 flag）。

4. 仓库结构导览

复制代码

aclshmem/
├── include/aclshmem/
│   ├── aclshmem.h            # Host 侧 API：init / malloc / team / 查询
│   ├── aclshmem_device.h     # Device 侧 API：put/get/signal/wait（__aicore__）
│   └── aclshmemx.h           # 扩展：putmem_signal_nbi 等融合原语
├── src/
│   ├── host/
│   │   ├── init/             # 初始化、bootstrap（交换 QPN/GID/heap 基址）
│   │   ├── mem/              # 对称堆：VA 预留、物理页映射、MR 注册、sub-alloc
│   │   ├── team/             # team 划分与资源复制
│   │   └── transport/
│   │       ├── rdma/         # RoCE QP 创建、RTS 建链、doorbell/SQ/CQ 资源
│   │       └── local/        # 节点内路径（IPC / HCCS 直拷）
│   ├── device/               # ★ AIV 直驱核心（编译进用户 kernel）
│   │   ├── *_rdma.*          # WQE 构造、SQ 管理、doorbell 写入
│   │   ├── *_signal.*        # signal / wait_until / quiet
│   │   └── include/          # WQE 格式、doorbell record、device context 定义
│   └── common/               # host/device 共享结构体与常量
├── examples/                 # put / signal 等双卡例程
└── tests/

5. Host 面：初始化与资源准备

5.1 `shmem_init_attr` 全流程

复制代码

shmem_init_attr(attr{my_rank, n_ranks, ip_port, ...})
│
├─ 1. bootstrap
│     基于 rank0 的 ip:port 建 TCP store（或复用 ranktable）
│     交换：device IP（RoCE v2）/ QPN / GID / PSN / heap 基址 / rkey
│
├─ 2. RDMA 建链
│     为每个对端 PE 创建 RC QP（1 个或多个以提升并发）
│     完成 INIT → RTR → RTS 状态机
│
├─ 3. 注册 MR
│     对称堆（或整片 HBM）注册给 NIC → 得到 lkey / rkey
│
├─ 4. 分配设备侧队列内存（全部在 HBM，NIC 与 AIV 双端可访问）
│     SQ ring / CQ ring / doorbell record / signal flag 池
│
├─ 5. 获取 device 侧 doorbell 地址
│     通过驱动私有接口拿到 AICORE 可写的 DB 映射地址
│
└─ 6. 组装 Device Context 并下发到 device 全局符号

前置条件 ：RoCE 链路需事先用 hccn_tool 配置好各 NPU 网卡 IP 并确认链路 up。

5.2 对称堆与地址翻译

shmem_malloc 是集合操作：

所有 PE 通过 aclrtReserveMemAddress + 物理页映射，预留相同的虚拟地址范围；
因此同一对象在所有 PE 上偏移一致，device 侧地址翻译只需查表：

c 复制代码

// 示意：本端 ptr 在 pe 上的对称地址
remote = ctx->heap_base[pe] + ((uint64_t)ptr - ctx->heap_base[my_pe]);

shmem_ptr(ptr, pe) 暴露的正是这一换算。

5.3 Device Context：直驱的"地图"

Host 把如下结构写入 device 全局符号（机制类比 cudaMemcpyToSymbol，CANN 通过 aclrtGetSymbolAddress + aclrtMemcpy 填充 __gm__ 全局变量）：

c 复制代码

// 示意重构：device 上下文
struct ShmemiDeviceCtx {
    int      my_pe, n_pes;
    uint64_t heap_base[MAX_PE];     // 各 PE 对称堆基址
    QpDev    qps[MAX_PE];           // 每对端 QP 的设备视图
    __gm__ uint64_t *sig_pool;      // signal flag 池
    __gm__ int32_t  *sq_lock;       // 每 QP 投递锁 / 原子计数
};

struct QpDev {
    __gm__ uint8_t *sq_buf;         // WQE ring（HBM，NIC 可读）
    uint32_t sq_depth, wqe_size;
    __gm__ uint32_t *sq_tail;       // 生产者指针（device 原子变量）
    uint64_t db_addr;               // ★ doorbell MMIO 地址（AICORE 可写）
    uint32_t qpn, lkey;
    uint64_t rheap_base;
    uint32_t rkey;
};

效果：kernel 内调用 shmem_putmem_nbi 无需任何初始化调用，直接读全局 context 即可。

6. Device 面：AIV 直驱数据通路

6.1 put 的四步曲（直驱本体）

c 复制代码

// 示意重构：shmem_putmem_nbi(dst, src, size, pe) 的 RDMA 路径
__aicore__ inline void shmemi_rdma_put(uint64_t laddr, int pe,
                                       uint64_t roff, uint32_t len)
{
    ShmemiDeviceCtx *ctx = &g_shmemi_ctx;      // device 全局符号
    QpDev *qp = &ctx->qps[pe];

    // ① 取 SQ 槽位（多 block 并发：GM 原子加）
    uint32_t idx = AtomicAdd(qp->sq_tail, 1);
    __gm__ Wqe *wqe = (__gm__ Wqe *)(qp->sq_buf +
                      (idx % qp->sq_depth) * qp->wqe_size);

    // ② 构造 WQE：RDMA_WRITE，SGE 直指源 buffer（零拷贝）
    wqe->ctrl.opcode = WQE_OP_RDMA_WRITE;
    wqe->sge.addr    = laddr;
    wqe->sge.len     = len;
    wqe->sge.lkey    = qp->lkey;
    wqe->raddr       = qp->rheap_base + roff;
    wqe->rkey        = qp->rkey;

    // ③ 一致性：确保 WQE 与数据对 NIC 可见
    dcci (wqe, sizeof(Wqe)) ;
    dsb();

    // ④ ★ 敲 doorbell：AIV 直接对 NIC 寄存器做一次 MMIO store
    DoorbellRec db = MakeDoorbell(qp->qpn, /*pi=*/idx + 1);
    *(__gm__ volatile uint64_t *)qp->db_addr = db.val;
}

四个要点：

零拷贝：SGE 直接指向用户/算子输出 buffer（只要在 MR 内），不做 staging；
批量摊薄：可连续构造多个 WQE 后只敲一次 doorbell（doorbell 携带最新 producer index，NIC 顺序消费），put + signal 可共一次 MMIO；
缓存一致性是阅读/移植最易踩坑处：AIV 的 GM 写需显式 clean + barrier 才对 NIC 可见；wait 侧读 flag 必须 volatile/绕过缓存；
doorbell = 一次 64bit MMIO store------这就是"直驱"的物理含义：kernel 里的 AIV 用一条 store 指令捅了一下 NIC 硬件。

6.2 signal 与 wait_until：靠 RC 序保证语义

关键洞察：同一 RC QP 上的 WQE 按序执行。

因此融合原语 putmem_signal_nbi 只需在同一 QP 上连发两个 WQE：

复制代码

WQE0: RDMA_WRITE  data → 对端 heap
WQE1: RDMA_WRITE  flag → 对端 sig_addr（或 WRITE_WITH_IMM）

对端看到 flag 更新 ⇒ 数据必然已落 HBM，零成本实现 "flag ⇒ data" 语义，无需额外 fence。

wait 侧为纯轮询：

c 复制代码

// 示意：shmem_signal_wait_until(sig, CMP_GE, expected)
__gm__ volatile int64_t *f = sig;
while (*f < expected) { /* nop backoff / timeout 检查 */ }

注意：

flag 必须在 GM（NIC 只能写 HBM）；
超时保护通过读 cycle counter 实现，避免链路故障时 kernel 死等。

6.3 fence / quiet：完成跟踪

语义	实现
`fence`（本地序）	保证 WQE/数据写先于 doorbell：cache clean + `PipeBarrier`
`quiet`（完成）	CQ ring 同在 device 内存，AIV 解析 CQE 推进 cons index（或读 NIC 回写的硬件 consumer index），等 `ci >= 投递时快照的 pi` 即返回

6.4 并发与多 QP

多 block 并发投递同一 QP：用 GM 原子加分配 SQ 槽位（或 per-QP device 锁）；doorbell 允许各敲各的（重复宣告 pi 语义幂等），也可约定"最后一个"统一敲；
每对端多 QP：按 block 哈希分流，牺牲跨 QP 序换并行度；
序敏感的 signal 固定走与数据相同的 QP（否则 "flag ⇒ data" 不成立）。

7. 端到端时序

以 MoE dispatch 一个 token 为例：

复制代码

PE0 AIV kernel              PE0 on-chip NIC              PE1 NPU
   │ 写数据到本端 heap buffer    │                            │
   │ 构造 WQE0 (WRITE data)    │                            │
   │ 构造 WQE1 (WRITE flag+1)  │                            │
   │ cache clean + barrier     │                            │
   │ MMIO doorbell ──────────► │ DMA 取 WQE                 │
   │ （kernel 继续算下一批）      │ DMA 读 data ── RoCE ─────► │ data 落 PE1 heap
   │                           │ 执行 WQE1 ──── RoCE ─────► │ flag 落 PE1 HBM
   │                           │                            │ PE1 AIV wait_until 命中
   │                           │                            │ → 读 data → 进入 expert 计算

要点：

Host CPU 全程无感知；
PE0 kernel 敲完门铃立即继续计算，实现粒度极细的通算重叠；
PE1 侧被动接收，单边语义，无需对端任何主动参与。

横向对比

对比项	说明
vs HCCL	HCCL 是 host 驱动的集合通信（粗粒度、流级）；aclshmem 是 device 驱动的 PGAS 原语库，两者互补。MC2 / MoE 融合算子可基于后者构建
vs NVSHMEM IBGDA	机制同构（GPU/AIV 造 WQE + 敲门铃）。差异：Ascend NIC 与 NPU 同 SoC，一致性模型与 doorbell 映射方式不同
vs A3 超节点 UB	同代际另有 UB（Unified Bus）内存语义互联，库内以独立 transport 存在；RDMA 路径主要覆盖跨节点/跨超节点

源码阅读路线

include/aclshmem/aclshmem.h + aclshmem_device.h ------ 建立 API 全景；
host 入口 shmem_init_attr → transport/rdma 建链 → heap 注册 ------ 理清"哪些资源进了 device context"；
src/common 中 WQE / doorbell 结构定义 ------ 理解硬件契约；
device 侧主线：putmem_nbi → 槽位分配 → WQE 填充 → cache clean → db 写，再读三个变体：
- putmem_signal_nbi（融合连发）
- signal_wait_until（轮询 + 超时）
- quiet（CQE 完成跟踪）
跑通 examples/ 双卡 put+signal 例程，再回头对照源码；
进阶：多 block 并发投递、多 QP 分流、team 资源复制。

常见问题 FAQ

Q1：kernel 里调用 shmem device API 前需要初始化吗？

不需要。context 在 shmem_init 阶段已由 host 写入 device 全局符号，kernel 直接引用。

Q2：为什么 signal 能保证数据先到？

同一 RC QP 上 WQE 严格按序执行，data WQE 先于 flag WQE，因此"flag 可见 ⇒ 数据已落 HBM"。跨 QP 无此保证，所以序敏感场景必须同 QP 连发。

Q3：普通（非对称堆）GM buffer 能直接 put 吗？

取决于 MR 注册策略：若注册覆盖整片 HBM 则可以；否则只有对称堆内地址可用。以实际版本的注册范围为准。

Q4：wait_until 死等怎么办？

实现内置超时（cycle counter）；超时后按"快速失败"语义上报。排查方向：链路状态、对端是否真的发了 signal、flag 地址是否对称。

Q5：多 block 同时投递会乱序吗？

同 QP 内，槽位由原子加分配，NIC 按 pi 顺序消费，数据一致性由 WQE 内容保证；但不同 block 的提交先后语义需用户自行约定（谁敲门铃是幂等的，顺序取决于槽位）。

参考资料

仓库：gitcode.com/cann/aclshmem（昇腾 CANN 开源镜像）
CANN 官方文档：SHMEM 通信 / AscendC 编程指南

再次提示：本文代码为按公开实现逻辑重构的示意代码，用于讲解方案；请以仓库实际源码为准

[AI][shmem] Aiv 直驱RDMA

CANN aclshmem 之 AIV 直驱 RDMA 源码分析报告

1. 概述

为什么由 AIV 执行

硬件前提（A2/A3 代际满足）

2. 术语与缩略语

3. 总体架构

3.1 分层视图

3.2 一句话概括

4. 仓库结构导览

5. Host 面：初始化与资源准备

5.1 shmem_init_attr 全流程

5.2 对称堆与地址翻译

5.3 Device Context：直驱的"地图"

6. Device 面：AIV 直驱数据通路

6.1 put 的四步曲（直驱本体）

6.2 signal 与 wait_until：靠 RC 序保证语义

6.3 fence / quiet：完成跟踪

6.4 并发与多 QP

7. 端到端时序

横向对比

源码阅读路线

常见问题 FAQ

参考资料

5.1 `shmem_init_attr` 全流程