FPGA GEMM算子设计

1. 1. 前言
2. 2. 基础版本
3. 3. 并行和流水线优化
   1. 3.1. Version 2
   2. 3.2. Version 3-5
   3. 3.3. Version 6
   4. 3.4. Version 7
      1. 3.4.1. Version 7.1
      2. 3.4.2. Version 7.2
      3. 3.4.3. Version 7.3
4. 4. 精度变化
   1. 4.1. Version 8

前言

因为项目需要，我尝试为FPGA设计了一个GEMM算子。这篇博客记录了算子的一部分迭代历程（因为项目才刚开始(❁´◡`❁)）

我们讨论的是矩阵运算：

$$\mathbf C = \mathbf A \mathbf B$$

我们先规定形状：

• C: $M \times N$
• A: $M \times D$
• B: $D \times N$

这个 GEMM 算子是 transformer 计算的一部分，所以M就是输入token数量，N是输出神经元的数量。

本文将只阐述整个设计思想的演进，不涉及具体的代码实现。

基础版本

按照矩阵乘法的定义，很容易想到最简单的解法，即依次取出A和B中的一行和一列，做内积，这样就可以求出C中的一个元素。

$$C_{i, j} = \sum_{k=0}^{D-1} A_{i, k} B_{k, j}$$

这种做法导致每次内积计算都需要从HBM加载A和B指定行列的权重，显然是十分低效的；并且没有做并行，计算速度非常缓慢。实际测试中计算速度甚至没有达到0.1GFLOP/s.

并行和流水线优化

既然一次只计算一个数效率低，那么我们就来做流水线并行。
但在此之前，我们需要增加单次读取到uram/bram的权重数量。

Version 2

单次从A中读取一块 $8 \times 64$ 的区域（图中紫色区域），从B中读取一块 $64 \times 16$ 的区域（图中紫色区域），由于这是早期的探索，为了因此没有完整读入D。

紫色窗口在A和B中分别按行和按列滑动，当窗口完整滑过D个元素时，就完成C中的一个 $8\times 16$的区域的计算。

在紫色窗口内部，这个版本还没有进行并行，因此，GEMM会先计算A的第0行乘B的第0列，再计算A的第1行乘B的第1列…以此类推。。

Version 3-5

观察发现，v2当中，每一次内积中每一次乘加所需的周期数是大于一的，实测大约是3-4拍，这意味着可以让单个乘法和加法操作做流水线计算。在这个版本中仍然只有一个乘加通路，它通过将单一一个内积计算的II变为8，将单个流水线完全填满。

Version 6

在这个版本中，我们引入了多级缓存

• HBM: 存储完整权重
• URAM：存储多块大块权重
• BRAM：存储当前正在计算部分的权重，即类似v2中的紫色区域
• LUT：暂存C中的计算结果

在URAM中我们开辟了多个64行/64列的区域用于存储大块的A和B，URAM和HBM之间采取乒乓缓存策略，当某一个URAM块完成计算后，将并行从HBM读取新的权重到这个块上。生产者是HBM，消费者是BRAM。BRAM会消费URAM块上的内容，每执行一个紫色计算块中的计算，BRAM将从URAM中读取相应的数据到BRAM的A、B缓冲区内，供计算核心消费。

Version 7

在做好了缓存的准备后，在这个版本中我们终于引入了并行计算。
仔细观察可以发现，在version3-5的图中，越靠后启动的列启动前浪费的周期就越多。只要将内积的计算通路增加一倍即可缩短一半的等待时间。同时也可以缩短同一行的两次乘加计算的启动间隔。当然内积通路也不是越多越好，若

$$\frac{\text{列数}}{\text{内积通路数量}} = \text{各内积通路负责的列数} \le \text{单次乘法计算所需周期数}$$

那么必然会导致周期和计算资源的双重浪费。因此我们首先考虑了

$$\text{内积通路数量} = \frac{\text{列数}}{4}$$

因为测得的乘法计算时长大约是4个周期。
下图是理想情况下A中各行的计算时序图，每一个深黄色方框表示A的一行的内积计算中的一次乘法操作结束，对应的加法操作开始。

对比发现，这样设计之后每一次计算的等待时间明显减少，计算效率明显提高。

Version 7.1

我们尝试扩大了单个紫色计算块的大小，从8x16扩大到8×32，16x16和16x32。但v7.2-v7.3是并行实现的，所以仍采用8x16的核心大小

Version 7.2

我们正式引入了「多核」的概念。
这里的一个核不是多个 GEMM kernel，而是指前面的紫色计算块。我们现在正式将其命名为「计算核」。

n个核可以同时处理不同的计算区域，从而理论上将计算速度扩大n倍（在不考虑布线资源等各种杂七杂八因素的情况下）

Version 7.3

在7.2的基础上，这个版本的BRAM上的A和B也实现了乒乓缓存。

精度变化

Version 8

在Version8当中，我们

• 将计算精度从fp32改为fp16乘法，fp32加法
• 增大了HBM读取的最大位宽到512bits
• 为了节省片上存储资源将A和B的URAM缩减为各两块
• 尝试增加DSP的使用率，降低LUT的使用率。
  初次尝试使用的是 $16 \times 64$的计算核大小，四个计算核，但最终布线资源不足失败了，因此有了8.1和8.2两个缩小版。
  实际测试下8.1和8.2的资源占用几乎没有差异，8.1略有优势，但大部分少于1各百分点。实测$4096 \times 4896 \times 4864$大小的矩阵，计算速度中位数分别为179.33GFLOPS/s和179.32GFLOPS/s。性能几乎没有差异。