Memory Wall in Neural Network Inference

news2024/9/21 10:35:50

Memory Wall in Neural Network Inference

神经网络推理的瓶颈在于访存带宽，通常无法发挥出加速器的全部算力。本文总结了目前常用的推理加速器及其设计，并分析了常用神经网络的访存瓶颈。文章大部分内容参考自Computer Architecture: A Quantitative Approach。

1 Compute centric accelerators

1.1 CPU

一般来说，CPU擅长于做控制逻辑复杂的workload。但是，CPU里面也存在SIMD的加速器设计。

Vector Architecture in RISC
- Basic idea
  - Read sets of data elements into “vector registers”
  - Operate on those registers
  - Disperse the results back into memory
- Y = a × X + Y
  
  左边使用Loop的原始汇编代码，右边是使用Vector Architecture优化后的汇编代码。
  - 如果 X 和 Y 均为32维向量，1 lane
    
    那么使用Vector Architecture加速，需要32×3个cycle
- Optimization
  
  可以优化的地方很多，讲两个有意思的优化。
  - Multiple Lanes
    
    如果是4 lane的话，上述例子仅需要32×3/4个cycle就可以完成。
  - Scatter-Gather
    
    读取的向量可以不存储在连续的内存中。
    
    例如，执行这段代码：
```
for (int i=0; i<n; i++){
    A[K[i]]=A[K[i]]+C[M[i]];
}
```
    使用Scatter-Gather优化后的汇编代码如下：
    
    这种优化技术用在稀疏矩阵场景很有用。
SIMD extensions for x86

简单来讲，就是利用长寄存器，一次处理多个数据。比如寄存器长度是256 bit，一个数据的长度是64 bit，这时候，就可以一次处理4个数据。

使用SIMD extensions优化Y = a × X + Y 后的汇编指令：
- 这种优化要求向量必须存储在连续的内存空间中
Vector Arch. vs. SIMD Extensions

看这张图，对于神经网络推理加速场景，x86完全干不过RISC（哈哈，连高性能计算都干不过RISC，x86是不是要凉啊）。

1.2 GPU

简单来讲，GPU的架构就是堆很多的寄存器和ALU，获取高算力。对应到下图，你可以看到很多的SIMD Lane 和 DP Unit。

在这里插入图片描述

但是，GPU里面的处理器算力很强，访存带宽跟不上怎么办？首先，它使用了HBM，内存带宽很高；另外，其中的SIMT (Single Instruction, Multiple Thread) 编程模型利用多线程来掩盖访存时延。

如下图，简单来讲，当thread 8运行的时候，其他的线程（比如thread 1, 3）也没闲着，它们正在取数据。这样的话，GPU里面的处理器时刻处于满负载状态，不需要傻傻地等到数据加载到寄存器才运行。（熟悉GPU的同学应该发现了，图里面的thread改成warp更合适，懒得改了）

在这里插入图片描述

1.3 TPU

听过好多人吹google的TPU多么牛逼，其实也就那样，就是一个专门用来加速GEMM的设备。下图是TPU的架构图，右边的Matrix multiply unit 就是TPU的核心部件。

在这里插入图片描述

通用的加速器GPU可以用来加速向量加法，甚至有人整活，直接在GPU上构建了一个数据库。然而，TPU是一个专用的加速器，只能用来加速GEMM算子，其他的活统统干不了。

不过，TPU里面加速GEMM的硬件设计很有意思。专业的名字叫做Systolic Execution，不过，我喜欢叫它脉冲阵列。如果矩阵的维数是N×N，TPU使用2N-1个cycle就可以把GEMM算子计算完。

在这里插入图片描述

2 Neural network and Memory wall

在当前以计算为中心的编程模型中，需要先从内存中把数据加载到cache，然后ALU进行计算任务，最后将结果保存至内存中。

目前，ALU算的很快，把内存中数据加载到cache的带宽有限，这时候，ALU的算力就会被浪费，性能瓶颈卡会在内存带宽上，这就是“内存墙”。如果我们读入cache一个数据，对它进行N次计算操作后才保存结果，N定义为Operational Density。如果Operational Density，就没有“内存墙”的问题。如果Operational Density比较低，程序就会卡在“内存墙”上。

接下来，我们分析一下MLP、CNN、LSTM和Transformer的Operational Density，判断它们是否会卡在内存墙上。