\

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

作者:

分享  

Share on weixin
Share on weibo
Share on linkedin

订阅

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

图1 GPU HBM bandwidth

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

片上分布式存储架构

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

图2 Bow IPU architecture

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

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

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

Tile架构

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

图3 Tile architecture

名词解释:

  • MAIN:用于执行memory load / store指令、整型运算指令、control flow指令的单元
  • AUX:用于执行浮点运算指令的单元
  • TILE MEMORY:Tile的local memory
  • MRF:MAIN单元的Register file
  • ARF:AUX单元的Register file
  • AMP:专门用于执行矩阵乘和卷积的硬件单元(accumulating matrix-vector product)

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

硬件线程(hardware context)

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

图4 instruction pipeline

instruction bundle

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

图5 instruction bundle
图6 example of 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大小。

图7 Memory architecture for Bow IPU

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

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

图8 同时执行两个64-bit 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操作。

图9 128-bit load

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

图10 同时执行两个64-bit load以及一个64-bit store

triple-pack address

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

图11 triple-pack address

其它技巧

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通过硬件支持零开销循环指令rpt,使得循环指令更高效,无需进行 unroll操作。

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

IPU访存带宽计算

通过以上的技术解析,我们可以看到,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直接搬到片上运行。

参考

More Posts

加速部署!Graphcore携手百度飞桨成功实现FastDeploy适配

新增FP8!Graphcore推出新一代PCIe加速卡C600 

两项第一!Graphcore在OGB-LSC中取得佳绩

稀疏化80%!Graphcore携手Aleph Alpha解锁人工智能计算效率新突破

助力科研!阿贡国家实验室新引入Graphcore Bow IPU人工智能系统

开始使用Poplar Triton后端

获取最新的GRAPHCORE资讯

在下方注册以获取最新的资讯和更新:




    获取最新的GRAPHCORE资讯

    在下方注册以获取最新的资讯和更新: