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 软件栈开销:

  1. 零调度延迟: 直接 doorbell 振铃 (~100ns) vs HCCL 调度 (~10-50μs)
  2. Kernel 内融合: FusedDeepMoE 可在同一 kernel 内交替执行 GEMM 和 RDMA,避免多次 kernel launch 和 HBM 读写
  3. 细粒度 P2P: 支持 token 级别的动态路由,无需预定义通信模式
  4. 硬件原子操作: 利用 RDMA 硬件原子操作实现无锁同步

9.4 启用路径

要将算子从 HCCL Window 模式切换到 SHMEM 模式:

  1. 修改 operator_registry.json: 将目标算子的 requires_shmem 设为 true
  2. 确保 SHMEM 安装: 设置 SHMEM_HOME_PATH 环境变量指向 SHMEM 安装目录
  3. 算子代码改造: 在 kernel 中调用 aclshmemx_roce_* API 替代 HCCL Window 操作
  4. 编译: 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 子系统,其核心优势在于:

  1. 硬件直通: AIV 核心直接操作 RoCE 硬件,消除 HCCL 软件栈开销
  2. 编译时多态: traits 系统实现零运行时开销的后端选择
  3. 原子操作丰富: 16 种原子操作变体,支持 8B/4B masked/unmasked 模式
  4. 版本兼容: V1/V2 API 版本编译期选择,支持多硬件世代
  5. 工程完备: 超时保护、错误处理、日志调试机制齐全

当前所有算子默认使用 HCCL Window 模式 (requires_shmem: false),但 SHMEM 模式基础设施已完整就绪,是 CAM 未来性能突破的关键路径。

← 测试体系与回归测试 📑 索引 非 SHMEM — HCCL Window →