AI芯片：Google TPU分析

In-Datacenter Performance Analysis of a Tensor Processing Unit

TPU系列芯片已经广泛应用于Google的各种服务中，加速其神经网络运算，例如2015年AlphaGo与李世石进行围棋人机对战中所用的就是TPU芯片。TPU1面向的是数据中心的推理应用，被设计为PCIe上的协处理器,可以像GPU一样插入到SATA硬盘插槽里使用，实现了12.5GB/s的有效带宽。

一、TPU1架构介绍

  

Weight FIFO：能容纳4个tile深的权重，主要负责从8GB的off-chip DRAM上读取权值；
Unified buffer：24MiB大小的片上缓存单元。用于存储输入数据（DMA传输）和产生的输出数据；
MMU：256 x 256 x 8bit的乘加单元，脉动阵列形式。每周期可以计算出256个乘法结果，计算出的结果是16bit的。矩阵单元中包含一个64KB的weight tile，以及一个双缓存单元；
Accumulators：4MiB，4098 x 256 x 32bit大小的累加单元，用来收集乘法产生的16bit结果（4096：从roofline模型得知，单次操作需要2048个单元才足够，双缓存则需要4096个单元才能满足需求）。

控制指令：指令通过PCIe总线从host主机端传输到TPU中，以4级流水线执行。采用CSIC指令集类型，总共有12条指令，下面列举几条关键指令：
1.Read_Host_Memory：CPU host->Unified Buffer
2.Read_Weights：Weight DRAM->Weight FIFO
3. MatrixMultiply /Convolve：执行乘法或卷积运算。举例：input（B * 256）*weight（256 * 256）=>output（B * 256）。
4.Activate：执行激活操作（Relu，Sigmoid）
5.Write_Host_Memory：Unified Buffer->CPU host

为什么要选择CISC指令集呢？：RISC指令集属于精简指令集，一条指令一般只花费一个时钟周期；CISC指令集属于复杂指令集，一条指令的执行可能需要花费1000个时钟周期。观察上面的运算指令，显然不是一个时钟周期就可以执行完成的。

流片：片上缓存Unified buffer和MMU乘法运算单元占用了2/3的空间。

二、脉动阵列

TPU1核心计算单元采用的是256*256大小的脉动阵列，属于数据流处理器。

脉动阵列相比于传统架构的两个优点
1.计算单元PE间共享数据，可以缩短关键路径，从而提高时钟周期；
2.SRAM读取次数减少，节省能耗。

脉动阵列示意图如图3所示，权值数据weight是预先加载的，输入数据input从左侧进入，输出数据从下侧输出。

三、CPU、GPU和TPU1服务器对比

TPU同CPU和GPU一样，需要具有一定的通用性。因此，Google为其设计了软件栈，可以将TensorFlow图转化为在TPU上运行的指令。这里也能看出，软件生态环境对于芯片发展具有非常重要的环境。下面表格是处理器周期运算量的数量级对比，图片是服务器上部署三种处理器的参数。

处理器	周期运算量数量级
CPU	几个
CPU（SIMD）	几十个
GPU	几万个
TPU	几十万个

四、实验结果

参数	TPU1
工艺	28nm
面积	<331mm2
频率	700MHz
运算单元	65536个8bit
性能	92Top/s
功耗	40W
片上缓存（UB）	(24+4=28)MiB

性能计算公式：92Top/s=700MHz * 65535 *2
其中65536是每周期执行的乘法操作数，乘2的原因是在MMU中同时还会有65536次加法。脉动阵列平稳运行时，每周期可以计算出256个乘法结果。

片上缓存：TPU1选择将片上缓存做的较大，从而节省片外访存消耗。

量化：训练阶段是32bit-float point，推理阶段是8bit-integer。这算是TPU一个值得注意的点。前面大多数研究中，推理阶段的量化只是到了16bit。

五、发展

名称	时间	性能	应用
TPUv1	2016年	92Tops	数据中心的推理应用
TPUv2	2017年	180TFlops（集成4块芯片），64GB（HBM）	数据中心的训练和推理，引入到Google Cloud中
TPUv3	2018年	420TFlops，128GB（HBM）	数据中心的训练和推理（进一步扩展）
Edge TPU	2018年	可处理高吞吐量的流式数据	用于loT设备中
TPUv2 pod	2019年	11.5千万亿次浮点运算/s，4TB（HBM），二维环面网状网络	深度学习领域
TPUv3 pod	2019年	>100千万亿次浮点运算/s，32TB（HBM），二维环面网状网络	深度学习领域

PS：高带宽内存(HBM)是一种用于支持内存设备数据吞吐量的高性能接口，其性能远超常规形式的内存。

Finished.