CAM SHMEM 模式代码深度剖析
1. SHMEM 模式整体架构
1.1 定位与使命
SHMEM 模式是 CAM 的 用户态 RDMA 直通路径。它完全绕过 HCCL 集合通信库,在 AIV(Ascend Intelligent Vector)核心上直接操作 RoCE RDMA 硬件的 WQE/CQE/SQ/CQ 寄存器,实现 kernel 内部的细粒度 P2P 数据传输。
┌──────────────────────────────────────────────────────────────────────┐
│ Host Side (CPU) │
│ │
│ RdmaTransportManagerV2 │
│ ├── OpenDevice() → phyId 获取 + RDMA IP 解析 │
│ ├── CreateEndpoint() → HcommEndpointCreate(COMM_PROTOCOL_ROCE) │
│ ├── RegisterMemoryRegion() → HcommMemReg (HBM 注册为 MR) │
│ ├── Connect() → HcommChannelCreate(COMM_ENGINE_AIV) │
│ └── FillRdmaInfo() → AiQpRMAQueueInfo 拷贝到 HBM │
│ │
│ 关键常量: MAX_RANKS_PER_NIC=16, ATOMIC_MAX_NUM=128 │
│ 底层协议: dl_hccp_def.h (HCCP) + dl_hcomm_def.h (Hcomm) │
└──────────────────────────────────────────────────────────────────────┘
│
│ qpInfo_ → AclrtMemcpy(Host→Device)
▼
┌──────────────────────────────────────────────────────────────────────┐
│ Device Side (AIV Core, Ascend C) │
│ │
│ 粗粒度 API 层 (shmem_device_rdma.hpp) │
│ ├── aclshmemx_roce_get_nbi<T>() ← RDMA READ (远端→本地) │
│ ├── aclshmemx_roce_put_nbi<T>() ← RDMA WRITE (本地→远端) │
│ ├── aclshmemx_roce_quiet<T>() ← 同步等待所有 WQE 完成 │
│ └── aclshmemx_roce_atomic_*() ← 16 种原子操作变体 │
│ │
│ 细粒度原语层 (rdma_device_backend_base.h) │
│ ├── aclshmemi_roce_fill_wqe<B, OP>() ← 构造 WQE │
│ ├── aclshmemi_roce_ring_sq_doorbell<B>() ← 振铃 SQ Doorbell │
│ ├── aclshmemi_roce_post_send<B, OP>() ← post 发送 (fill+bell) │
│ ├── aclshmemi_roce_ring_cq_doorbell<B>() ← 振铃 CQ Doorbell │
│ └── aclshmemi_roce_poll_cq<B>() ← 轮询 CQE 完成 │
│ │
│ 后端特化层 (compile-time 选择) │
│ ├── IN_DIE — 片内互联 (HCCS, 当前未完整实现) │
│ └── XSCALE — 跨机 RoCE (DDRM Diamond 引擎, 完整实现) │
│ │
│ 硬件寄存器: SQ Doorbell / CQ Doorbell / WQE(128B) / CQE(64B) │
└──────────────────────────────────────────────────────────────────────┘
1.2 数据流全景
一次 aclshmemx_roce_put_nbi (RDMA WRITE) 的完整路径:
用户调用 aclshmemx_roce_put_nbi(dst, src, buf, elem_size, pe)
│
├─ 1. aclshmem_ptr(dst, pe) ← 本地对称地址 → 远端对称地址
│ 计算: remote = p2p_device_heap_base[pe] + (dst - heap_base)
│
├─ 2. aclshmemi_roce_write<B>() ← 填充 send_wr 结构体
│ wr.remote_addr = ptr
│ wr.local_addr = src
│ wr.message_len = elem_size * sizeof(T)
│
├─ 3. aclshmemi_roce_post_send<B, OP_RDMA_WRITE>()
│ │
│ ├─ 3a. 读取 SQ head/tail (dcci_cachelines 刷新缓存行)
│ │ 若 SQ 即将满 (head+16 == tail),先 poll_cq 等待空闲
│ │
│ ├─ 3b. aclshmemi_roce_fill_wqe<B, OP_RDMA_WRITE>()
│ │ → 在 HBM 上构造 128B WQE (ctrl_seg + remote_seg + local_seg)
│ │
│ ├─ 3c. dcci_cachelines(wqe_addr, 128) ← 刷回 HBM 确保硬件可见
│ │
│ └─ 3d. aclshmemi_roce_ring_sq_doorbell<B>()
│ → 更新 head_addr (producer index)
│ → 写 doorbell register: {next_pid, qp_id}
│
└─ 4. 硬件自动执行 RDMA WRITE,数据通过 RoCE 网络直达远端 HBM
2. 核心数据结构深度分析
2.1 aclshmemi_rdma_info — 全局 RDMA 信息表
文件: shmemi_device_rdma.h:38-45
struct aclshmemi_rdma_info {
uint32_t qp_num; // 每个连接的 QP 数量 (支持多 QP 并行)
uint64_t sq_ptr; // SQ 上下文数组地址 [PE_NUM][qp_num]
uint64_t rq_ptr; // RQ 上下文数组地址 [PE_NUM][qp_num]
uint62_t scq_ptr; // Send CQ 上下文数组地址 [PE_NUM][qp_num]
uint64_t rcq_ptr; // Recv CQ 上下文数组地址 [PE_NUM][qp_num]
uint64_t mem_ptr; // Memory Region 信息数组地址 [MAX_PE_NUM]
};
设计要点:
- 这是 device 端唯一的全局入口,通过 aclshmemi_qp_info_fetch() 从 HBM 固定地址读取
- 数组维度 [PE_NUM][qp_num] 支持多 rank 多 QP 场景
- mem_ptr 存储每个 rank 的 MR (Memory Region) 信息,包含 lkey/rkey 用于 RDMA 访问权限控制
2.2 aclshmemi_rdma_sq_ctx — Send Queue 上下文
文件: shmemi_device_rdma.h:54-66
struct aclshmemi_rdma_sq_ctx {
uint32_t wqn; // Work Queue Number (硬件 QP 编号)
uint64_t buf_addr; // SQ ring buffer 起始地址 (HBM)
uint32_t wqe_size; // 每个 WQE 的大小 (XSCALE=128B)
uint32_t depth; // ring buffer 深度 (WQE 数量)
uint64_t head_addr; // Producer Index 地址 (HBM, 软件写)
uint64_t tail_addr; // Consumer Index 地址 (HBM, 硬件写)
aclshmemi_rdma_db_mode_t db_mode; // HW_DB(0) / SW_DB(1)
uint64_t db_addr; // Doorbell register 地址
uint32_t sl; // Service Level (RDMA 优先级)
uint64_t amo_addr; // 原子操作的本地结果缓冲区地址
uint32_t amo_lkey; // 原子操作缓冲区的 local key
};
关键字段解读:
- head_addr / tail_addr: 分离的 producer/consumer index,支持硬件异步消费。head_addr 由软件通过 aclshmemi_roce_ring_sq_doorbell 更新,tail_addr 由硬件通过 CQE 完成事件推进。
- db_mode: 区分硬件 doorbell (HW_DB) 和软件 doorbell (SW_DB),适应不同硬件版本。
- amo_addr / amo_lkey: 原子操作的专用缓冲区。RDMA atomic fetch 操作的结果会写回此地址,软件通过 dcci_cachelines 刷新后读取。
2.3 aclshmemi_rdma_cq_ctx — Completion Queue 上下文
文件: shmemi_device_rdma.h:68-77
struct aclshmemi_rdma_cq_ctx {
uint32_t cqn; // Completion Queue Number
uint64_t buf_addr; // CQ ring buffer 起始地址
uint32_t cqe_size; // 每个 CQE 的大小 (XSCALE=64B, 含 32B padding)
uint32_t depth; // ring buffer 深度
uint64_t head_addr; // Producer Index (硬件写)
uint64_t tail_addr; // Consumer Index (软件写)
aclshmemi_rdma_db_mode_t db_mode;
uint64_t db_addr; // CQ Doorbell register 地址
};
与 SQ 的对称性: CQ 的 head_addr 由硬件更新(新 CQE 写入时),tail_addr 由软件更新(aclshmemi_roce_ring_cq_doorbell 在 poll 完成后更新)。
2.4 aclshmemi_rdma_send_wr — 发送工作请求
文件: shmemi_device_rdma.h:121-141
struct aclshmemi_rdma_send_wr {
// ---- Operation ----
uint32_t send_flags; // 发送标志位掩码
uint32_t imm_data; // immediate data (带内通知数据)
// ---- Remote side ----
__gm__ uint8_t *remote_addr; // 远端 HBM 目标地址
uint32_t rkey; // 远端 Memory Region key
// ---- Local side ----
__gm__ uint8_t *local_addr; // 本地 HBM 源地址
uint64_t message_len; // 消息长度 (Bytes)
uint32_t lkey; // 本地 Memory Region key
// ---- Multi-SGE ----
uint32_t num_sge; // SGE 数量 (scatter-gather)
aclshmemi_rdma_sge *sg_list; // SGE 列表
// ---- Atomic ----
aclshmemi_rdma_atomic_params atomic; // 原子操作参数
};
设计要点:
- __gm__ 修饰符标记全局内存指针(HBM 地址空间)
- num_sge / sg_list 支持 scatter-gather 传输(当前 XSCALE 后端主要使用单 SGE)
- atomic 联合体支持多种原子操作参数格式
2.5 aclshmemi_rdma_atomic_params — 原子操作参数联合体
文件: shmemi_device_rdma.h:85-119
struct aclshmemi_rdma_atomic_params {
union {
// 标准 8B FA/CAS (unmasked)
struct { uint64_t compare_add; uint64_t swap; };
// 通用 masked 格式 (XSCALE 使用)
struct {
uint64_t swap_add_data;
uint64_t compare_data;
uint64_t swap_add_mask;
uint64_t compare_mask;
} masked_common;
// 64-bit masked CAS (双 64-bit 段: swap+compare, mask_swap+mask_compare)
struct { uint64_t swap_data; uint64_t compare_data;
uint64_t swap_mask; uint64_t compare_mask; } masked_cas64;
// 64-bit masked FA (add_data + field_boundary)
struct { uint64_t add_data; uint64_t field_boundary; } masked_fa64;
// 32-bit masked CAS
struct { uint32_t swap_data; uint32_t compare_data;
uint32_t swap_mask; uint32_t compare_mask; } masked_cas32;
// 32-bit masked FA
struct { uint32_t add_data; uint32_t field_boundary;
uint64_t _reserved; } masked_fa32;
};
};
设计哲学: 使用 union 实现参数格式的零开销复用。masked_common 是 XSCALE 后端实际使用的格式,其 4 个 64-bit 字段直接映射到 WQE 的 atomic segment。
2.6 aclshmemi_rdma_mem_info — Memory Region 信息
文件: shmemi_device_rdma.h:47-52
struct aclshmemi_rdma_mem_info {
uint64_t size; // MR 大小
uint64_t addr; // MR 起始地址
uint32_t lkey; // Local Key (本地访问权限)
uint32_t rkey; // Remote Key (远端 RDMA 访问权限)
};
3. API 调用链
3.1 粗粒度 → 细粒度调用层次
┌─────────────────────────────────────────────────────────────┐
│ 用户代码 / 算子 kernel │
│ aclshmemx_roce_put_nbi(dst, src, buf, len, pe) │
│ aclshmemx_roce_get_nbi(dst, src, buf, len, pe) │
│ aclshmemx_roce_quiet(pe, buf, sync_id) │
│ aclshmemx_roce_atomic_fetch_add(dst, val, pe) │
└────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ 后端无关转发层 (shmem_device_rdma.hpp) │
│ aclshmemi_roce_write<T, ACLSHMEMI_K_RDMA_BACKEND>(...) │
│ → 根据 ACLSHMEMI_K_RDMA_BACKEND 宏选择 XSCALE 或 IN_DIE │
└────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ 细粒度原语层 (rdma_device_backend_base.h/.hpp) │
│ aclshmemi_roce_post_send<B, OP_RDMA_WRITE>(wr, pe, ...) │
│ → aclshmemi_roce_fill_wqe<B, OP_CODE>(wr, sq, addr, hd) │
│ → aclshmemi_roce_ring_sq_doorbell<B>(sq, head, ...) │
│ aclshmemi_roce_poll_cq<B>(pe, qp_idx, target, ...) │
│ aclshmemi_roce_ring_cq_doorbell<B>(pe, qp, tail, ...) │
└────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ 后端特化层 (rdma_device_backend_xscale.hpp) │
│ aclshmemi_xscale_post_send_read_write<OP_CODE>(...) │
│ → aclshmemi_roce_xscale_fill_wqe_ctrl_seg(...) │
│ → aclshmemi_roce_xscale_fill_wqe_data_seg(...) │
│ → aclshmemi_roce_write_ub_to_gm_with_sync(...) │
│ aclshmemi_rdma_xscdv_fill_wqe_fa<T, IS_MASKED>(...) │
│ aclshmemi_rdma_xscdv_fill_wqe_cas<T, IS_MASKED>(...) │
└─────────────────────────────────────────────────────────────┘
3.2 RDMA WRITE 完整调用链
// Step 1: 用户调用
aclshmemx_roce_put_nbi<uint32_t>(dst, src, buf, 256, pe);
// Step 2: 地址转换 + 转发
// shmem_device_rdma.hpp:183-197
aclshmem_ptr(dst, pe); // 本地对称地址 → 远端对称地址
aclshmemi_roce_write<uint32_t, XSCALE>(ptr, src, pe, 0, 1024, ...);
// Step 3: 填充 wr + post_send
// rdma_device_backend_base.h:137-149
wr.remote_addr = ptr;
wr.local_addr = src;
wr.message_len = 1024;
aclshmemi_roce_post_send<XSCALE, OP_RDMA_WRITE>(wr, pe, 0, ...);
// Step 4: XSCALE 后端 post_send
// rdma_device_backend_xscale.hpp:609-667
aclshmemi_xscale_post_send_read_write<OP_CODE>(wr, pe, qp_idx, ...) {
// 4a. 读取 SQ head/tail
dcci_cachelines(sq_pi_addr, 8);
cur_head = *(uint32_t*)sq_pi_addr;
// 4b. SQ 满检查 (阈值=16)
if ((cur_head + 16) % depth == (depth + cur_tail) % depth)
aclshmemi_roce_poll_cq<XSCALE>(...);
// 4c. 获取 rkey/lkey
wr.rkey = remote_mem_info->rkey;
wr.lkey = local_mem_info->lkey;
// 4d. 填充 WQE
aclshmemi_roce_fill_wqe<XSCALE, OP_RDMA_WRITE>(wr, sq, wqe_addr, cur_head);
// 4e. 刷回 + doorbell
dcci_cachelines(wqe_addr, 128);
aclshmemi_roce_ring_sq_doorbell<XSCALE>(sq, cur_head+1, ...);
}
3.3 aclshmemx_roce_quiet 同步机制
// shmem_device_rdma.hpp:253-276
template <typename T>
void aclshmemx_roce_quiet(uint32_t pe, __ubuf__ T *buf, uint32_t sync_id) {
for (qp_idx = 0; qp_idx < qp_num; qp_idx++) {
sq_context = ...;
dcci_cachelines(sq_pi_addr, 8); // 刷新 head 缓存行
cur_head = *(uint32_t*)sq_pi_addr; // 读取当前 producer index
aclshmemi_roce_poll_cq<XSCALE>(pe, qp_idx, cur_head, ...);
// poll_cq 会阻塞直到 cur_tail == cur_head (所有 WQE 完成)
}
}
关键设计: quiet 遍历所有 QP,对每个 QP 读取当前 SQ head(即已 post 的 WQE 总数),然后 poll_cq 等待 CQ tail 追上该值。
4. 编译时多态 traits 系统
4.1 设计理念
CAM 使用 C++ 模板特化 实现零运行时开销的后端选择。核心思想:
在编译期通过
ACLSHMEMI_K_RDMA_BACKEND宏确定后端,所有aclshmemi_roce_*<B>()调用在编译期绑定到具体后端实现,运行时无 switch/if 分支。
4.2 后端选择宏
// shmem_device_rdma.hpp:25-29
#if defined(ACLSHMEMI_RDMA_K_BACKEND_XSCALE)
#define ACLSHMEMI_K_RDMA_BACKEND (aclshmemi_rdma_backend_t::XSCALE)
#else
#define ACLSHMEMI_K_RDMA_BACKEND (aclshmemi_rdma_backend_t::IN_DIE)
#endif
4.3 traits 系统架构
// rdma_device_backend_base.h:263-315 — 主模板 (未定义, 强制特化)
template <aclshmemi_rdma_backend_t B>
struct aclshmemi_backend_traits {
template <typename T, bool IS_MASKED>
struct atomic_op_traits {
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE uint32_t fill_wqe(...);
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE void post_send(...);
};
};
// rdma_device_backend_xscale.hpp:883-905 — XSCALE 特化
template <>
struct aclshmemi_backend_traits<aclshmemi_rdma_backend_t::XSCALE> {
template <typename T, bool IS_MASKED>
struct atomic_op_traits {
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE uint32_t fill_wqe(...) {
return aclshmemi_roce_xscale_fill_wqe_atomic<T, IS_MASKED, ATOMIC_OP_CODE>(...);
}
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE void post_send(...) {
aclshmemi_rdma_xscdv_post_send_atomic<T, IS_MASKED, ATOMIC_OP_CODE>(...);
}
};
};
4.4 编译期错误防护
// rdma_device_backend_base.h:318-323
template <aclshmemi_rdma_atomic_op_t>
struct aclshmemi_atomic_op_dependent_false : std::false_type {};
template <aclshmemi_rdma_backend_t>
struct aclshmemi_rdma_backend_dependent_false : std::false_type {};
当调用未特化的后端或未实现的原子操作时,static_assert 触发编译错误而非运行时崩溃。
4.5 零运行时开销验证
所有关键路径函数标记为 ACLSHMEM_DEVICE(展开为 __forceinline__ __device__),配合 if constexpr 实现编译期分支消除:
// shmem_device_rdma.hpp:75-81
if constexpr (ACLSHMEMI_K_RDMA_BACKEND == aclshmemi_rdma_backend_t::XSCALE) {
return aclshmemi_roce_atomic_fetch_and_add<T, IS_MASKED, XSCALE>(...);
} else {
ACLSHMEM_DEBUG_FUNC(aclshmemi_kernel_abort, "ROCE atomic add is only supported on XSCALE backend.\n");
return T(0);
}
编译后,XSCALE 后端下 else 分支被完全消除,无运行时判断。
5. XSCALE 后端 WQE/CQE 格式
5.1 WQE 整体结构 (XSCALE, 128 Bytes)
XSCALE 后端的 WQE 为 128 Bytes (ACLSHMEMI_XSCALE_SND_WQE_SIZE = 1 << 7 = 128),由以下部分组成:
┌──────────────────────────────────────────────────────┐
│ Control Segment (16 Bytes) │
│ ├── msg_opcode (8 bits) — 操作码 │
│ ├── with_imm (1 bit) — 是否带 immediate │
│ ├── csum_en (2 bits) — 校验和使能 │
│ ├── ds_data_num (5 bits) — Data Segment 数量 │
│ ├── wqe_id (16/20 bits) — WQE 标识 │
│ ├── msg_len (32 bits) — 消息长度 │
│ ├── opcode_data (32 bits) — 操作码扩展数据 │
│ ├── se (1 bit) — Solicitation Event │
│ ├── ce (1 bit) — Completion Enable │
│ ├── in_line (1 bit) — Inline 模式 │
│ ├── fence_mode (2 bits) — Fence 模式 │
│ └── mask (2 bits) — 掩码 │
├──────────────────────────────────────────────────────┤
│ Remote Data Segment (16 Bytes) │
│ ├── seg_len (31 bits) — 远端数据长度 │
│ ├── rkey (32 bits) — 远端 Memory Region Key │
│ └── va (64 bits) — 远端虚拟地址 │
├──────────────────────────────────────────────────────┤
│ Local Data Segment (16 Bytes) │
│ ├── seg_len (31 bits) — 本地数据长度 │
│ ├── lkey (32 bits) — 本地 Memory Region Key │
│ └── va (64 bits) — 本地虚拟地址 │
├──────────────────────────────────────────────────────┤
│ Atomic Segment (16 Bytes, 仅原子操作) │
│ ├── swap_add (64 bits) — 交换/加数据 │
│ └── compare (64 bits) — 比较数据 │
├──────────────────────────────────────────────────────┤
│ Padding (64 Bytes, 补齐到 128B) │
└──────────────────────────────────────────────────────┘
5.2 Control Segment 版本差异
XSCALE 支持两个 API 版本,通过 ACLSHMEMI_XSCALE_API_VERSION_VAR 宏选择:
// V1: wqe_id 占 16 bits (bits [0:15] of the last 32-bit field)
struct aclshmemi_xscdv_wqe_ctrl_seg_v1 {
uint8_t msg_opcode; // [0:7]
uint8_t with_imm : 1; // [8]
uint8_t csum_en : 2; // [9:10]
uint8_t ds_data_num : 5; // [11:15]
uint16_t wqe_id; // [16:31] — 16 bits
uint32_t msg_len; // [32:63]
uint32_t opcode_data; // [64:95]
uint32_t se : 1; // [96]
uint32_t ce : 1; // [97]
uint32_t in_line : 1; // [98]
uint32_t fence_mode : 2; // [99:100]
uint32_t mask : 2; // [101:102]
uint32_t rsv : 25; // [103:127]
};
// V2: wqe_id 扩展到 20 bits, 重新排列字段
struct aclshmemi_xscdv_wqe_ctrl_seg_v2 {
uint8_t msg_opcode;
uint8_t with_imm : 1;
uint8_t csum_en : 2;
uint8_t ds_data_num : 5;
uint16_t rsv1; // 保留
uint32_t msg_len;
uint32_t opcode_data;
uint32_t se : 1;
uint32_t ce : 1;
uint32_t in_line : 1;
uint32_t fence_mode : 2;
uint32_t mask : 2;
uint32_t rsv : 5;
uint32_t wqe_id : 20; // [27:47] — 20 bits (支持更多 WQE)
};
版本选择机制:
using aclshmemi_xscdv_wqe_ctrl_seg_t =
std::conditional_t<ACLSHMEMI_XSCALE_API_VERSION_VAR == 1,
aclshmemi_xscdv_wqe_ctrl_seg_v1,
aclshmemi_xscdv_wqe_ctrl_seg_v2>;
5.3 CQE 格式 (XSCALE Diamond, 64 Bytes)
struct aclshmemi_xscdv_cqe64_t {
aclshmemi_xscdv_diamond_cqe_t cqe; // 32 Bytes
uint8_t padding[32]; // 32 Bytes padding
};
// V2 CQE (32 Bytes):
struct aclshmemi_xscdv_diamond_cqe_v2 {
union {
struct {
uint32_t error_code : 8; // [0:7] 错误码
uint32_t qp_id : 15; // [8:22] 完成的 QP 编号
uint32_t rsv : 1; // [23]
uint32_t se : 1; // [24] Solicitation Event
uint32_t has_pph : 1; // [25] 是否有 PPH
uint32_t type : 1; // [26]
uint32_t with_imm : 1; // [27] 是否有 immediate
uint32_t csum_err : 4; // [28:31] 校验和错误
};
uint32_t flags_qp_id_err_code;
};
uint32_t imm_data; // immediate value
uint32_t msg_len; // 消息长度
uint32_t vni; // Virtual Network Identifier
uint32_t ts_l; // 时间戳低 32 位
uint32_t ts_h; // 时间戳高 16 位 (共 48 位)
union {
struct {
uint32_t msg_opcode : 8; // 对应 WQE 的 opcode
uint32_t rsv1 : 4;
uint32_t wqe_id : 20; // 对应 WQE 的 wqe_id
};
uint32_t wqe_id_rsv_opcode;
};
union {
struct {
uint32_t rsv2 : 31;
uint32_t owner : 1; // Owner bit: 软件解析权限标志
};
uint32_t owner_rsv;
};
};
5.4 Owner Bit 机制
XSCALE 使用 owner bit 实现硬件/软件同步,避免昂贵的 atomic 操作:
// rdma_device_backend_xscale.hpp:389-397
bool aclshmemi_roce_xscale_check_cqe_owner(__gm__ aclshmemi_xscdv_cqe64_t *cqe64,
uint32_t cur_tail, uint32_t depth) {
bool owner_bit = (cqe64->cqe.owner & ACLSHMEMI_XSC_CQE_OWNER_MASK);
bool expect_bit = ((cur_tail / depth) & 1); // 每翻转一轮 depth,owner bit 翻转
return owner_bit == expect_bit;
}
原理: 硬件每写完一轮 CQ(depth 个 CQE),自动翻转 owner bit。软件通过比较 owner_bit == (cur_tail / depth) & 1 判断 CQE 是否已写入。
5.5 Doorbell 格式
SQ Doorbell (8 Bytes):
union aclshmemi_xscdv_diamond_send_doorbell_v2 {
struct {
uint64_t next_pid : 21; // 下一个 WQE 的 producer index (已左移)
uint64_t qp_id : 16; // QP 编号
};
uint64_t raw;
};
CQ Doorbell (8 Bytes):
union aclshmemi_xscdv_diamond_cq_doorbell_t {
struct {
uint64_t cq_next_cid : 23; // 下一个待处理的 CQE ID
uint64_t cq_id : 16; // CQ 编号
uint64_t cq_sta : 2; // CQ 状态
};
uint64_t raw;
};
5.6 操作码映射
enum class aclshmemi_xscdv_msg_type_t : uint32_t {
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_WRITE = 1,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_READ = 2,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_CMP_AND_SWAP = 26,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_FETCH_AND_ADD = 27,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_8B_MSK_CMP_AND_SWAP = 31,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_8B_MSK_FETCH_AND_ADD = 32,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_4B_MSK_CMP_AND_SWAP = 33,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_4B_MSK_FETCH_AND_ADD = 34,
};
6. 原子操作机制
6.1 原子操作 API 全景
CAM 提供 16 种 原子操作 API,覆盖 fetch/set/compare-swap/inc/add/and/or/xor 及其组合:
| 类别 | API | 同步性 | 实现方式 |
|---|---|---|---|
| Fetch | atomic_fetch(src, pe) |
同步 | FA(add=0) + quiet + read |
| Set | atomic_set(dst, val, pe) |
异步 | CAS(swap=val, comp=0, mask_swap=MAX) |
| CAS | atomic_compare_swap(dst, cond, val, pe) |
同步 | CAS + quiet + read |
| Swap | atomic_swap(dst, val, pe) |
同步 | CAS(swap=val, comp=0) + quiet + read |
| Inc | atomic_inc(dst, pe) |
异步 | FA(add=1) |
| Add | atomic_add(dst, val, pe) |
异步 | FA(add=val) |
| And | atomic_and(dst, val, pe) |
异步 | CAS(swap=val, mask_swap=~val) |
| Or | atomic_or(dst, val, pe) |
异步 | CAS(swap=val, mask_swap=val) |
| Xor | atomic_xor(dst, val, pe) |
异步 | FA(add=val, boundary=MAX) |
| FetchInc | atomic_fetch_inc(dst, pe) |
同步 | FA(add=1) + quiet + read |
| FetchAdd | atomic_fetch_add(dst, val, pe) |
同步 | FA(add=val) + quiet + read |
| FetchAnd | atomic_fetch_and(dst, val, pe) |
同步 | CAS + quiet + read |
| FetchOr | atomic_fetch_or(dst, val, pe) |
同步 | CAS + quiet + read |
| FetchXor | atomic_fetch_xor(dst, val, pe) |
同步 | FA + quiet + read |
6.2 Masked vs Unmasked 变体
XSCALE 后端对原子操作区分 masked 和 unmasked 两种模式:
// 4B 操作强制使用 masked 模式
static_assert(!(TYPE_BYTES <= BYTES_32 && !IS_MASKED),
"XSCALE backend doesn't support fetch and add operation with type size of 4B without mask\n");
Masked 模式 允许对 64-bit 字中的特定位字段进行原子操作,通过 swap_add_mask / compare_mask 控制哪些位参与操作。
6.3 原子操作 WQE 构造
以 aclshmemi_rdma_xscdv_fill_wqe_fa 为例:
template <typename T, bool IS_MASKED>
uint32_t aclshmemi_rdma_xscdv_fill_wqe_fa(aclshmemi_rdma_send_wr &wr, ...) {
constexpr uint32_t DS_DATA_NUM = 3; // ctrl + data_seg + atomic_seg
constexpr uint32_t MSG_LEN = (TYPE_BYTES == BYTES_32 && IS_MASKED) ? 8 : 16;
constexpr uint32_t DATA_LEN = (TYPE_BYTES <= BYTES_32) ? BYTES_32 : BYTES_64;
// 根据类型选择操作码
constexpr aclshmemi_xscdv_msg_type_t OPCODE =
(TYPE_BYTES == BYTES_64 && IS_MASKED) ? ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_8B_MSK_FETCH_AND_ADD
: (TYPE_BYTES == BYTES_64 && !IS_MASKED) ? ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_FETCH_AND_ADD
: ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_4B_MSK_FETCH_AND_ADD;
// 填充 ctrl_seg (inline=1, 因为数据在 WQE 内)
cur_wqe_addr = aclshmemi_roce_xscale_fill_wqe_ctrl_seg(wqe_addr, DS_DATA_NUM, cur_head, OPCODE, MSG_LEN, 1);
// 填充 data_seg (remote + local)
cur_wqe_addr = aclshmemi_roce_xscale_fill_wqe_data_seg(cur_wqe_addr, wr.rkey, wr.remote_addr, wr.lkey, wr.local_addr, DATA_LEN);
// 填充 atomic_seg
return aclshmemi_rdma_xscdv_fill_fa_atomic_seg<T, IS_MASKED>(cur_wqe_addr, wr);
}
6.4 字节序转换
XSCALE 硬件要求 Big-Endian 格式的原子操作数据:
// rdma_device_backend_xscale.hpp:333-343
ACLSHMEM_DEVICE uint64_t aclshmemi_htobe64(uint64_t host_val) {
return ((host_val >> 0) & 0xFF) << 56 |
((host_val >> 8) & 0xFF) << 48 |
((host_val >> 16) & 0xFF) << 40 |
((host_val >> 24) & 0xFF) << 32 |
((host_val >> 32) & 0xFF) << 24 |
((host_val >> 40) & 0xFF) << 16 |
((host_val >> 48) & 0xFF) << 8 |
((host_val >> 56) & 0xFF) << 0;
}
读取 fetch 结果时需反向转换:
// shmem_device_rdma.hpp:293-297
if constexpr (sizeof(T) == 4 && ACLSHMEMI_K_RDMA_BACKEND == aclshmemi_rdma_backend_t::XSCALE) {
uint32_t fetch_bytes = *fetch_addr;
return (T)aclshmemi_htobe32(fetch_bytes); // 4B 结果需要字节序转换
}
6.5 原子操作结果读取
// shmem_device_rdma.hpp:278-301
uint64_t aclshmemi_roce_get_atomic_fetch_addr(uint32_t pe, uint32_t qp_idx) {
// 从 sq_ctx->amo_addr 获取原子操作结果缓冲区地址
return sq_context->amo_addr;
}
template <typename T>
T aclshmemi_roce_get_atomic_fetch_data(uint32_t pe, uint32_t qp_idx) {
auto amo_addr = aclshmemi_roce_get_atomic_fetch_addr(pe, qp_idx);
dcci_cachelines((__gm__ uint8_t *)amo_addr, sizeof(T)); // 刷新缓存行
__gm__ T *fetch_addr = reinterpret_cast<__gm__ T *>(amo_addr);
return *fetch_addr; // 读取 fetch 前的原始值
}
7. Host 端传输管理
7.1 RdmaTransportManagerV2 生命周期
OpenDevice() → RegisterMemoryRegion() × N → Prepare() → Connect() → [运行期] → CloseDevice()
7.2 OpenDevice 流程
// device_rdma_transport_manager_v2.cpp:44-101
Result OpenDevice(const TransportOptions& options) {
// 1. 获取设备 ID
DlAclApi::AclrtGetDevice(&userId);
DlAclApi::AclrtGetPhyDevIdByLogicDevId(userId, &phyId_);
// 2. 解析 RDMA 网络地址
TopoReader::ParseRdmaNetAddr(phyId_, deviceIp_);
devicePort_ = RDMA_PORT_PREFIX(60032) + (rankId_ % MAX_RANKS_PER_NIC(16));
nicInfo_ = GenerateDeviceNic(deviceIp_, devicePort_);
// 3. 创建 RoCE 端点
CreateEndpoint(); // → HcommEndpointCreate(&endpointDesc, &endpointHandle_)
// 4. 预留 device 内存 (qpInfo + atomic buffer)
ReserveRdmaInfoSpace();
}
7.3 Connect 流程
// device_rdma_transport_manager_v2.cpp:280-307
Result Connect() {
// 1. 验证单 NIC rank 数不超过 16
ValidateRanksPerNic();
// 2. 准备 channel 描述符
uint32_t channelNum = rankCount_ - 1;
PrepareChannelDescs(channelDescs, roceTc, roceSl);
// 对每个 remote rank:
// - 设置 protocol = COMM_PROTOCOL_ROCE
// - 设置 IP 地址 (IPv4/IPv6)
// - 设置 tc/sl (Traffic Class / Service Level)
// - 设置 port = 60032 + serverRank * 16 + clientRank
// - 设置 role = SERVER (rankId < remoteRank) / CLIENT
// 3. 创建 channel 并填充 QP 信息
CreateChannelsAndFillInfo(channelDescs, channelNum);
}
7.4 FillRdmaInfo — 核心数据提取
// device_rdma_transport_manager_v2.cpp:507-531
Result FillRdmaInfo() {
// 1. 分配临时 buffer
std::vector<uint8_t> qpInfoBuffer(qpInfoSize_);
auto copyInfo = (AiQpRMAQueueInfo*)qpInfoBuffer.data();
// 2. 设置指针布局 (sq/rq/scq/rcq/mr 紧跟在 AiQpRMAQueueInfo 之后)
FillQpPreSettingCopyInfo(copyInfo);
// 3. 从 ChannelEntity 读取 SQ/CQ/MR 信息
GetRdmaInfoFromChannelEntity(copyInfo, channelPtrs_);
// ├── ReadLocalBufferInfo() → 读取本地 MR 的 lkey/rkey
// └── ReadRemoteBufferInfo() → 遍历每个 remote rank:
// ├── ReadSingleRemoteRank() → D2H 拷贝 ChannelEntity
// │ ├── 读取 remote MR 的 lkey/rkey
// │ └── FillSqCqAtomicInfo() → 读取 SQ/CQ 上下文
// │ ├── CopyAiWQInfo() → 提取 wqn/bufAddr/wqeSize/depth/head/tail/dbAddr
// │ ├── CopyAiCQInfo() → 提取 cqn/bufAddr/cqeSize/depth/head/tail/dbAddr
// │ └── 设置 atomicAddr/atomicLkey
// └── 设置 remote MR 的 size/addr
// 4. 拷贝到 device HBM
DlAclApi::AclrtMemcpy(qpInfo_, qpInfoSize_, copyInfo, qpInfoSize_, ACL_MEMCPY_HOST_TO_DEVICE);
}
7.5 内存注册策略
// device_rdma_transport_manager_v2.cpp:191-224
Result RegisterMemoryRegion(const TransportMemoryRegion& mr) {
CommMem commMem{};
// 区分 HBM 和 Host DDR
if ((mr.flags & REG_MR_FLAG_HBM) || IsVirtualAddressNpu(mr.addr))
commMem.type = COMM_MEM_TYPE_DEVICE; // HBM
else
commMem.type = COMM_MEM_TYPE_HOST; // DDR
// 注册到 HCCP 协议栈
HcommMemReg(endpointHandle_, "HcclBuffer", &commMem, &memHandle);
registeredMRs_.emplace(mr.addr, entry);
}
7.6 原子操作共享内存
// device_rdma_transport_manager_v2.cpp:703-735
bool ReserveRdmaInfoSpace() noexcept {
// qpInfo 大小: sizeof(AiQpRMAQueueInfo) + rankCount * [2*(sizeof(AiQpRMAWQ) + sizeof(AiQpRMACQ)) + sizeof(RdmaMemRegionInfo)]
qpInfoSize_ = sizeof(AiQpRMAQueueInfo) + oneQpSize * rankCount_;
DlAclApi::AclrtMalloc(&ptr, qpInfoSize_, 0);
qpInfo_ = (AiQpRMAQueueInfo*)ptr;
// 原子操作共享内存: 128 * 8B * rankCount, 512B 对齐
uint32_t atomicSize = ATOMIC_MAX_NUM(128) * sizeof(uint64_t) * rankCount;
atomicSize = ALIGN_UP(atomicSize, MEMORY_ALIGNMENT(512));
DlAclApi::AclrtMalloc(&atomicSharedMemory_, atomicSize, 0);
}
设计要点: 原子操作共享内存用于存储 fetch-before-write 结果。每个 rank 占用 128 × 8B = 1KB,支持最大 128 个并发原子操作。
7.7 底层协议结构
HCCP 协议 (dl_hccp_def.h):
- HccpMrInfo: Memory Region 信息 (addr/size/access/lkey/rkey)
- HccpQpExtAttrs: QP 扩展属性 (qpMode/cqAttr/qp_attr/version/mem_align)
- ai_data_plane_wq / ai_data_plane_cq: 硬件队列描述 (wqn/cqn/buf_addr/wqebb_size/depth/head/tail/db_reg)
- AiQpRMAWQ / AiQpRMACQ: 运行时 SQ/CQ 信息 (含 atomicAddr/atomicLkey)
Hcomm 协议 (dl_hcomm_def.h):
- ChannelEntity: 通道实体 (engine/protocol/bufferNum/notifyNum/sqNum/cqNum)
- SqContext / CqContext: 队列上下文 (roceSq/roceCq 包含 qpn/cqn/sqVa/cqVa/dbVa 等)
- RegedBufferEntity: 注册缓冲区 (addr/size/protectionInfo)
- ProtectionInfo: 保护信息 (type=PROTECTION_TYPE_ROCE, lkey/rkey)
8. 编译注入机制
8.1 SHMEM_HOME_PATH 环境变量
# ascendc_impl_build.py:167-178
SHMEM_HOME_PATH = os.environ.get('SHMEM_HOME_PATH')
if SHMEM_HOME_PATH is None:
SHMEM_HOME_PATH = os.path.realpath("/usr/local/Ascend/shmem/latest")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/device")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/host_device")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/internal")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/src")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/src/device")
options.append("-I" + SHMEM_HOME_PATH + "/memfabric_hybrid/include/smem/device")
8.2 编译期后端选择
# compile_ascend_proj.sh:179-192
KERNEL_COMPILE_OPTS=""
if [ -n "${SHMEM_HOME_PATH}" ]; then
KERNEL_COMPILE_OPTS="-I${SHMEM_HOME_PATH}/shmem/include -I${SHMEM_HOME_PATH}/shmem/src/device"
fi
# 通过 sed 注入到 build.make 的 --compile-options 中
find build_out/op_kernel/CMakeFiles -name "build.make" -exec sed -i \
"s|--compile-options=\"\"|--compile-options=\\\"${KERNEL_COMPILE_OPTS}\\\"|g" {} +
8.3 算子选择过滤
# select_ops.py:55-60, 116-121
def requires_shmem(registry, op):
return bool(get_meta(registry, op).get("requires_shmem", False))
# 在 resolve() 中:
if not shmem:
dropped = [o for o in candidates if requires_shmem(registry, o)]
for o in dropped:
sys.stderr.write(f"NOTE: dropping {o} (requires SHMEM, which is not installed)\n")
candidates = [o for o in candidates if not requires_shmem(registry, o)]
8.4 当前算子注册表状态
// operator_registry.json
"operator_meta": {
"dispatch_layout": { "requires_shmem": false },
"fused_deep_moe": { "requires_shmem": false, "family": "fused_deep_moe", "variant": "base" },
"moe_combine_normal": { "requires_shmem": false },
"moe_dispatch_normal":{ "requires_shmem": false },
"notify_dispatch": { "requires_shmem": false }
}
当前所有算子 requires_shmem: false, 即默认使用 HCCL Window 模式。SHMEM 模式基础设施已完整就绪,但尚未在算子层启用。
9. 与 HCCL Window 模式的本质区别
9.1 架构层面对比
| 维度 | SHMEM 模式 | HCCL Window 模式 |
|---|---|---|
| 通信路径 | AIV 核心直接操作 RoCE 硬件 WQE/CQE | 通过 HCCL Window 机制 (winContext_->localWindowsIn) |
| 运行时开销 | 单次 kernel 内完成 RDMA 操作,无 HCCL 调度 | 需经过 HCCL 调度层、Window 资源管理 |
| 延迟 | 微秒级 (直接 doorbell 振铃) | 数十微秒级 (HCCL 软件栈开销) |
| 灵活性 | 支持任意地址、任意长度的 P2P 传输 | 受限于 Window 预分配区域 |
| 原子操作 | 16 种 RDMA 原子操作 (CAS/FA/AND/OR/XOR) | 依赖 HCCL 集合通信原语 |
| 计算-通信重叠 | 同一 kernel 内交替执行 GEMM 和 RDMA | 需多 kernel 协作,HBM 读写开销 |
| 后端选择 | 编译期 traits 系统 (IN_DIE/XSCALE) | HCCL 内部自动选择 |
| 内存注册 | 显式 MR 注册 (HcommMemReg) | HCCL 隐式管理 |
| QP 管理 | 显式 SQ/CQ 上下文管理 | HCCL 抽象 |
9.2 代码层面对比
SHMEM 模式 (直接 RDMA):
// 算子 kernel 内部直接发起 RDMA
aclshmemx_roce_put_nbi(remote_addr, local_addr, ub_buf, len, remote_pe);
// ↑ 直接构造 WQE → doorbell → 硬件执行
HCCL Window 模式 (当前默认):
// 通过 HCCL Window 机制
winContext_->localWindowsIn[peer_rank] // 预分配的 Window 区域
HcclRankRelationResV2 // HCCL 管理的 rank 关系
// ↑ 需经过 HCCL 调度、Window 资源同步
9.3 性能影响分析
SHMEM 模式的核心优势在于 消除 HCCL 软件栈开销:
- 零调度延迟: 直接 doorbell 振铃 (~100ns) vs HCCL 调度 (~10-50μs)
- Kernel 内融合: FusedDeepMoE 可在同一 kernel 内交替执行 GEMM 和 RDMA,避免多次 kernel launch 和 HBM 读写
- 细粒度 P2P: 支持 token 级别的动态路由,无需预定义通信模式
- 硬件原子操作: 利用 RDMA 硬件原子操作实现无锁同步
9.4 启用路径
要将算子从 HCCL Window 模式切换到 SHMEM 模式:
- 修改 operator_registry.json: 将目标算子的
requires_shmem设为true - 确保 SHMEM 安装: 设置
SHMEM_HOME_PATH环境变量指向 SHMEM 安装目录 - 算子代码改造: 在 kernel 中调用
aclshmemx_roce_*API 替代 HCCL Window 操作 - 编译:
select_ops.py会自动包含该算子,编译注入 SHMEM include 路径
10. 关键设计洞察
10.1 用户态 RDMA 的极致优化
CAM 的 SHMEM 模式是 用户态 RDMA 在 AI 芯片上的极致实现:
- AIV 核心直接驱动: 不同于传统 RDMA 由 CPU 或专用 NIC 驱动,CAM 让 AIV 向量核心直接构造 WQE 和振铃 doorbell
- Cache-line 精确刷新: 使用
dcci_cachelines()精确刷新 64B 缓存行,确保 HBM 数据对硬件可见 - Inline 模式: 原子操作使用 inline 模式 (数据在 WQE 内),避免额外的 DMA 读取
10.2 编译时多态的工程价值
traits 系统实现了 零成本抽象:
- 无虚函数表 (vtable) 开销
- 无运行时 switch/if 分支
- 未使用的后端代码在编译期完全消除
if constexpr确保编译期分支消除
10.3 版本兼容性的优雅处理
XSCALE 后端的 V1/V2 API 版本通过 std::conditional_t 实现编译期选择:
using aclshmemi_xscdv_wqe_ctrl_seg_t =
std::conditional_t<ACLSHMEMI_XSCALE_API_VERSION_VAR == 1,
aclshmemi_xscdv_wqe_ctrl_seg_v1,
aclshmemi_xscdv_wqe_ctrl_seg_v2>;
这允许同一份代码支持不同硬件世代的 WQE 格式差异。
10.4 同步机制的精心设计
- SQ head/tail 分离: 软件写 head (producer), 硬件写 tail (consumer via CQE)
- CQ owner bit: 硬件/软件同步无 atomic 操作,每轮 depth 翻转一次
- SQ 满检查: 阈值 16 的预检查避免死锁,同时最大化 SQ 利用率
- 超时保护: 5 分钟超时防止硬件异常导致无限等待
10.5 原子操作的统一抽象
通过 aclshmemi_rdma_atomic_params union + IS_MASKED 模板参数,实现了:
- 8B/4B 数据类型统一处理
- masked/unmasked 模式编译期选择
- 多种原子操作 (FA/CAS/AND/OR/XOR) 的统一 WQE 构造路径
11. 总结
CAM 的 SHMEM 模式是一个 完整的用户态 RDMA 子系统,其核心优势在于:
- 硬件直通: AIV 核心直接操作 RoCE 硬件,消除 HCCL 软件栈开销
- 编译时多态: traits 系统实现零运行时开销的后端选择
- 原子操作丰富: 16 种原子操作变体,支持 8B/4B masked/unmasked 模式
- 版本兼容: V1/V2 API 版本编译期选择,支持多硬件世代
- 工程完备: 超时保护、错误处理、日志调试机制齐全
当前所有算子默认使用 HCCL Window 模式 (requires_shmem: false),但 SHMEM 模式基础设施已完整就绪,是 CAM 未来性能突破的关键路径。