IPU如何做到访存带宽65TB/s

AI任务具有计算密集型和访存密集型并存的特点。针对计算密集型的特点，现代AI处理器一般通过增加特有的硬件处理单元（如NVIDIA GPU的Tensor Core单元，或者Graphcore拟未 IPU的AMP单元）来处理如MatMul或Convolution等计算，而对于访存密集型的特点，现代AI处理器一般会不断提高片外存储的带宽。当前访存带宽最高的存储介质是HBM^[1]（high bandwidth memory），在NVIDIA最新架构Hopper中应用的HBM3，带宽已经达到3TB/s，刷新了AI处理器片外存储带宽的记录。访存带宽通常是制约整个模型性能最关键的因素。

那么，拟未IPU处理器65TB/s的访存带宽又是如何实现的呢？下面本文就来为大家揭晓答案。

本文拿拟未最新的Bow IPU举例，下面是Bow IPU的架构图。

IPU芯片没有高速的片外存储，而是把存储放到了片上。整个芯片由1472个核心组成，我们称作Tile。每个Tile由独立的计算单元和存储单元组成，整个片上存储是分布式的。每个Tile中有624KB的SRAM，因此整个芯片的存储大小为624KB * 1472 = 900MB。

IPU芯片采用纯分布式的架构，每个Tile有自己的存储和计算资源，采用MIMD^[2]的计算架构（与NVIDIA CUDA的SIMT不同），每个Tile可以独立地执行不同的指令，可以独立地访存。Tile和Tile之间的memory不能共享访问，只能访问自己Tile内部的memory，我们叫做local memory。因此整个芯片的访存带宽 = Tile 访存带宽 * Tile数量。

下面我们介绍单个Tile的访存带宽。

IPU的Tile是一个高确定性、异步双LIW处理器，支持多个硬件线程来做指令流水，掩盖指令执行的时延。每个Tile包含一个local memory，用于存储程序执行所需要的代码和数据。

下面分几个子章节来陈述一下Tile架构的实现细节。

和其它典型处理器架构一样，IPU Tile每条指令的执行分为多个阶段，每个阶段消耗1个时钟周期。每个Tile中有6个worker线程组成指令流水^[3]，来掩盖指令执行的延时。通过指令流水，大部分的指令都可以在1个时钟周期执行。

IPU Tile有两个独立的执行单元，MAIN和AUX，MAIN单元主要负责执行内存的读写指令，以及整型计算指令；AUX单元主要负责执行浮点计算指令。IPU Tile通过采用VLIW^[4] （Very long instruction word）技术，可以实现1个时钟周期并行执行MAIN和AUX指令。通常我们写汇编代码时，会把MAIN指令和AUX指令用一个括号括起来，称作一个instruction bundle，代表这两个指令是并行执行的。通过这种技术，Tile的访存和浮点运算可以并行执行。

IPU目前没有复杂的内存分层等级（譬如NVIDIA GPU的Global memory、L2 cache、L1 cache），只有一层local memory可用，降低了软件设计难度。

IPU的每个Tile都包含一个紧密耦合的local memory，用于存储执行上下文的所有代码和数据。Tile的local memory是唯一可由Tile指令集直接访问的memory。

Tile的内存地址用21位表示，因此Tile的寻址空间为2MB，从地址 0x0 开始，每个线程都可以看到整个空间。在实现时，local memory只映射到其中一部分寻址空间，比如Bow IPU的local memory映射到寻址空间0x4C000~0xE7FFF，共计624KB大小。

local memory被组织为两个region，region0和region1，每个region由多个64-bit的bank组成。IPU Tile的内存读写接口共有三个，其中有两个64位的读接口，一个64位的写接口。三个接口可以并行寻址。为了不产生memory conflict，并行寻址时需要在不同的memory bank。

例如，通过ld2x64pace 指令，可以在一个时钟周期内执行两个64-bit的memory load操作。

region1比较特殊，其中的memory是由两个相邻的64-bit memory bank组成一个128-bit 的memory element，因此可以直接进行128-bit的寻址（例如ld128指令）。由于在连续寻址128-bit时，一个memory element中的奇偶bank是交错的选择的，因此这个region中的memory也叫做interleaved memory。

例如下图表示通过ld128指令做128-bit memory load操作。

Tile的两个读接口和一个写接口可以同时操作，这样就可以在一个时钟周期内进行3个64-bit访存操作。例如ld2xst64pace指令，在一个时钟周期内，允许执行两个64-bit的load的同时，执行一个64-bit的store。

需要说明的是，为了在一条指令内操作三个地址，拟未在设计指令时采用了一个小技巧，因为IPU Tile的寻址空间是21位，因此我们可以把三个21位地址封装成1个64位数，存储到两个32位寄存器里，我们称作triple-pack。通过triple-pack，我们可以在1条指令里同时操作三个地址。

AI计算中常用的二元浮点运算，都是通过以下流程计算（通常参数和结果都用Tensor表示，是在memory中连续存储的数据）：

1. 从memory address 0中load参数1到通用寄存器

2. 从memory address 1中load参数2到通用寄存器

3. 执行浮点计算指令，结果存储到通用寄存器

4. 结果从通用寄存器中store到memory address2

5. 把参数和结果要操作的address + 1，计算下一组数据

6. 循环以上

Tile的很多寻址指令，例如ld2xst64pace和ld2x64pace等带pace后缀的指令，都支持自动对基地址进行加偏移的操作。

另外，IPU Tile通过硬件支持零开销循环^[5]指令rpt，使得循环指令更高效，无需进行 unroll操作。

通过三个接口并行寻址技术，instruction bundle技术，以及以上提到的一些硬件特性，可以把上述6步操作在一个时钟周期做完。

通过以上的技术解析，我们可以看到，IPU的Tile可以在一个时钟周期内对local memory进行3次64-bit（8 bytes）的访问（例如ld2xst64pace指令），因此整个IPU芯片的访存带宽可以通过以下公式计算：

Bandwidth= 3 * 8(byte) * Freq(Hz) * numTiles

例如Bow IPU的memory bandwidth等于：

Bandwidth(Bow IPU) = 3 * 8 * 1850(MHz) * 1472 = 65TB/s

同时，通过instruction bundle技术，在访存的同时，还可以并行进行浮点运算。

由于IPU拥有900MB的超大片上memory，因此能够把整个模型放到片上运行；由于IPU拥有超高的65TB/s的访存带宽，因此能够打破memory wall的限制，极大地释放AI的算力。另外，模型性能优化也被大大简化了，在传统AI硬件上常用的kernel fusion优化，在IPU上通常是不需要做的，因为IPU的运行模式相当于kernel fusion的极限：整个模型融合成一个kernel直接搬到片上运行。

• 拟未官网：

https://www.graphcore.cn/

• 拟未文档：

https://docs.graphcore.ai/en/latest/

• IPU ISA文档：

https://docs.graphcore.ai/projects/isa/en/latest/index.html

• Poplar编程指导：

https://docs.graphcore.ai/projects/poplar-user-guide/en/latest/index.html