跳转至

NPU 原理

本节从矩阵乘法出发，介绍脉动阵列 NPU 的基本原理：为什么要把 B 矩阵驻留在 PE 中、A 矩阵如何斜入、C 矩阵如何斜出，以及片上存储和矩阵分块的作用。理解这些内容后，再看硬件和软件实现会容易得多。

从矩阵乘法到脉动阵列

从矩阵乘法说起

首先回到矩阵乘法的定义，从 2x2 矩阵开始，设 A、B、C 都是 2x2 矩阵：

A = ┌ a00 a01 ┐      B = ┌ b00 b01 ┐      C = ┌ c00 c01 ┐
    └ a10 a11 ┘          └ b10 b11 ┘          └ c10 c11 ┘

约定 aij 代表 A 矩阵的第 i 行第 j 列元素（约定下标从 0 开始，下同），bij 和 cij 同理。按矩阵乘法的定义，C 矩阵的每个元素是一组乘积的和：

c00 = a00*b00 + a01*b10
c01 = a00*b01 + a01*b11
c10 = a10*b00 + a11*b10
c11 = a10*b01 + a11*b11

观察矩阵乘法定义中，B 矩阵元素是如何参与到计算当中的：

  • B 的第 0 行（b00, b01）只会和 A 的第 0 列（a00, a10）相乘；
  • B 的第 1 行（b10, b11）只会和 A 的第 1 列（a01, a11）相乘。

可见 B 矩阵的每个元素都会和 2 个 A 矩阵的元素相乘，并且同一行的 B 矩阵元素，与之相乘的 A 矩阵元素是相同的，后面会用到这个性质。

处理单元 PE

下面转向硬件的视角，为了提高计算的性能，应当避免数据在内存和寄存器之间来回传输，尽量把频繁使用的数据保存在芯片内部的寄存器上并重复使用。既然 B 矩阵的每个元素都要和 2 个 A 矩阵的元素相乘，那就把 2x2 的 B 矩阵直接用寄存器保存下来。考虑到矩阵乘后续还要做乘法和加法，就在每个寄存器旁边配备一个乘法器和加法器。这里把一个寄存器、一个乘法器和一个加法器组合起来，叫做一个 PE（Processing Element，处理单元）。把四个 PE 按照 2x2 矩阵的顺序排列，也按照类似矩阵元素的方法，给它们编号，即 PEij 是在第 i 行第 j 列的 PE，内有寄存器、乘法器和加法器，其中 PEij 的寄存器保存了 bij 的值：

根据前面的分析，已知 B 矩阵的第 0 行只会和 A 矩阵的第 0 列相乘，那为了进行这些乘法运算，就让 A 的元素按列进入对应的 PE 行，具体地：

  • A 的第 0 列（a00、a10）进入第 0 行的 PE00 和 PE01，每个元素（a00、a10）都在 PE00 与 b00 相乘，也在 PE01 与 b01 相乘；
  • A 的第 1 列（a01、a11）进入第 1 行的 PE10 和 PE11，每个元素（a01、a11）都在 PE10 与 b10 相乘，也在 PE11 与 b11 相乘。

这样就可以得到计算 C 矩阵所需要的所有的 8 个乘积：

  • 第 0 行第 0 列的 PE00 计算出 a00*b00 和 a10*b00
  • 第 0 行第 1 列的 PE01 计算出 a00*b01 和 a10*b01
  • 第 1 行第 0 列的 PE10 计算出 a01*b10 和 a11*b10
  • 第 1 行第 1 列的 PE11 计算出 a01*b11 和 a11*b11

回到矩阵乘的定义，得到乘积以后，还需要对乘积进行累加，才能得到 C 矩阵的元素。首先观察 c00 的定义 c00 = a00*b00 + a01*b10，求和的两项 a00*b00 和 a01*b10 分别来自 PE00 和 PE10，它们都是第 0 列 PE。其余元素同理：

  • c00 = a00*b00 + a01*b10 ，求和的两项来自第 0 列 PE00 和 PE10
  • c10 = a10*b00 + a11*b10 ，求和的两项来自第 0 列 PE00 和 PE10
  • c01 = a00*b01 + a01*b11 ，求和的两项来自第 1 列 PE01 和 PE11
  • c11 = a10*b01 + a11*b11 ，求和的两项来自第 1 列 PE01 和 PE11

那么为了计算 c00，只需要在 PE00 计算出 a00*b00 的结果，在 PE10 计算出 a01*b10 的结果，加起来就可以了。那么这个加法在哪里做呢？正好 PE10 自带了加法器，就用它的加法器求和：PE00 计算出 a00*b00 后，向下传给 PE10，PE10 内部首先用乘法器计算 a01*b10，再用加法器计算 a00*b00+a01*b10，就得到了 c00。这个流程画出来，就是下图：

类似地，为了计算 c01，则要在 PE01 计算出 a00*b01 的结果，在 PE11 计算出 a01*b11 的结果，最后在 PE11 的加法器内计算 a00*b01+a01*b11，就得到了 c01。既然 PE00 和 PE01 都需要 a00 作为输入，那就让 a00 先进入 PE00，再从 PE00 转给 PE01；类似地，a01 首先进入 PE10，再从 PE10 转给 PE11，这样整体的数据流动如下：

其他元素同理：

  • 为了计算 c00：a00 进入 PE00，和 b00 做乘法，得到 a00*b00，向下传到 PE10；另一方面，a01 进入 PE10，和 b10 做乘法，得到 a01*b10；利用 PE10 的加法器，得到 a00*b00+a01*b10，即 c00，输出计算结果
  • 为了计算 c10：a10 进入 PE00，和 b00 做乘法，得到 a10*b00，向下传到 PE10；另一方面，a11 进入 PE10，和 b10 做乘法，得到 a11*b10；利用 PE10 的加法器，得到 a10*b00+a11*b10，即 c10，输出计算结果
  • 为了计算 c01：a00 经过 PE00，进入 PE01，和 b01 做乘法，得到 a00*b01，向下传到 PE11；另一方面，a01 经过 PE10，进入 PE11，和 b11 做乘法，得到 a01*b11；利用 PE11 的加法器，得到 a00*b01+a01*b11，即 c01，输出计算结果
  • 为了计算 c11：a10 经过 PE00，进入 PE01，和 b01 做乘法，得到 a10*b01，向下传到 PE11；另一方面，a11 经过 PE10，进入 PE11，和 b11 做乘法，得到 a11*b11；利用 PE11 的加法器，得到 a10*b01+a11*b11，即 c11，输出计算结果

这样就实现了矩阵乘，我们把这样由 NxN 个 PE 组成的阵列，叫做脉动阵列（Systolic Array）。具体为啥叫脉动，后面会进行揭晓。PE 之间的数据流动非常规律，A 矩阵的元素从左向右流动，乘积从上往下流动，且两个方向的流动总是在相邻的 PE 之间。由于每个 PE 内只有一个加法器，只能做一次加法，对于更大的脉动阵列，比如 4x4，列方向上的乘积累加就需要在多个 PE 内的加法器内完成：PE00 计算 a00*b00，PE10 计算 a00*b00+a01*b10，PE20 计算 a00*b00+a01*b10+a02*b20，PE30 计算 a00*b00+a01*b10+a02*b20+a03*b30 并输出。所以实际在列方向上流动的，不只是乘积，实际上是乘积的部分和（Partial Sum）或者说前缀和（Prefix Sum），通过多步完成整个累加过程。

接下来，思考 PE 内部的设计。为了完成 C = AxB 的计算，只需要事先把 B 矩阵的元素保存在 PE 内部，让 A 矩阵的第 i 列进入在第 i 行的 PE 里做乘法，再把乘积沿列方向向下累加，就得到了 C 的全部元素。为了做到这一点，PE 只需要实现三种功能：

  1. 做乘法：从左边接收一个 A 的元素，和内部保存的 B 元素相乘；
  2. 做加法并向下传：从上方接收一个部分和，和内部乘法器输出的乘积求和，把结果传给下方的 PE；
  3. 向右传：把 A 的元素原样传给右边的 PE，因为同一行的 A 还要在右边 PE 做乘法。

但这时候就出现了一个问题：在硬件里面，如果只有如上图的组合逻辑（乘法、加法以及 PE 之间的连线都是组合逻辑），那么输入一旦固定，输出也就不会变化，然而每个 PE 都需要做多次计算：例如 PE00 不仅需要负责 a00 的计算，还需要负责 a10 的计算。所以势必要在不同的时钟周期内，执行不同的计算。

一种简单的方法就是模仿单周期的 CPU，用时钟来驱动巨大的组合逻辑：第 0 拍给第 0 行输入 a00，给第 1 行输入 a01，等待巨大的组合逻辑把 c00 和 c01 算出来；第 1 拍给第 0 行输入 a10，给第 1 行输入 a11，把 c10 和 c11 算出来。但是，前面的实验里写流水线 CPU 的经验告诉我们，单周期会导致时钟频率上不去，需要把这个组合逻辑拆分为多个周期实现，再进一步流水起来。

为了实现多周期，观察这里的计算过程，它天然分为了多步，既有 PE 内部的计算，又有 PE 之间的数据传递，那么很自然地，在每个 PE 内部设置输出寄存器，把组合逻辑分成多段，以 c00 的计算为例，把它分成两拍完成：

  • 第 0 拍，在 PE00 里输入 a00，计算 a00*b00，暂存到输出寄存器内；同时在 PE10 里输入 a01，计算 a01*b10，暂存在输出寄存器内
  • 第 1 拍，PE10 从 PE00 的输出寄存器取出 a00*b00 的计算结果，从 PE10 自己的输出寄存器取出 a01*b10 的计算结果，通过加法器求和 a00*b00+a01*b10，下一拍输出 c00 的结果

但这种设计，没法扩展到更大的脉动阵列，比如 4x4 的脉动阵列下，第 0 列有四个 PE，第 0 拍出了四个乘法结果，第 1 拍没法求和，因为同一列的 PE 之间离得更远了，而只有相邻 PE 之间有连接。怎么解决这个问题呢？答案是，延迟 a01 的输入到第 1 拍，让加法和乘法在第 1 拍内同时在 PE10 中完成：

  • 第 0 拍，在 PE00 里输入 a00，计算 a00*b00，暂存到输出寄存器内；此时 PE10 闲置
  • 第 1 拍，PE10 从 PE00 的输出寄存器得到 a00*b00 的计算结果，同时输入 a01，计算 a01*b10，通过加法器求和 a00*b00+a01*b10，下一拍输出 c00 的结果

这样即使是更大的脉动阵列，比如 4x4，也可以通过 4 拍完成计算：第 0 拍 PE00 计算 a00*b00，第 1 拍 PE10 计算 a00*b00+a01*b10，第 2 拍 PE20 计算 a00*b00+a01*b10+a02*b20，第 3 拍 PE30 计算 a00*b00+a01*b10+a02*b20+a03*b30 并输出。

引入了 PE 内部的寄存器以后，内部结构如下图：

此时有一个问题：b_reg 只有输出没有输入，它的数值从哪里来？上面的分析里省略了这个细节，但解决起来也并不复杂：类似部分和的处理，从上方输入 B 矩阵的元素，打一拍再往下传，当整个阵列的 PE 都得到对应元素后，就停止更新：

后面把这个操作叫做 LDTILE，就是 Load Tile 的意思，Tile 指的是子矩阵，因为实际计算的 B 肯定不止 2x2 的大小，所以每次只会加载它的一个子矩阵。这种简化设计下，在加载 B 矩阵时，就不能同时进行计算了。实际的 NPU 设计，通常会引入 Double Buffering 来解决这个问题，从而掩盖 LDTILE 的延迟。用类似的命名方法，把矩阵乘计算命名为 TMUL（Tile Multiply）。

这样就通过多周期实现了计算，诀窍就是，延迟 A 矩阵元素的输入：第 0 拍给 PE00 输入 a00，依此类推，第 i 拍给 PEi0 输入 a0i，我们把这种延迟输入的模式形象地称为斜入。这时候，我们可以惊奇地发现，实际上，这个多周期自然组成了流水线，不同元素的计算正好错开，且可以正确进行。具体这是怎么做的呢？

斜入斜出

刚才提到，为了计算 c00，A 矩阵的第一列是按照下面的顺序进入脉动阵列：

第 0 行 第 1 行
第 0 拍 a00 进入 空闲
第 1 拍 空闲 a01 进入

如果把 A 矩阵的第二列紧跟着第一列的后面，进入脉动阵列：

第 0 行 第 1 行
第 0 拍 a00 进入 空闲
第 1 拍 a10 进入 a01 进入
第 2 拍 空闲 a11 进入

在模拟器里观察一下每一拍发生了什么事情：

PE00 PE01 PE10 PE11
第 0 拍 计算 a00*b00 空闲 空闲 空闲
第 1 拍 计算 a10*b00 计算 a00*b01 计算 c00=a00*b00+a01*b10 空闲
第 2 拍 空闲 空闲 计算 c10=a10*b00+a11*b10，输出 c00 计算 c01=a00*b01+a01*b11
第 3 拍 空闲 空闲 输出 c10 计算 c11=a10*b01+a11*b11，输出 c01
第 4 拍 空闲 空闲 空闲 输出 c11

可以看到这些 PE 自然构成了流水线，在 5 拍的时间内，完成了 C = AxB 的矩阵运算。可以看到，不仅 A 是斜着进入阵列的，C 也是斜着离开阵列的。这种随着时钟周期，数据在阵列中就像水波一样前进，所以这种 PE 阵列就叫作脉动阵列。特别地，这里的设计里，B 矩阵是驻留在阵列当中的，所以它就叫 weight stationary 脉动阵列。

这就是一个完整的矩阵乘单元：B 的一个子矩阵（tile）驻留在阵列里，A 从左边流入，C 从底部流出。事实上，这里的 A 不一定是方阵，可以是长条形的，只需要一列一列地连续送入，这样就可以更好地填补流水线里的空泡。具体地，对于一个 NxN 的脉动阵列，它计算 C=AxB 的矩阵乘，其中 A 是 MxN 矩阵，B 是 NxN 矩阵，C 是 MxN 矩阵，每一拍完成的任务如下：

  • 第 0 拍到第 M+N-2 拍：输入 A 矩阵，A 矩阵元素 A[i][j] 在第 i+j 拍从第 j 行 PE 进入阵列；
  • 第 N 拍到第 M+2N-2 拍：输出 C 矩阵，C 矩阵元素 C[i][j] 在第 i+j+N 拍第 j 列 PE 离开阵列。

一共需要 M+2N-1 拍的时间，完成 MxNxNx2 次运算。

NPU 上的数据流动

片上存储 L0A/L0B/L0C

通过上文的分析，已知 4x4 的脉动阵列在计算时，每个周期需要从左侧输入 4 个 int8 数，从底部输出 4 个 int32 数，并且在计算前，还需要从阵列顶部每拍输入 4 个 int8 数。由于脉动阵列的特性，计算时数据的移动是要求严格满足斜入斜出的顺序的，否则会计算出错误的结果。而内存的延迟通常是不确定的（虽然本课程所使用的内存是 SRAM，可以实现为固定延迟，但实际的 NPU 都是用 DDR 内存，而即使这里用的是 SRAM，由于需要和 CPU 争抢总线，也会出现延迟的波动），那就不能直接把数据从内存导入到脉动阵列，而且随着阵列的扩大，其需要的带宽也很容易超过内存能提供的带宽。所以必须要在片内设置若干片上存储，用来提供稳定的高带宽的数据访问，保证可以在正确的拍传正确的数据。

既然矩阵乘法涉及到三个矩阵，并且三个矩阵在脉动阵列的不同方位：A 矩阵从脉动阵列左侧流入，B 矩阵从脉动阵列顶部流入，C 矩阵从脉动阵列底部流出，因此就针对这三个矩阵分别设计对应的片上存储：L0A 用来保存 A 矩阵，L0B 用来保存 B 矩阵，L0C 用来保存 C 矩阵。在 FPGA 上，片上存储通常用 LUTRAM 或 BlockRAM 来实现，考虑到未来可能更大的阵列需要更大的片上存储容量，本实验采用 BlockRAM 来实现，此时就需要考虑这三块片上存储都应该是什么样的规格：

  • L0A：计算（TMUL）时，每周期需要给脉动阵列提供 4 个 int8，不妨用一个 32bit 宽的 BlockRAM 来实现，每周期读出 A 矩阵的一行数据，打拍后斜入脉动阵列
  • L0B：加载 B 矩阵（LDTILE）时，每周期需要给脉动阵列提供 4 个 int8，因此也用一个 32bit 宽的 BlockRAM 来实现
  • L0C：计算（TMUL）时，脉动阵列每周期会输出 4 个 int32，那就用一个 128bit 宽的 BlockRAM 来实现

L0A/L0B/L0C 的命名借鉴了华为 Ascend NPU：在 Google TPU 原论文里，与之对应的是 Unified Buffer/Weight FIFO/Accumulators，对初学者来说不够直观，因此采用了 L0A/L0B/L0C 的命名方式。

上面讨论了 BlockRAM 的宽度，那么深度呢？这就是一个可以调优的参数了，例如更大的 L0A 和 L0C，首先可以把 A/C 矩阵的规模扩大（前面提到，A 和 C 实际上可以是长条形矩阵，例如 A 是 Mx4，B 是 4x4，C 是 Mx4，随着片上存储增大，M 也可以增大），这样流水线的空泡会减少，其次还可以进一步做 Double Buffering，当 L0A 的一部分在参与计算的时候，另一部分在进行数据的搬运，掩盖数据搬运的延迟。类似地，L0B 也可以扩大到能放下多个 B 矩阵，做 Double 甚至 Triple Buffering。提供的代码是一个比较简单的实现，没有 Double Buffering，所以这些片上存储都做得比较小。

有了片上存储以后，接下来就是思考数据如何搬运了。

数据搬运

从数据的视角来看，起点是 CPU 侧的内存：内存里已经有了 A、B 和 C 矩阵。所以首先要做的是，把数据从内存搬运到 NPU 的片上存储，也就是上面所说的 L0A/L0B/L0C。在实际 NPU 上，这一步是硬件完成的，而目前提供的代码，为了理解方便，直接把 NPU 的 L0A/L0B/L0C 向 CPU 暴露为 MMIO，类似串口，CPU 可以直接通过特定内存地址来访问 NPU 上的这些片上存储，因此在 CPU 和 NPU 之间的数据搬运，实际上就是在 CPU 上执行 Load/Store 指令，通过 MMIO 在 Wishbone 总线上搬运数据。

当数据搬运到了 L0A/L0B/L0C 以后，就需要由控制逻辑完成后续的数据搬运，包括：

  • 执行 LDTILE 指令时：逆序按行加载 B 矩阵，从 L0B 读取一行行元素，从脉动阵列顶部传入
  • 执行 TMUL 指令时：顺序按行加载 A 矩阵，通过一系列的寄存器实现斜入，从脉动阵列左边缘传入；与此同时，脉动阵列底部会逐渐斜出 C 矩阵的值，通过一系列的寄存器对齐后，写入到 L0C

矩阵分块

到这里，就可以加速矩阵运算了：CPU 先通过 MMIO 把数据搬运到 L0A 和 L0B，调用 NPU 的 LDTILE 指令和 TMUL 指令，然后从 L0C 取出结果。但是，问题出现了：如果要计算的矩阵很大，超出了 L0A/L0B/L0C 的容量范围，该如何进行计算？

答案是矩阵分块：在线性代数课上学过矩阵的一个性质，那就是可以对矩阵进行分块，对子块按照类似矩阵乘法的规则进行矩阵乘，结果与原矩阵乘法相同，比如把 A 和 B 分别切成四个子块，分别记为 Aij 和 Bij：

A = ┌ A00 A01 ┐      B = ┌ B00 B01 ┐      C = ┌ C00 C01 ┐
    └ A10 A11 ┘          └ B10 B11 ┘          └ C10 C11 ┘

那么矩阵乘的结果 C 矩阵的子块也符合如下关系：

C00 = A00*B00 + A01*B10
C01 = A00*B01 + A01*B11
C10 = A10*B00 + A11*B10
C11 = A10*B01 + A11*B11

与元素的运算过程是类似的，只不过这里是把一个子块作为整体来进行运算。通过分块，就可以把大的矩阵分成很多个小的子块，使得这些子块可以在 NPU 内高效地进行运算。分块在英文里是 Tiling，子块就是一个 Tile，LDTILE 和 TMUL 指令的命名方式也是从这来的。

采用分块以后，就会出现一个新的需求：计算 C00 的时候，不仅需要进行两次矩阵乘 A00*B00 和 A01*B10，还需要把这两个乘积加起来。因此，实际上 NPU 计算的，不只是 C = AxB，实际上是 C += AxB，实现方法就是在写入 L0C 之前，读取 L0C 中原来的值，累加上新从脉动阵列流出的乘积，再写入 L0C。那第一次矩阵乘怎么办？那就给 TMUL 指令加一个参数 acc（accumulate 的缩写），acc=0 代表不累加，把结果写入 L0C；acc=1 代表累加，先读 L0C，累加后写入。有了这项能力后，要计算 C00 = A00*B00 + A01*B10，可以拆分成这些步骤：

  • 把 A00、A01 搬运到 L0A、把 B00、B10 搬运到 L0B
  • 执行 LDTILE B00
  • 执行 TMUL C00 = A00*B00
  • 执行 LDTILE B10
  • 执行 TMUL C00 += A01*B10
  • 把 C00 从 L0C 搬运回 CPU

小结

脉动阵列的关键在于三点：

  1. 权重驻留：B 矩阵的子块（tile）加载到 PE 内部的寄存器中，计算期间不再访问 L0B，从而把带宽压力转移到复用率最高的权重上；
  2. 斜入斜出：A 矩阵按列、按拍错开进入阵列，C 矩阵的部分和沿列向下流动，整个阵列自然形成流水线；
  3. 片上存储 + 分块：L0A/L0B/L0C 提供稳定的高带宽访问，软件用分块和 acc 累加把大矩阵拆成 4x4 的子块计算。

一个 4x4 阵列计算 Mx4x4 的矩阵乘需要 M+2N-1 拍，完成 MxNxNx2 次运算。理论上，阵列每拍可以做 2xNxN 次运算；实际利用率取决于数据搬运、指令开销和流水线气泡，这正是后续性能优化的重点。

接下来可以继续阅读 NPU 硬件实现，看看这些原理如何落到 SystemVerilog 代码中。

作者：Jiajie Chen