理解 Ascend C 的向量编程,可以先从一个问题入手:输入已经存放在 GM 中,为什么计算前还需要 UB 缓冲?

本文以 Ascend 910B 的分离架构和 LocalTensor 向量编程方式为范围,依次介绍计算单元、存储层次、数据搬运与同步,并与 CUDA 的常见编程模型对照。图示用于解释编程关系,不是物理布局图;不同代际的硬件与接口不能直接套用。

1. GM:存放输入和输出的设备内存

写加法算子时,需要先有地方存放输入数组 xy,以及输出数组 z。在 NPU 上,这些数组可以存放在 GM(Global Memory,全局内存) 中。

先把 GM 理解为设备上的“大容量数据存放区”:它保存输入和输出,计算核通过地址访问其中的数据。“全局”表示它不属于某一个计算核独占的局部缓冲,并不是说 CPU 可以随意直接访问它。

GM 和 HBM 经常一起出现,但它们回答的是两个不同的问题:

名称 回答的问题 本例中的含义
GM,全局内存 程序通过什么内存空间访问数据? 核函数访问 x、y、z 的全局地址
HBM,高带宽内存 实际用什么硬件存储数据? 910B 设备上承载这些数据的物理内存

同一份数组,从程序角度可以说“存放在 GM 中”,从硬件角度可以说“存放在 HBM 中”。 这不是两份数据,也不需要先从 GM 再拷贝到 HBM。

为什么有了 GM,还需要 UB?

存放数据和执行计算是两件事。GM 用来保存数组,但在本文使用的 Ascend C 向量编程方式中,Vector 计算单元需要从 UB(Unified Buffer,统一缓冲区) 读取数据,并把计算结果写到 UB。

UB 是靠近计算单元的一块片上存储,容量比设备的大容量内存小,可以理解为“这一次计算使用的工作区”。计算 z = x + y 时,先把输入从 GM 搬到 UB,再计算,最后把结果从 UB 搬回 GM。参与计算的输入数据也叫“操作数”,例如 1 + 2 中的 1 和 2。

Host 与 Device 的分工

Host 是运行控制程序的 CPU 一侧,Device 是执行核函数的 NPU 一侧。Host 准备输入、申请设备内存、提交任务,再取回结果;核函数处理已经交给设备的数据。

图 1:Host 与 Device 的分工

图 1:Host 与 Device 的分工。 Host 与 Device 之间的数据传输,以及 Device 内部 GM 与片上缓冲之间的搬运,是两个不同层次的步骤。

把 Host、GM 与计算核串起来看

图 2:当前 Add 的简化架构图

图 2:当前 Add 的简化架构图。 实线表示数据读写或搬运,虚线表示启动、控制或指令发射。先沿着 GM → MTE2 → UB → Vector → UB → MTE3 → GM 阅读,再看 Scalar 如何组织这些单元。AIC 画在旁边用于区分矩阵计算职责,当前 Add 不使用它。

这张图只展示一个 AIV,不表示设备只有一个向量核;图中也省略了缓存与互连细节。查看 Mermaid 图源

2. 计算资源为什么分成 Scalar、Vector 和 Cube

Scalar 负责标量计算和程序控制,例如循环、分支、地址计算以及指令发射。它可以帮助我们理解“是谁组织这段核函数”,但它不是运行 Host 程序的那颗 CPU。

Vector 执行向量运算。同一类操作作用于多个元素,很适合示例中的逐元素加法。Cube 负责矩阵运算,例如矩阵乘加;左、右矩阵分别由 L0A、L0B 提供,结果或累加中间值保存在 L0C。示例中的 Add 不需要经过 Cube。官方计算单元说明

还要避免把三者简单画成一个固定的“大核”。910B 所采用的分离架构中,矩阵计算核 AIC 和向量计算核 AIV 分开,各有 Scalar 控制。下图只画两类核的职责,不表示核数量、配比或一对一绑定关系。官方基本架构说明

图 3:910B 的 AIC 与 AIV 职责示意

图 3:910B 的 AIC 与 AIV 职责示意。 图中不表示核数量或配比。

需要注意的是:看到 C++ 风格的顺序代码,不代表里面的事情都由同一个执行单元逐条完成。控制、搬运和计算有不同的执行资源。

3. UB、L1 和 L0:名字相似,职责不同

先保留与入门有关的几种存储:

存储 这阶段应该怎样理解
GM 保存设备侧输入、输出等全局数据
UB(Unified Buffer) 本文向量编程方式中的片上操作数缓冲
L1 Buffer 矩阵数据等的片上中转与复用空间
L0A / L0B Cube 的两个矩阵输入缓冲
L0C Cube 的结果与累加缓冲

这里的 L1 Buffer、L0 Buffer 不能直接套用 CPU 的透明缓存层级。程序需要通过相应接口和数据通路组织数据。L0C 也不是“更小一级的通用缓存”。官方存储单元说明,第 4.3 节

这里暂时不需要把矩阵路径的全部转换和搬运细节展开。以 Add 为例,数据在 GM 与 UB 之间流动,计算发生在 Vector 上。

4. UB 可以类比 GPU 的 shared memory 吗

可以类比的是:两者都能帮助我们理解“容量有限、靠近计算资源、需要程序考虑如何使用的片上存储”。但这个类比到这里就应该停一下。

在常见的 CUDA 逐元素加法中,线程从 global memory 加载数值,算术操作使用寄存器,再把结果写回。简单 Add 并不要求先经过 shared memory。 Shared memory 常用于线程块内的数据共享或复用;寄存器则是线程执行时的重要操作数存储。CUDA 的 local memory 名字里虽然有 local,也不意味着它就在片上。CUDA 内存与核函数说明

而我们本例使用的 Ascend C 向量接口接受 LocalTensor,数据缓冲放在 UB。LocalTensor 是对这块数据的访问对象,并不是一个 CUDA thread,也不是每线程的一组私有寄存器。

熟悉的 GPU 概念 可以建立的联系 不能直接等同的地方
global memory 与 GM 都描述全局数据访问 地址空间不等于具体物理介质
shared memory 与 UB 都有片上缓冲的作用 使用规则与计算接口不同
CUDA thread / warp 描述线程和执行分组 不能把 256 元素解释成 256 个 NPU 线程
Tensor Core 与 Cube 都面向矩阵计算 指令、数据格式和存储通路不同
SM 都涉及局部计算资源 不能把一个 SM 直接对应一个 AIV 或 AIC

CUDA 的 SIMT 模型向程序暴露线程;线程按 warp 组织执行。示例中的 Ascend C 代码则显式组织一段向量数据及其流水依赖。比较这两个模型,是为了读懂当前代码,不是在断言所有 GPU 或所有 NPU 只能采用一种编程方式。NVIDIA 对 SIMT 的说明

所以 Add(zLocal, xLocal, yLocal, 256) 里的 256 是处理的元素数,不是线程数,也不能据此断言硬件恰好用一条底层指令或一个时钟周期完成它。

先对照简单加法的数据路径

NPU:本文的 Ascend C 向量编程方式
GM 中的 x、y
    ↓ 搬入
UB 中的 x、y → Vector 执行加法 → UB 中的 z
                                   ↓ 搬出
                               GM 中的 z

GPU:常见的 CUDA 逐元素加法
Global Memory 中的 x、y
    ↓ 线程加载
寄存器中的数值 → 执行加法 → 寄存器中的结果
                                   ↓ 线程写回
                          Global Memory 中的 z

图 4:NPU 与 GPU 简单加法的数据路径对照。 省略缓存、互连和具体指令细节。这个 Ascend C 示例显式准备 UB 缓冲;常见的 CUDA 简单加法由线程加载数值到寄存器后计算,不要求先经过 shared memory。

再看两张整体概念图

下面两张概念图用于对照设备的整体组成。它们不是具体型号的官方结构图,不能把每个框的位置和箭头当作所有 GPU、NPU 都遵循的固定设计。点击图片可打开原图查看细节。

图 5:通用 GPU 架构概念图

图 5:通用 GPU 架构概念图。 阅读时重点看计算单元、寄存器、共享内存、L2 缓存与显存。图中混用了不同厂商的术语,例如 NVIDIA 的 SM、AMD 的 CU;RT Core 和显示输出也并非所有计算卡都具备。图形渲染路径不是本文 Add 算子的必经路径。

图 6:通用 NPU 架构概念图

图 6:通用 NPU 架构概念图。 重点观察矩阵/向量计算、数据搬运和局部存储的分工。图中把 Unified Buffer 画成多个计算簇共享的一层,这不能直接对应 910B 的 AIV 局部 UB;“图编译器”也不能据此理解为一定在芯片内执行的硬件模块。学习 910B 时,以图 3 的核组织和官方文档为准。

对照时可以沿着“数据存在哪里 → 怎么搬到计算附近 → 谁完成运算 → 结果写到哪里”这四个问题阅读,而不是逐个寻找完全对应的方框。

5. 一个 Add 的数据到底怎样走

对示例中的例子,只记住这条路径:

图 7:Ascend C Add 的三个阶段

图 7:Ascend C Add 的三个阶段——搬入、计算、搬出。

图 7 中,Vector 是执行运算的单元;xLocal、yLocal、zLocal 对应的缓冲仍在 UB。箭头表示读写,不表示 Vector 里面又存了一份完整数组。

MTE 是数据搬运引擎。在这条路径上,MTE2 负责输入搬入,MTE3 负责输出搬出。Host 到设备的 aclrtMemcpy 不应被混进这张核内流水图里。

这就解释了最初的疑问:GM 中保留整段输入输出,UB 中准备本次计算要使用的数据。对当前三个独立的 256 元素 float32 缓冲而言,每个是 256 × 4 = 1024 字节,UB 数据缓冲合计 3072 字节。这不包含 Host 数组,也不包含 GM 的三块内存。

如果数据变大,不能无限扩大 UB。Tiling,就是让有限的局部缓冲分批处理更大的全局数据。

6. 为什么搬完还要同步

“代码已经写到下一行了”与“前一项硬件工作已经完成了”不是同一件事。不同流水有各自的指令队列,输入搬运、向量计算和输出搬运可以独立推进;存在数据依赖时,就必须建立完成顺序。

图 8:搬运与计算之间的同步

图 8:搬运与计算之间的两次同步。

HardEvent::MTE2_V 的方向是 MTE2 → Vector:SetFlag 安排在源流水前面的工作完成后发信号,WaitFlag 让目标流水等待这个信号。V_MTE3 则表示 Vector → MTE3。事件类型说明方向,event ID 区分相应事件;同一对 Set/Wait 要使用一致的类型和编号。官方流水与同步说明

因此单块 Add 的依赖顺序是:

  1. 把 x、y 搬进 UB。
  2. 输入搬运完成,才允许 Vector 读取。
  3. Vector 计算 z。
  4. 结果计算完成,才允许 MTE3 读取 z 并写回 GM。

这里等待的是流水之间的数据依赖,不是 CUDA 那种线程块集合意义上的 __syncthreads()。Host 等待整个 stream 完成又是更外层的等待。

7. Tiling 引入的缓冲复用依赖

前面的 Add 示例只处理一块数据。加入循环后,会出现另一种问题:下一轮搬入能否覆盖上一轮还在读取的输入缓冲?下一轮计算能否覆盖上一轮尚未搬出的结果?

这是使用 Tiling 时需要分析的缓冲复用依赖。当前的两次前向同步不能被当成任意循环的完整同步方案。理解谁生产数据、谁消费数据、什么时候允许复用缓冲,比先记一串事件名字更有用。

配套实操记录:从零跑通 Ascend C Add:我的第一天算子开发

官方参考资料

以下是本文参考的昇腾与 NVIDIA 官方资料。实验环境为 CANN 8.5.0;部分概念说明引用了较早版本的文档,具体 API 约束应以所用 CANN 版本及芯片型号为准。外部资料在新标签页打开。