AI 芯片架构基础：数据流、脉动阵列与张量核心

欢迎来到本教程。我们将深入AI芯片架构的核心概念，理解为何现代处理器能高效执行大规模神经网络计算。本教程面向初学者，无需硬件设计背景，只需对AI计算有基本认知。我们将聚焦三大基石：数据流、脉动阵列与张量核心。

理解AI计算的特殊需求

在进入具体架构之前，我们需要明确AI芯片要解决什么根本问题。传统CPU擅长处理复杂的控制逻辑与分支预测，但在面对神经网络计算时显得力不从心。神经网络的底层操作主要是大规模矩阵乘法和卷积操作，其核心是高密度的乘加运算（Multiply-Accumulate, MAC）。这就催生了专用架构的需求：

• 并行度要求极高：单次推理可能需要数十亿次运算，但运算间无复杂数据依赖。
• 数据移动功耗远高于计算功耗：芯片大部分能量消耗在数据搬运上，而非计算本身。
• 可预测的数据流：神经网络计算图是静态的，数据流模式可被提前分析和优化。

因此，AI芯片设计的核心命题是：如何在最小化数据移动的前提下，最大化并行运算单元的利用率。

数据流：芯片的血液输送系统

数据流定义了数据在计算单元、存储层级（寄存器、片上缓存、外部内存）之间移动的路径与调度方式。一个低效的数据流会阻塞高能效的计算引擎，让乘加单元空转等待数据。AI芯片中，常见的数据流策略有以下几种：

权重固定数据流

• 工作方式：将卷积核或权重矩阵固定在每个处理单元（PE）的本地寄存器中。
• 数据流动：输入特征图、部分和或输出数据在阵列中流动传递。
• 优势：最大化权重数据复用，因为权重被固定在计算单元旁，读取一次即可被多次使用。非常适用于神经网络推理场景，因为推理时权重恒定。
• 挑战：权重需要预加载，适合批量处理相同权重的任务，但对权重极多的模型可能导致固定存储不足。

输出固定数据流

• 工作方式：将部分和或输出结果留在处理单元本地，不断累加直到完成完整输出值。
• 数据流动：输入数据和权重被广播或流动进入阵列，而部分和保持在寄存器中不移动。
• 优势：消除了部分和频繁读取写入外部缓存的能耗，对训练中间数据的管理极为高效。
• 挑战：需要处理单元有足够的本地累加能力，且对输入和权重的复用模式要求更高。

行固定数据流

• 工作方式：这是学术界著名架构Eyeriss提出的能量最优数据流。它在1D卷积运算中，将一行滤波器权重和一行输入特征图保持在寄存器，通过行内广播和累加完成计算。
• 核心思想：通过最大化行卷积复用和最小化不同数据类型（特征图、权重、部分和）的访问功耗，达到接近理论最低的能耗水平。
• 实际应用：启发了许多低功耗端侧推理芯片的设计，强调通过数据流优化而非单纯堆计算单元来提升能效。

数据流对架构的影响

选择何种数据流，决定了芯片中片上网络的设计。片上网络包括连接方式（如网格、环形、总线）和数据路由规则。优化数据流的目标永远是“让数据离计算越近越好，让数据移动尽可能少且重复”。

脉动阵列：经典的高效计算架构

脉动阵列是一种专用硬件结构，其灵感来源于心脏有节奏的泵血过程。数据像脉搏一样，被规律地“泵”过处理单元阵列，每个单元完成单一乘加操作后将结果传递给相邻单元。

脉动阵列的基本原理

• 结构：通常是一个2D网格，每个PE只与其直接邻居相连。PE内部包含一个乘法器、一个加法器和一个寄存器。
• 数据脉动：数据从左侧和顶部流入阵列，在每个时钟周期，PE接收来自左/上部的数据，完成乘累加，然后将输入数据传递到右/下邻居，将累加结果保存在本地或传递至右侧。
• 特点：
  • 高度规则化：唯有简单本地连接，无全局连线，易于扩展频率与规模。
  • 极高吞吐：用流水线方式实现极高的数据重用与并行度。
  • 确定性行为：所有运算、数据移动完全由编译器预调度，无运行时缓存缺失问题。

经典权重固定脉动阵列（如Google TPUv1）

以Google的首款AI加速器TPUv1为例，它使用了一个256x256的脉动阵列（共65,536个MAC单元），采用权重固定数据流：

1. 权重预加载：将卷积核权重从顶部预填充进阵列，每个PE存放一个权重值。
2. 输入从左侧流入：特征图数据以脉动方式从左侧流入，每一拍向右移动一列。
3. 部分和从上向下流动：每个PE将当前乘法结果与从上邻居流入的部分和相加，并传递给下邻居。
4. 输出从底部收集：最终计算结果从阵列底部流出。

整个过程如同一个超级高效的“二维乘加流水线”，时钟频率700 MHz时，峰值算力可达92 TOPS（以INT8计算）。TPU的成功证明了脉动阵列在云端推理的巨大潜力。

脉动阵列的变体与局限

现代设计中也有“输出固定”脉动阵列，将部分和存储在本地，而广播输入和权重。脉动阵列的刚性结构也存在挑战：

• 对稀疏计算不友好：无法有效跳过零值权重，因为数据脉动是强制同步的。
• 映射灵活性受限：非规则形状计算（如非标准卷积核大小或步长）需要复杂的映射策略，可能导致利用率下降。
• 片上缓冲区需求大：若要平衡数据供给，需要精确设计的缓冲区尺寸。

张量核心：迈向灵活高效的SIMD

如果说脉动阵列是以阵列结构实现规整的数据流，那**张量核心（Tensor Core）**则是NVIDIA在GPU中引入的可编程矩阵乘加单元，代表了另一种思路：将矩阵运算原语直接集成到可编程SM（流多处理器）中。

张量核心的运作机制

第一代Volta张量核心专为4x4矩阵乘加设计，执行公式：

D = A x B + C

其中A、B、C、D均为4x4矩阵。在单个时钟周期内，张量核心完成：

• 输入矩阵A（FP16） 和 矩阵B（FP16） 进行乘法。
• 结果与矩阵C（FP32）进行加法累加，输出矩阵D（FP32）。

这是通过一个专门的硬件单元实现的，并非在常规CUDA核心上模拟。从Ampere架构开始，张量核心增加了对稀疏模式（2:4结构化稀疏）和TF32、BF16等数据格式的支持，大幅提升吞吐。

张量核心与数据流的关系

张量核心本质上是线程块级别的并行原语：

• 程序员（或cuBLAS库）将大的矩阵问题分割成多个4x4（或更大瓦片）的子矩阵。
• 每个张量核心操作需要从寄存器或共享内存中获取数据。
• GPU的线程束调度器与内存层次（全局内存、L2缓存、共享内存、寄存器）协同工作，实现数据流供给。

在张量核心之上，CUDA编程模型提供了灵活的数据流控制。不同的矩阵乘法算法（如阶梯式瓦片算法）可以最大化数据重用，隐藏内存延迟。

对比脉动阵列

两者并非对立，而是效率与灵活性的权衡。部分新型架构（如Tesla Dojo的D1芯片）甚至融合了张量核心的可编程性与脉动阵列的规整数据流。

现代AI芯片的融合趋势

现实中，顶级AI芯片往往综合运用上述概念，例如：

• 数据流与脉动阵列结合：近存计算（near-memory computing）架构，在内存单元附近布置脉动阵列，降低数据移动。
• 数据流引擎+张量核心：有些芯片用可配置的数据流引擎作为互连，连接多个张量核心集群，实现芯片级扩展。
• 稀疏数据流：针对激活和权重的稀疏性，设计专门的跳过逻辑，打破传统脉动的强制同步，仅在非零数据到达时才触发PE工作。

总结生态与学习路径

理解“数据流决定架构、脉动阵列提供高效规整计算、张量核心引入可编程专用原语”这三驾马车，便抓住了AI芯片架构的核心。

• 若关注能效和固定推理：深入脉动阵列与数据流设计。
• 若关注训练灵活性和快速发展：掌握张量核心的编程与优化。
• 下一步学习：建议研究具体芯片的微架构文档（如NVIDIA白皮书、Google TPU论文），从数据流模拟和小规模矩阵乘法优化入手，逐步建立起从算法到硬件的完整映射思维。

AI芯片架构仍在快速演进，掌握基础是应对未来创新的底牌。