Ascend C Reg API 教程与高性能实现

第一部分:基础概念

1.1 什么是 Reg API

Reg API(寄存器级向量计算 API)是 Ascend C 在 CANN v9.0.0 引入的面向 SIMD Vector 寄存器的编程接口,对应硬件 SIMD Vector Core 内部的 VF Reg(Vector Function Register)文件。所有 Reg API 都位于 AscendC::Reg 命名空间,使用 __simd_callee__ inline 修饰。

与传统基础 API(AscendC::AddAscendC::Mul 等 LocalTensor 接口)相比:

维度 基础 API(MemBase) Reg API(RegBase)
数据载体 LocalTensor<T>(UB 内存) RegTensor<T>(VF 寄存器)
单次处理粒度 任意长度(硬件内部 tiling) 一个 VL(Vector Length,950PR=256B)
软件显式分块 不需要 需要:循环每次处理 VL/sizeof(T) 个元素
数据搬运 DataCopy(MTE2/MTE3) LoadAlign/StoreAlign(直接 UB↔VF Reg)
控制流 __aicore__ 内联 必须__simd_vf__ 函数 + asc_vf_call 调用
Mask 控制 通过 count 参数自动 显式 MaskReg 寄存器
寄存器复用 硬件决定 软件可控:dst=src 时直接复用
性能上限 受 MTE 流水限制 减少 UB 往返,提升指令并发
编程复杂度 简单 较高,需理解寄存器/调用层级

1.2 何时使用 Reg API

推荐使用 Reg API 的场景

  1. 多步融合计算:同一个数据要经过多个算子(如 Cast → Mul → Add → Cast),中间结果可留在寄存器中避免反复 Load/Store UB
  2. 计算密集型算子:如 Exp/Ln/Sqrt/Div 等指令周期较长,减少 UB 往返收益大
  3. RegTrait 双宽场景:64 位整数运算(int64_t)需要 RegTraitNumTwo 把 2 个 VL 拼成 2×VL 的逻辑寄存器
  4. 极致性能调优:基础 API 已达瓶颈,需要软件掌控指令调度
  5. 非连续数据访问Gather/ScatterSqueeze、Block Strided Load 等 MemBase 难以高效表达的模式

不推荐使用 Reg API 的场景

  1. 简单单步算子(如纯 Add/Sub),基础 API 已足够快
  2. 团队不熟悉寄存器编程模型,调试成本高
  3. 数据搬运本身就是瓶颈(mte2_ratio > 90%)的算子,Reg API 收益有限

来自实测数据:Ascend 950 系列 Add 算子 RegBase 相对 MemBase 的主要收益来自"减少冗余 Load/Store、寄存器复用、提升无依赖并发指令比例";但当前 Add 样例核心计算基本是单条 Add,计算链短、可融合/双发空间有限,难以充分释放 RegBase 优势。这说明 Reg API 在多步融合场景才能体现价值。

1.3 存储层级与编程模型

                 ┌─────────────────────────────────────────────┐
   GM (HBM)      │  ~1.6-1.8 TB/s                              │
   (Global)      └──────────────┬──────────────────────────────┘
                                 │ DataCopy (MTE2/MTE3)
                                 ▼
                 ┌─────────────────────────────────────────────┐
   UB            │  Unified Buffer, 950PR=256KB                │
   (LocalMem)    │  bank=16group×3bank×4KB, vec一拍访问256B    │
   __ubuf__      └──────────────┬──────────────────────────────┘
                                 │ LoadAlign / StoreAlign (vld/vst)
                                 ▼
                 ┌─────────────────────────────────────────────┐
   VF Reg        │  Vector Function Register file              │
   RegTensor<T>  │  单个寄存器宽度 = VL = 256B                 │
                 │  RegTraitNumOne: 1 reg; NumTwo: 2 reg (int64)│
                 │  MaskReg (vector_bool): 32B (256 bits)      │
                 │  AddrReg (vector_address): 地址偏移寄存器   │
                 │  UnalignRegForLoad/Store (vector_align):    │
                 │     非对齐访存辅助寄存器                    │
                 └──────────────┬──────────────────────────────┘
                                │
                                ▼
                       AscendC::Reg::Add / Mul / Cast / ...
                       (VALU 算术逻辑单元执行)

1.4 调用层级

Reg API 的调用链有严格层级,违反会编译报错:

__global__ __aicore__ kernel (Host 入口)
   │
   ├─ 可调用: __aicore__ 函数、constexpr aicore
   │
   ▼
__aicore__ 函数(算子主流程)
   │
   ├─ 可调用: __simd_vf__ 函数 (通过 asc_vf_call)、__aicore__ 函数
   │
   ▼
__simd_vf__ 函数(SIMD Vector Function,Reg API 的唯一执行环境)
   │
   ├─ 可调用: __simd_callee__ 函数(所有 AscendC::Reg::* 都是 __simd_callee__)
   │
   ▼
__simd_callee__ 函数(叶子节点,仅能调用其他 __simd_callee__)

关键规则

  • Reg API(AscendC::Reg::*)全部标注 __simd_callee__ inline
  • __simd_vf__ 标记的函数只能通过 asc_vf_call<VFFunc>(args...) 调用
  • __simd_callee__ 只能被 __simd_vf__、其他 __simd_callee__constexpr aicore 调用
  • Reg API 不能在 __aicore__ 函数中直接调用,必须经 asc_vf_call 进入 VF 上下文

合法调用模板

// 1. 定义 SIMD VF 函数
template <typename T>
__simd_vf__ inline void AddVF(__ubuf__ T* dst, __ubuf__ T* src0, __ubuf__ T* src1,
                              uint16_t repeatTimes, uint32_t oneRepeatSize)
{
    AscendC::Reg::RegTensor<T> r0, r1, r2;
    AscendC::Reg::MaskReg mask = AscendC::Reg::CreateMask<T>();
    for (uint16_t i = 0; i < repeatTimes; ++i) {
        AscendC::Reg::LoadAlign(r0, src0 + i * oneRepeatSize);
        AscendC::Reg::LoadAlign(r1, src1 + i * oneRepeatSize);
        AscendC::Reg::Add(r2, r0, r1, mask);          // __simd_callee__ 调用
        AscendC::Reg::StoreAlign(dst + i * oneRepeatSize, r2, mask);
    }
}

// 2. 在 __aicore__ 中调用
__aicore__ inline void Process()
{
    // ... DataCopy / SetFlag / WaitFlag 等 MTE 操作 ...
    asc_vf_call<AddVF<T>>(dstAddr, src0Addr, src1Addr, repeatTimes, oneRepeatSize);
    // ... 后续 MTE3 ...
}

第二部分:Reg API 数据类型与寄存器

2.1 RegTensor —— 向量数据寄存器

RegTensor<T> 是 Reg API 的核心数据载体,定义在 kernel_reg_compute_struct_intf.h

struct RegTrait { int REG_NUM = 1; };
constexpr RegTrait RegTraitNumOne = {1};   // 1 个 VL 宽寄存器(默认)
constexpr RegTrait RegTraitNumTwo = {2};   // 2 个 VL 宽寄存器(int64 等需要)

template <typename T, const RegTrait& regTrait = RegTraitNumOne>
struct RegTensor {
    using ActualT = T;
    using RegType = typename TypeGet<T>::T;
    RegType reg[regTrait.REG_NUM];   // 物理寄存器数组
    static constexpr int REG_NUM = regTrait.REG_NUM;
};

容量RegTensor<T> 一次可装载 VL / sizeof(T) 个元素:

  • float/half/int32_t:VL=256B → 64 个 float / 128 个 half / 64 个 int32
  • int64_t:需用 RegTraitNumTwo 拼接 2 个 VL → 64 个 int64(逻辑上 512B 宽寄存器)

使用要点

  • RegTensor<T> 应声明在 __simd_vf__ 函数内部,作为栈上变量
  • 寄存器生命周期由硬件 rename 表管理,软件声明只是逻辑寄存器
  • dstRegsrcReg 可以是同一个变量(如 Add(r, r, r, mask)),硬件会正确处理 RAW 依赖

2.2 MaskReg —— 谓词掩码寄存器

MaskReg = vector_bool,宽度为 VL/8(950PR 都是 32B = 256 bit)。每一位对应 RegTensor 中的一个 lane(最小粒度为字节级 mask 的部分场景除外)。

创建 Mask 的两种方式

// 1. CreateMask: 编译期固定模式
AscendC::Reg::MaskReg maskAll  = AscendC::Reg::CreateMask<T, AscendC::Reg::MaskPattern::ALL>();   // 全 1
AscendC::Reg::MaskReg maskVL1  = AscendC::Reg::CreateMask<T, AscendC::Reg::MaskPattern::VL1>();    // 仅最低 1 lane
AscendC::Reg::MaskReg maskVL64 = AscendC::Reg::CreateMask<T, AscendC::Reg::MaskPattern::VL64>();   // 低 64 lane
AscendC::Reg::MaskReg maskM3   = AscendC::Reg::CreateMask<T, AscendC::Reg::MaskPattern::M3>();     // 每 4 lane 取 3
AscendC::Reg::MaskReg maskH    = AscendC::Reg::CreateMask<T, AscendC::Reg::MaskPattern::H>();      // 高半
AscendC::Reg::MaskReg maskQ    = AscendC::Reg::CreateMask<T, AscendC::Reg::MaskPattern::Q>();      // 1/4

// 2. UpdateMask: 运行期根据剩余 count 生成
uint32_t count = 200;   // 需要处理的元素数
AscendC::Reg::MaskReg mask = AscendC::Reg::UpdateMask<T>(count);
// 当 count >= VL/sizeof(T) 时返回 ALL mask
// 当 count < VL/sizeof(T) 时返回低 count 位为 1 的 mask

MaskPattern 枚举全集(来自 kernel_reg_compute_utils.h):
ALL, VL1, VL2, VL3, VL4, VL8, VL16, VL32, VL64, VL128, M3, M4, H, Q, ALLF

MaskMergeMode —— mask 未命中位的处理策略

模式 行为 适用场景
MaskMergeMode::ZEROING(默认) mask=0 的 lane dst 写 0 第一次写入、不需要保留旧值
MaskMergeMode::MERGING mask=0 的 lane dst 保留原值 多次部分写入累加、保留之前结果

典型场景:处理非对齐尾段时,最后一拍 repeat 用 UpdateMask(remainder) 生成部分 mask,配合 MERGING 模式避免越界写 0 污染已写入数据。详见样例 mergemode/mergemode.asc

AscendC::Reg::Duplicate(dstReg, (T)2);                                       // 先填默认值
mask = AscendC::Reg::UpdateMask<T>(count);                                   // 生成部分 mask
AscendC::Reg::LoadAlign<T, AscendC::Reg::PostLiteral::POST_MODE_UPDATE>(src0Reg, src0Addr, oneRepeatSize);
AscendC::Reg::Max<T, AscendC::Reg::MaskMergeMode::MERGING>(dstReg, src0Reg, src1Reg, mask);
AscendC::Reg::StoreAlign(dstAddr + i * oneRepeatSize, dstReg, allMask);

2.3 AddrReg —— 地址偏移寄存器

AddrReg = vector_address,存储 (index, stride) 二元组,配合 LoadAlign/StoreAlign 的 AddrReg 重载,可在循环中自动累加地址偏移,省去软件计算 addr + i * size 的开销。

// 4 个 (index, stride) 对的 CreateAddrReg 重载
AddrReg CreateAddrReg<T>(uint16_t index0, uint32_t stride0);
AddrReg CreateAddrReg<T>(uint16_t i0, uint32_t s0, uint16_t i1, uint32_t s1);
AddrReg CreateAddrReg<T>(uint16_t i0, uint32_t s0, uint16_t i1, uint32_t s1,
                          uint16_t i2, uint32_t s2);
AddrReg CreateAddrReg<T>(uint16_t i0, uint32_t s0, ..., uint16_t i3, uint32_t s3);

// 使用
AscendC::Reg::AddrReg areg;
for (uint16_t i = 0; i < repeatTimes; ++i) {
    areg = AscendC::Reg::CreateAddrReg<T>(i, oneRepeatSize);   // index=i, stride=oneRepeatSize
    AscendC::Reg::LoadAlign(xReg, xAddr, areg);                // 等价于 LoadAlign(xReg, xAddr + i * oneRepeatSize)
    AscendC::Reg::StoreAlign(zAddr, zReg, areg, mask);
}

适用场景:循环中地址按固定 stride 递增,且 stride 较大需要地址计算外提时。多数情况下 addr + i * oneRepeatSize 已足够,AddrReg 在 post-update 模式或跨多段数据交错访问时更优。

2.4 UnalignRegForLoad / UnalignRegForStore —— 非对齐辅助寄存器

UnalignRegForLoad = UnalignRegForStore = vector_align,用于跨 32B/256B 边界的非对齐数据搬移。

非对齐搬移分两步:先用 LoadUnAlignPre 把跨边界数据缓存到 ureg,再用 LoadUnAlign 拼接成完整 VL 数据;写回侧用 StoreUnAlign + StoreUnAlignPost

AscendC::Reg::UnalignRegForLoad ureg0;
AscendC::Reg::UnalignRegForStore ureg1;
AscendC::Reg::LoadUnAlignPre(ureg0, srcAddr, areg);                    // 预取跨边界部分
AscendC::Reg::LoadUnAlign(srcReg, ureg0, srcAddr, areg, 0);            // 完整加载
AscendC::Reg::StoreUnAlign(dstAddr, srcReg, ureg1, areg);              // 非对齐写(缓存尾部)
// 循环结束后:
AscendC::Reg::StoreUnAlignPost(dstAddr, ureg1, areg);                  // 刷出最后残留

适用场景:输出数据长度不是 VL/sizeof(T) 的整数倍,且不能用 MaskReg 简单处理的非对齐连续场景,如 Squeeze 压缩输出。

2.5 特殊寄存器访问(GetSpr / ClearSpr)

少数场景需要读写特殊功能寄存器(System Special Purpose Register):

AscendC::Reg::ClearSpr<AscendC::SpecialPurposeReg::AR>();   // 清地址寄存器
// 详见 docs/api/SIMD-API/基础API/Reg矢量计算/系统变量访问/

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐