转载Google TPU论文

选自 Google Drive

作者：Norman P. Jouppi 等

痴笑@矽说 编译

该论文将正式发表于 ISCA 2017

从去年七月起，Google就号称了其面向深度学习的专用集成电路（ASIC）产品——Tensor Processing Unit （TPU），然而其神秘面纱一直未被揭开。直至本周，Google公开了其向ISCA（国际计算机体系架构年会）投稿的的预录取论文——In Datacenter Performance Analysis of a Tensor Processing Unit，TPU的技术细节才公开发表，令我们才有幸见识其真面目。虽然可能只是“犹抱琵琶半遮面”，但其作为CPU/GPU/FPGA后的另一深度学习选项，特别是TPU和tensorflow间可能存在的微妙联系值得我们特别关注。矽说在第一时间选编、翻译 了其中的重要部分，以飨读者 。

论文地址： https://arxiv.org/abs/1704.04760?context=cs.AR

前言

文章具体描述了Tensor Processing Unit (TPU)的体系结构，并与目前主流的CPU（Intel Haswell Xeon）和GPU（Nvidia K80）的性能做出了比较，采用的benchmark包含了CNN，RNN（LSTM）和全链接（MLP）神经网络。其特点包括：

（1） 面向inference的专用app与硬件，强调了吞吐率上的性能

（2） TPU的不仅在面积和功耗上低于GPU，而且在乘累加的数量和存储器容量是K80的25倍和3.5倍

（3） TPU的速度上的优势明显，达到GPU和CPU的15到30倍

（4） 在6个NN架构中，4种神经网络的性能瓶颈在于存储器带宽。若存储器带宽达到K80的性能，作者相信性能能提升到30到50倍

（5） TPU的能效值（TOPS/W）达到在目前其他产品的30到80倍

（6） CNN在TPU中工作量的比列只有5%

起源、架构与实现

早在2006年开始，Google就开始讨论在数据中心部署GPU，FPGA或定制ASIC。在2013年，当DNN已经展露头脚，以致可能会使我们的数据中心的计算需求加倍时，传统的CPU已经被认为不合时宜了。因此，我们开始了一个高度优先的项目，以快速生成用于Inference的定制ASIC，（训练仍使用GPU）。 目标是将性能提升10倍以上。 根据这一任务，Google的TPU的设计和验证在短短15个月内完成，并构建并部署在数据中心。

而不是与CPU紧密集成，为了减少延迟部署的可能性，TPU被设计为PCIe I/O总线上的协处理器，允许它像GPU那样插入现有的服务器。 此外，为了简化硬件设计和调试，主机服务器发送TPU指令来执行，而不是自己提取它们。 因此，TPU在精神上比FPU（浮点单元）协处理器更接近于GPU。

TPU的目标是运行整体深度学习的神经网络模型，以减少与主机CPU的交互，并且具有足够的灵活性，以满足2015年及其后的NN需求，而不仅仅是2013年NN所需。

图1是TPU的整体体系架构。TPU指令通过PCIe Gen3 x16总线从主机发送到指令缓冲区。 内部块通常通过256字节宽的路径连接在一起。 从右上角开始，Matrix Multiply Unit是TPU的核心。 它包含256×256个乘累加单元（MAC），可以对有符号或无符号整数执行8位乘法和加法。 求和单元（Accumulator）的输入为16位乘积，输出和的位宽为32，共有4906个，每个的输入个数为256。 因此，矩阵单元每个时钟周期产生一个256元素的部分和。

当使用8位权重和16位激活函数的混合计算结构时，MMU的计算速度将减半，当它们都是16位，计算速度将减为四分之一。 它每个时钟周期读取和写入256个值，并且可以执行矩阵乘法或卷积。 矩阵单元采用两个64KiB的权重tile的双缓冲设计，其中的一个仅在非稀疏模式下才会被激活，这样就提升了TPU对于稀疏网络的性能支持。Google相信稀疏性将在未来的设计中占有优先地位。

矩阵单元的权重通过片上Weight FIFO进行分级，该FIFO从片外8 GiB DRAM读取（在inference中，权重是只读的）。Weight FIFO存储了四个4个tile。 中间结果保存在24 MiB片上Unified Buffer中，可作为未来的Matrix单元的输入。 可编程DMA控制器向CPU主机内存和Unified Buffer传输数据。

图2显示了TPU芯片的布局图。 24 MiB Unified Buffer几乎是芯片的三分之一，Matrix Multiply Unit是四分之一，因此数据路径占到了整个芯片的三分之二。由于开发时间短，部分设计选取了简单的值以简化编译器设计。 控制逻辑占到总面积的只有2％。 图3显示了搭载在PCB上的TPU实现图，期接口类似SATA磁盘，可通过PCIe直接接入数据中心。

TPU的指令是通过PCI额发射的，遵循CISC传统，包括重复字段。 这些CISC指令的每个指令的平均时钟周期（CPI）在10到20之间。它总共有大约十多个指令，这里罗列最关键的五个：

（1） Read_Host_Memory

（2） Read_Weights

（3） MatrixMultiply/Convolve

（6） Activation

（5） Write_Host_Memory

其他指令是备用主机内存读/写，设置配置，两个版本的同步，中断主机，调试Jtag，空操作和停止。 CISC MatrixMultiply指令为12个字节，其中3个为Unified Buffer地址; 2是累加器地址; 4是长度（有时是卷积的2个维度）; 其余的是操作码和标志。TPU微架构的理念是保持MMU的繁忙。 因此，MMC的CISC指令使用4级流水线结构，其中每条指令在其中的单独一级执行。 其目标是通过将其执行与MatrixMultiply指令重叠来隐藏其他指令（如Read_Weights等）。但，当激活的输入或权重数据尚未就绪，矩阵单元将进入等待模式。

由于读取大型SRAM数据的功耗比算术功耗高的多，所以MMU通过systolic方式减少对Unified Buffer的方位，节省功耗。 图4显示了TPU数据流程，输入从左侧流入，权重从顶部加载。 给定的256元乘法累加运算通过矩阵作为对角波前移动。 此过程中权重是预先加载的，并且与数据同步。通过对控制逻辑和数据流水线操作，使得对于MMU而言，256个输入一次读取的，并且它们立即更新到256个累加器中的一个个具体位置。 从正确的角度来看，软件不知道硬件的systoli特性，但是却需要对于systoilc引起的延时效果有正确的评估。

TPU软件堆栈必须与为CPU和GPU开发的软件栈兼容，以便应用程序可以快速移植到TPU。在TPU上运行的应用程序的一部分通常写在TensorFlow中，并被编译成可以在GPU或TPU上运行的API。像GPU一样，TPU堆栈分为用户空间驱动程序和内核驱动程序。内核驱动程序是轻量级的且长期稳定的，只处理内存管理和中断。用户空间驱动程序频繁更改。它设置和控制TPU执行，将数据重新格式化为TPU命令，将API调用转换为TPU指令，并将其转换为应用程序二进制文件。用户空间驱动程序在首次评估模型时编译模型，缓存程序映像并将重量映像写入TPU的重量存储器; 第二次和以下评估全速运行。 对于大部分模型，TPU可以从输入到输出完整实现。实际操作中，同一时间通常每次进行一层的MMC计算，而其非关键路径操作将被隐藏在MMU操作下。

CPU、GPU与TPU的性能比较

在TPU与CPU和GPU的性能比较部分，Google选择了两款并非最新但具有代表性的平台——Intel Haswell架构的Xeon 5处理器和Nvidia的K80处理器，其集成板卡后的性能对比如下：

以MLP0为例子，7ms的反应时间内TPU可得到的数据处理吞吐率是CPU和GPU的数十倍（但值得指出的是，该benchmark下TPU的计算占用率要远远高于GPU和CPU）。

图9将K80 GPU和TPU相对于Haswell CPU的几何平均功耗性能进行了对比，包括两种计算量程——第一个是“总计”性能/瓦，包括在计算GPU和TPU的性能/瓦特时由主机CPU服务器消耗的功率。；第二个“增量”性能/瓦，仅包括预先从GPU和TPU中减去主机CPU服务器的电源余下的功耗。TPU服务器的总性能/瓦特比Haswell高出17到34倍，这使得TPU服务器的性能/ Katt服务器的功率是14到16倍。 TPU的相对增量性能/瓦特（Google致力于定制ASIC的核心价值）为41到83，这将TPU提升到GPU的性能/瓦特的25到29倍。