跳转至

NPU 软件实现

本节介绍如何在 CPU 上用 C 代码操作 NPU：通过 MMIO 访问寄存器和 L0，实现 4x4 基础矩阵乘和通用的分块矩阵乘。对应的代码在 thu-cs-lab/cod-npu 仓库根目录下的 software/npu.c 和 software/npu.h 中。

MMIO

有了硬件以后，还需要有配套的软件，把它驱动起来。上文提到，NPU 已经挂载到了 Wishbone 总线上，而 CPU 可以通过 Load 和 Store 指令直接进行 MMIO，因此操作 NPU，实际上就是对指定地址的读写来完成。为了编程的方便，这里使用 C 代码来编写，那么 C 代码里，如何编写 MMIO 呢？

首先是定义一些常量，用于计算各 MMIO 寄存器的地址，以及 L0A/L0B/L0C 内数据的地址：

#define NPU_BASE 0x20000000ul

// 寄存器
#define NPU_REG_CMD     (NPU_BASE + 0x00)
#define NPU_REG_STATUS  (NPU_BASE + 0x04)
#define NPU_REG_A_BASE  (NPU_BASE + 0x08)
#define NPU_REG_B_BASE  (NPU_BASE + 0x0C)
#define NPU_REG_C_BASE  (NPU_BASE + 0x10)
#define NPU_REG_M_ROWS  (NPU_BASE + 0x14)

// L0A/L0B/L0C 数据地址
#define NPU_L0A_DATA(w) (NPU_BASE + 0x100 + (w) * 4)
#define NPU_L0B_DATA(w) (NPU_BASE + 0x200 + (w) * 4)
#define NPU_L0C_DATA(w) (NPU_BASE + 0x300 + (w) * 4)

有了这些地址定义后，下一步就是进行实际的写入了。在 C 代码里，对 MMIO 地址的写入必须通过 volatile 进行，否则编译器可能会把它优化掉，这里定义了两个工具函数来保证 volatile 的读写：

#include <stdint.h>

static inline uint32_t npu_read(uint32_t addr) {
  return *(volatile uint32_t *)(uintptr_t)addr;
}
static inline void npu_write(uint32_t addr, uint32_t val) {
  *(volatile uint32_t *)(uintptr_t)addr = val;
}

接下来，要完成一些对 NPU 的操作，实际上就是用上述两个函数来进行 MMIO，比如：

执行 LDTILE 指令（B 已写入 L0B，b_base 是 4x4 B 矩阵在 L0B 中的起始地址，单位是 4 字节）：

#define NPU_CMD_LDTILE   (1u << 0)
npu_write(NPU_REG_B_BASE, b_base);   // LDTILE 从 L0B 的 b_base 开始取 4x4 矩阵
npu_write(NPU_REG_CMD, NPU_CMD_LDTILE);

执行 TMUL 指令：

#define NPU_CMD_TMUL     (1u << 1)
#define NPU_CMD_ACC      (1u << 2)

npu_write(NPU_REG_A_BASE, a_base);
npu_write(NPU_REG_C_BASE, c_base);
npu_write(NPU_REG_M_ROWS, m_rows - 1); // 减去 1 是因为寄存器保存的值是实际 m_rows 减一
npu_write(NPU_REG_CMD, NPU_CMD_TMUL | (acc ? NPU_CMD_ACC : 0)); // 开始执行 TMUL 指令，acc 代表是否累加

判断 NPU 是否正在执行指令、是否已经完成了指令的执行：

#define NPU_STATUS_BUSY  (1u << 0)
#define NPU_STATUS_DONE  (1u << 1)

int npu_busy(void) {
  return npu_read(NPU_REG_STATUS) & NPU_STATUS_BUSY;
}

int npu_done(void) {
  return npu_read(NPU_REG_STATUS) & NPU_STATUS_DONE;
}

类似地，写入数据到 L0A/L0B：

// 向 L0A 下标为 index 的 uint32 写入 value
npu_write(NPU_L0A_DATA(index), value);
// 向 L0B 下标为 index 的 uint32 写入 value
npu_write(NPU_L0B_DATA(index), value);

计算完成后，从 L0C 读取数据：

// 从 L0C 第 index 个 uint32 读取
npu_read(NPU_L0C_DATA(index));

在这些 NPU 操作函数的基础上，就可以进一步实现矩阵乘法了。

基础矩阵乘法

基于以上函数，已经可以实现一个最简单的 4x4x4（即 4x4 矩阵乘以 4x4 矩阵）矩阵乘法：A、B、C 矩阵都是 4x4 的大小，计算过程如下：

// 把 A 矩阵的值写入 L0A
npu_write(NPU_L0A_DATA(0), a_row_0); // 写入 A 第 0 行的四个 int8，合并成一个 int32，一次写入，下同
npu_write(NPU_L0A_DATA(1), a_row_1);
npu_write(NPU_L0A_DATA(2), a_row_2);
npu_write(NPU_L0A_DATA(3), a_row_3);

// 把 B 矩阵的值写入 L0B，原理同上
npu_write(NPU_L0B_DATA(0), b_row_0);
npu_write(NPU_L0B_DATA(1), b_row_1);
npu_write(NPU_L0B_DATA(2), b_row_2);
npu_write(NPU_L0B_DATA(3), b_row_3);

// 执行 LDTILE 指令，从 L0B 的地址 0 加载 B 矩阵
npu_write(NPU_REG_B_BASE, 0);
npu_write(NPU_REG_CMD, NPU_CMD_LDTILE);

// 等待 LDTILE 指令完成执行
while (npu_busy())
  ;

// 执行 TMUL 指令，从 L0A 的地址 0 加载 A 矩阵，计算结果写入 L0C 的地址 0，不累加
npu_write(NPU_REG_A_BASE, 0);
npu_write(NPU_REG_C_BASE, 0);
npu_write(NPU_REG_M_ROWS, 4 - 1);
npu_write(NPU_REG_CMD, NPU_CMD_TMUL);

// 等待 TMUL 指令完成执行
while (npu_busy())
  ;

// 读取 C 矩阵结果。
// c_base=0 时 C[i][j] 对应 NPU_L0C_DATA(i*4 + j)
int32_t c00 = (int32_t)npu_read(NPU_L0C_DATA(0));
int32_t c01 = (int32_t)npu_read(NPU_L0C_DATA(1));
// ...
int32_t c33 = (int32_t)npu_read(NPU_L0C_DATA(15));

通用矩阵乘法

但是，只能计算 4x4x4 的矩阵乘还不够，实际的矩阵乘法，其矩阵的大小可能是各种各样的。约定矩阵乘 MxNxK 代表用 MxK 的矩阵 A，乘以 KxN 的矩阵 B，得到 MxN 的矩阵 C。此时的 M、N 和 K 可能会特别巨大，比如 512x512x512，一次进行 512*512*512*2 次运算，但是硬件的脉动阵列又只有 4x4，怎么实现更大的 MxNxK 呢？先不考虑 M、N、K 不能整除 4 的情况。

这就需要软件来补齐，核心的思想就是对矩阵进行分块。在大一的线性代数课程中，学过分块矩阵的乘法规律，假如把 A、B 和 C 矩阵分别划分成四个小的矩阵：

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

这里的 Aij、Bij 和 Cij 都是对应矩阵的子矩阵。那么，分块矩阵的乘法，满足：

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

其实就是把原来矩阵乘法里的元素，换成了子矩阵，公式是一样的。有这个性质以后，就可以用小矩阵乘法，来实现大矩阵乘法了：例如，把 512x512 矩阵，分成 128x128 个 4x4 矩阵，然后这 4x4 矩阵内部的矩阵乘法，就用 NPU 来加速；而软件要做的，就是按照一定的顺序，把大矩阵里的小分块，搬运到 NPU 上，在 NPU 上执行 LDTILE 和 TMUL 指令，再把结果搬运到大矩阵的对应位置。

此时，TMUL 指令在 L0C 上累加能力就有用武之地了：例如在计算 C00 的时候，先计算 A00*B00，把结果保存到 L0C 里，此时可以继续计算 A01*B10，指定 TMUL 指令在计算矩阵乘的同时，把结果累加到 L0C 的同一个位置，那么计算完成以后，L0C 里的数据就是 C00，最后只需要再进行一次搬运。

第一版通用矩阵乘法

这就推导出来了第一个版本的通用矩阵乘法：

  1. 把 A、B、C 矩阵都划分成若干个 4x4 的子矩阵；
  2. 第一层循环：遍历 C 矩阵的子矩阵行，i 从 0 到 M/4；
  3. 第二层循环：在 C 矩阵的子矩阵行里，遍历子矩阵，j 从 0 到 N/4，此时 Cij 就是 C 矩阵里的一个 4x4 子矩阵；
  4. 第三层循环：针对矩阵乘法的公式，Cij = Sum_k(Aik * Bkj)，遍历 k，k 从 0 到 K/4；
  5. 第三层循环内，每次循环做如下工作：
    1. 把 Bkj 子矩阵从 SRAM 搬运到 L0B；
    2. 执行 LDTILE 指令，把 Bkj 子矩阵加载到脉动阵列；
    3. 把 Aik 子矩阵从 SRAM 搬运到 L0A；
    4. 执行 TMUL 指令，对于第一次内层循环 k=0，执行 Cij = Ai0 * B0j，即 acc = 0；此后的内层循环 k>0，执行 Cij += Aik * Bkj，即 acc = 1；
  6. 内层循环结束后，把 Cij 子矩阵从 L0C 搬运回 SRAM。

对应的 C 代码（部分细节略去）：

// A: MxK, B: KxN, C: MxN
void npu_matmul_v1(int8_t *A, int8_t *B, int32_t *C, int M, int N, int K) {
  // 第一层循环：遍历 C 矩阵的子矩阵行
  for (int i = 0;i < M / 4;i ++) {
    // 第二层循环：在 C 矩阵的子矩阵行里遍历子矩阵
    for (int j = 0;j < N / 4;j ++) {
      // 此时通过两层循环，会遍历 C 矩阵的所有子矩阵 Cij
      for (int k = 0;k < K / 4;k ++) {
        // 搬运 Bkj 子矩阵到 NPU
        npu_transfer_l0b(B + (k * N + j) * 4, N);
        // 执行 LDTILE 指令并等待其完成
        npu_ldtile();
        // 搬运 Aik 子矩阵到 NPU
        npu_transfer_l0a(A + (i * K + k) * 4, K);
        // 如果 k > 0，启用累加模式
        bool acc = k > 0;
        // 执行 TMUL 指令并等待其完成
        npu_tmul(acc);
      }

      // 此时 Cij 已经计算完成，搬运回 SRAM
      npu_transfer_l0c(C + (i * N + j) * 4, N);
    }
  }
}

可以看到，这个计算过程和经典的三层循环矩阵乘法是差不多的，只是把元素的乘积，变成了在 NPU 上实现的矩阵乘法。

第一版矩阵乘可视化模拟器

第一版矩阵乘的交互式可视化模拟器：

打开模拟器（新窗口）

第二版通用矩阵乘法

但是，上述第一版通用矩阵乘法也有一些明显的性能问题：

  1. 数据搬运次数过多：A 和 B 的子矩阵，各需要搬运 (M / 4) * (N / 4) * (K / 4) 次；C 的子矩阵需要搬运 (M / 4) * (N / 4) 次
  2. LDTILE 命令执行次数过多：每次 TMUL 前都要进行一次 LDTILE 指令，缺少数据复用，如果一次 LDTILE 可以用于多次 TMUL 运算，就可以节省 LDTILE 指令的开销

围绕着这两点问题，可见核心就是要复用。考虑到 LDTILE 可以把 B 的子矩阵加载到脉动阵列里，然后进行多次 TMUL 运算，就来分析一下，一次 LDTILE 可以对应哪些 TMUL 运算。还是上面的例子，假设 A、B、C 矩阵都分成了四个子矩阵：

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

假如 LDTILE 指令加载 B00 矩阵到了脉动阵列当中，那么它可以做哪些运算呢？根据定义：

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

B00 会与 A00、A10 进行运算，分别得到 C00、C10 结果的一部分。这样就意味着，一次 LDTILE 加载 B00，可以做两次 TMUL，第一次是 C00 = A00 * B00，第二次是 C10 = A10 * B00；接着，再一次 LDTILE 加载 B10，再做两次 TMUL，C00 += A01 * B10 和 C10 += A11 * B10。这样就完成了 C00 和 C10 的计算，再对 B01、B11 故技重施，就完成了整个计算过程。这样，LDTILE 指令的执行次数从原来的 8 次降低到了 4 次，B 子矩阵的搬运次数也是同理。

通过这种方式，就节省了 LDTILE 指令的执行次数和 B 子矩阵的搬运次数，不过代价就是需要更大的 L0A/L0C 片上存储空间。由于 L0A/L0C 空间是有限的，可能不能一下把 C00、C10、C20 等等一系列子矩阵都存下来。不妨设 L0C 和 L0A 都能存下来 2 个子矩阵，就得到了第二个版本的通用矩阵乘法：

  1. 把 A、B、C 矩阵都划分成若干个 4x4 的子矩阵；
  2. 第一层循环：遍历 C 矩阵的子矩阵行，但这次是两行一起处理，i 从 0 到 M/4，每次增加 2，这里没考虑 M/4 是奇数的情况；
  3. 第二层循环：在 C 矩阵的子矩阵行里，遍历子矩阵，j 从 0 到 N/4，此时 Cij 就是 C 矩阵里的一个 4x4 子矩阵，内层循环需要把 Cij 和 C(i+1)j 都一起计算出来；
  4. 第三层循环：针对矩阵乘法的公式，Cij = Sum_k(Aik * Bkj)，遍历 k，k 从 0 到 K/4；
  5. 第三层循环内，每次循环做如下工作：
    1. 把 Bkj 子矩阵从 SRAM 搬运到 L0B；
    2. 执行 LDTILE 指令，把 Bkj 子矩阵加载到脉动阵列；
    3. 把 Aik 子矩阵从 SRAM 搬运到 L0A；
    4. 执行 TMUL 指令，对于第一次内层循环 k=0，执行 Cij = Ai0 * B0j，即 acc = 0；此后的内层循环 k>0，执行 Cij += Aik * Bkj，即 acc = 1；
    5. 把 A(i+1)k 子矩阵从 SRAM 搬运到 L0A；
    6. 执行 TMUL 指令，对于第一次内层循环 k=0，执行 C(i+1)j = A(i+1)0 * B0j，即 acc = 0；此后的内层循环 k>0，执行 C(i+1)j += A(i+1)k * Bkj，即 acc = 1；
  6. 内层循环结束后，把 Cij 和 C(i+1)j 子矩阵从 L0C 搬运回 SRAM。

对应的 C 代码（部分细节略去）：

// A: MxK, B: KxN, C: MxN
void npu_matmul_v2(int8_t *A, int8_t *B, int32_t *C, int M, int N, int K) {
  // 第一层循环：遍历 C 矩阵的子矩阵行，但是两行两行一起处理
  for (int i = 0;i < M / 4;i += 2) {
    // 第二层循环：在 C 矩阵的子矩阵行里遍历子矩阵
    for (int j = 0;j < N / 4;j ++) {
      // 此时通过两层循环，会遍历 C 矩阵的所有子矩阵 Cij 和 C(i+1)j
      for (int k = 0;k < K / 4;k ++) {
        // 搬运 Bkj 子矩阵到 NPU
        npu_transfer_l0b(B + (k * N + j) * 4, N);
        // 执行 LDTILE 指令并等待其完成
        npu_ldtile();
        // 搬运 Aik 子矩阵到 NPU
        a_base = 0;
        npu_transfer_l0a(A + (i * K + k) * 4, K, a_base);
        // 如果 k > 0，启用累加模式
        bool acc = k > 0;
        c_base = 0;
        // 执行 TMUL 指令并等待其完成
        npu_tmul(acc, a_base, c_base);
        // 搬运 A(i+1)k 子矩阵到 NPU
        a_base = 4; // 避免覆盖 Aik 矩阵的内容
        npu_transfer_l0a(A + ((i + 1) * K + k) * 4, K, a_base);
        // 执行 TMUL 指令并等待其完成
        c_base = 4; // 避免覆盖 Cij 矩阵的内容
        npu_tmul(acc, a_base, c_base);
      }

      // 此时 Cij 和 C(i+1)j 已经计算完成，搬运回 SRAM
      c_base = 0;
      npu_transfer_l0c(C + (i * N + j) * 4, N, c_base);
      c_base = 4;
      npu_transfer_l0c(C + ((i + 1) * N + j) * 4, N, c_base);
    }
  }
}

第二版矩阵乘可视化模拟器

第二版矩阵乘的交互式可视化模拟器：

打开模拟器（新窗口）

更多优化与算子库

实际上，通用矩阵乘法还有更多的实现方法，根据硬件的设计和矩阵的不同规格，可能会出现不同的最优软件实现。例如，上述的版本里，A 矩阵并没有被复用，同时脉动阵列实际上支持 Mx4x4 的矩阵乘，其中 M 大于 4，增大 M 可以减少脉动阵列的气泡等等。

那么，这些优化过的矩阵乘法实现，其实就是算子：把常用的操作（比如矩阵乘），在特定硬件和矩阵规格上进行调优，然后把算子打包成一个现成的软件交给上层软件去使用，那么这些算子就可以让上层软件在不知道硬件设计的情况下，也能发挥出硬件的性能。这里的这两版代码，在提供的 npu.c 里也有具体实现，当然具体实现里会有更多的细节，因此部分函数的定义和含义和上面并不完全相同。

驱动 API

cod-npu 仓库根目录下的 software/npu.h 中提供了以下接口（完整定义以头文件为准）：

函数 作用
npu_fill_l0a(A, stride, a_base) 把一个 4x4 的 A 子块写入 L0A
npu_fill_l0b(B, stride) 把一个 4x4 的 B 子块写入 L0B（固定写到 L0B 字 0..3）
npu_ldtile() 执行 LDTILE，把 L0B 中的 B 子块加载到脉动阵列
npu_tmul(a_base, c_base, m_rows, acc) 执行 TMUL：C += A x B，m_rows 为行数，acc 选择是否累加
npu_read_l0c(C, stride, c_base) 从 L0C 读回一个 4x4 的 C 子块
npu_busy() / npu_done() / npu_clear_done() 查询/清除状态
npu_matmul_v1(A, B, C, M, N, K) 第一版通用矩阵乘，返回 LDTILE/TMUL 次数
npu_matmul_v2(A, B, C, M, N, K) 第二版通用矩阵乘（复用 B 子块）
npu_matmul_sw(A, B, C, M, N, K) CPU 标量参考实现，用于对拍

示例代码与真实 API 的差异

上面的教程为了突出算法，使用了 npu_transfer_l0a、npu_tmul(acc) 等简化写法；真实驱动中的函数名和参数顺序略有不同，例如 npu_tmul(a_base, c_base, m_rows, acc)。以仓库根目录下的 software/npu.h 为准。

先写对，再写快

建议先用 npu_matmul_v1 和 npu_matmul_sw 对拍，确认 4x4、8x8、16x16、32x32 全部正确，再尝试 v2 和后续优化。调试时可以先固定一组小矩阵，把 L0A/L0B 的内容通过串口打印出来，逐行核对。

作者：Jiajie Chen