Google TPU 揭密

Google TPU（Tensor Processing Unit）问世之后，大家一直在猜测它的架构和性能。Google的论文“In-Datacenter Performance Analysis of a Tensor Processing Unit”让我们有机会一探究竟。

首先让我们看看摘要：

Many architects believe that major improvements in cost-energy-performance must now come from domain-specific hardware. This paper evaluates a custom ASIC—called a Tensor Processing Unit (TPU)—deployed in datacenters since 2015 that accelerates the inference phase of neural networks (NN). The heart of the TPU is a 65,536 8-bit MAC matrix multiply unit that offers a peak throughput of 92 TeraOps/second (TOPS) and a large (28 MiB) software-managed on-chip memory. The TPU’s deterministic execution model is a better match to the 99th-percentile response-time requirement of our NN applications than are the time-varying optimizations of CPUs and GPUs (caches, out-of-order execution, multithreading, multiprocessing, prefetching, …) that help average throughput more than guaranteed latency. The lack of such features helps explain why, despite having myriad MACs and a big memory, the TPU is relatively small and low power. We compare the TPU to a server-class Intel Haswell CPU and an Nvidia K80 GPU, which are contemporaries deployed in the same datacenters. Our workload, written in the high-level TensorFlow framework, uses production NN applications (MLPs, CNNs, and LSTMs) that represent 95% of our datacenters’ NN inference demand. Despite low utilization for some applications, the TPU is on average about 15X - 30X faster than its contemporary GPU or CPU, with TOPS/Watt about 30X - 80X higher. Moreover, using the GPU’s GDDR5 memory in the TPU would triple achieved TOPS and raise TOPS/Watt to nearly 70X the GPU and 200X the CPU.

这段摘要的信息量非常大，我把重点高亮出来了。首先，这个TPU芯片是面向datacenter inference应用。它的核心是由65,536个8-bit MAC组成的矩阵乘法单元（matrix multiply unit），峰值可以达到92 TeraOps/second (TOPS) 。有一个很大的片上存储器，一共28 MiB。它可以支持MLP，CNN和LSTM这些常见的NN网络，并且支持TensorFLow框架。摘要里面还可以看出，传统CPU和GPU使用的技术（caches, out-of-order execution, multithreading, multiprocessing, prefetching）它都没用，原因是它面向的应用都是deterministic execution model，这也是它可以实现高效的原因。它的平均性能（TOPS）可以达到CPU和GPU的15到30倍，能耗效率（TOPS/W）能到30到80倍。如果使用GPU的DDR5 memory，这两个数值可以达到大约GPU的70倍和CPU的200倍。

到此我们可以看出，Google的TPU采用了一个专用处理器或者硬件加速器的架构，比之前的想象的在GPU架构上改进的方法要激进的多，当然这样做实现的效率也高得多。

整体架构

上图是TPU的架构框图，按文章的说法，这个芯片的目标是：

“The goal was to run whole inference models in the TPU to reduce interactions with the host CPU and to be flexible enough to match the NN needs of 2015 and beyond, instead of just what was required for 2013 NNs.”

具体的架构信息如下：

“TPU指令通过PCIe Gen3 x16总线从主机发送到指令缓冲区（instruction buffer）。内部模块一般是通过256 bit宽的路径连接在一起的。从右上角开始，矩阵乘法单元（Matrix Multiply Unit）是TPU的核心。它包含256x256 MAC，可以对有符号或无符号整数执行8位乘法和加法。 16位运算结果由下面的4MiB 32bit的累加器处理。 4MiB代表4096个, 256-element, 32-bit累加器。矩阵单元每个时钟周期产生一个256-element的部分和（partial sum）。我们选择了4096，首先是注意到每个字节的操作需要达到峰值性能为大约1350，所以我们将其舍入到2048；然后复制了它，以便编译器可以在峰值运行时使用双缓冲方式。”

“当使用8 bit权重（weight）和16 bit激活（activation）（或反之亦然）的混合时，矩阵乘法单元以半速计算，并且当它们都是16bit时以四分之一速度计算。它每个时钟周期读取和写入256个值，可以执行矩阵乘法或卷积。矩阵单元保存一个64KiB的权重块（ tile of weights），并使用了双缓冲（这样可以在处理时同时读入权重）。该单元设计用于密集矩阵，而没有考虑稀疏的架构支持（部署时间的原因）。稀疏性将在未来的设计中占有优先地位。”

“矩阵乘法单元的权重通过片上的权重FIFO（Weight FIFO）进行缓存，该FIFO从片外8 GiB DRAM读取。由于是用于推论，权重为只读。8 GiB可以同时支持许多模型。权重FIFO的深度是4个tile。中间结果保存在24 MiB片上统一缓冲区（Unified Buffer）中，可以作为矩阵乘法单元的输入。可编程DMA实现和CPU Host内存以及统一缓冲区传输数据。”

下图给出了TPU芯片的布局，可以粗略看出各部分面积的比例。

指令集

按作者的说法，TPU的指令集是CISC类型（TPU instructions follow the CISC tradition），平均每个指令的时钟周期CPI（clock cycles per instruction）是10到20。指令总共应该有10几个，重要的指令如下：

1. Read_Host_Memory reads data from the CPU host memory into the Unified Buffer (UB).

2. Read_Weights reads weights from Weight Memory into the Weight FIFO as input to the Matrix Unit.

3. MatrixMultiply/Convolve causes the Matrix Unit to perform a matrix multiply or a convolution from the Unified Buffer into the Accumulators. A matrix operation takes a variable-sized B*256 input, multiplies it by a 256x256 constant weight input, and produces a B*256 output, taking B pipelined cycles to complete.

4. Activate performs the nonlinear function of the artificial neuron, with options for ReLU, Sigmoid, and so on. Its inputs are the Accumulators, and its output is the Unified Buffer. It can also perform the pooling operations needed for convolutions using the dedicated hardware on the die, as it is connected to nonlinear function logic.

5. Write_Host_Memory writes data from the Unified Buffer into the CPU host memory.

其它指令还包括：

The other instructions are alternate host memory read/write, set configuration, two versions of synchronization, interrupt host, debug-tag, nop, and halt. The CISC MatrixMultiply instruction is 12 bytes, of which 3 are Unified Buffer address; 2 are accumulator address; 4 are length (sometimes 2 dimensions for convolutions); and the rest are opcode and flags.

这是一个非常专用的指令集，主要就是运算和访存指令。从另一个角度也说明了TPU是一个非常专用的处理器。

微结构

首先是最重要的一句话“The philosophy of the TPU microarchitecture is to keep the matrix unit busy”，相信这也是所有设计NN加速器的同学的philosophy。

“TPU为CISC指令使用4级流水线，其中每条指令在单独的阶段中执行。设计是是通过将其它指令的执行与MatrixMultiply指令重叠，来隐藏其他指令的执行时间。为此，Read_Weights指令遵循解耦访问/执行（decoupled-access/execute）原理。它可以在发送地址之后，从权重存储器中取出权重之前就完成。如果输入激活或权重数据尚未就绪，矩阵单元就会暂停（stall）。”

“因为TPU的CISC指令可以占用站数千个时钟周期，这与传统的RISC流水线不同（每个阶段一个时钟周期）,所以TPU没有一个很清晰的流水线overlap diagrams。如果一个NN网络的层的激活操作必须在下一层的矩阵乘法开始之前完成，就会发生有趣的情况；我们看到一个“延迟槽”（delay slot），其中矩阵单元会在从Unified Buffer中安全读取数据之前等待显式同步信号（ explicit synchronization）。”

由于读取大的SRAM的能耗比算术运算高的多，所以矩阵单元通过减少统一缓冲区的读写来降低能耗，所谓脉动运行（systolic execution）。下图显示了数据从左侧流入，权重从顶部加载。给定的256-element乘累加运算通过矩阵作为对角波前（diagonal wavefront）移动。权重被预加载，并且与新块的第一数据一起在前进波上生效。控制和数据是流水线方式，所以好像给出256个输入是一次读取的，并且它们可以立即更新256个累加器中的每一个的一个位置。从正确性的角度来看，矩阵单元的脉动性对软件的是透明的。但是对于性能，它的延迟是软件需要考虑的因素。

软件

“TPU软件栈必须与为CPU和GPU开发的软件栈兼容，以便应用程序可以快速移植到TPU。在TPU上运行的应用程序的一部分通常写在TensorFlow中，并被编译成可以在GPU或TPU上运行的API。像GPU一样，TPU栈分为用户空间驱动程序和内核驱动程序。内核驱动程序是轻量级的，只处理内存管理和中断。它是为长时间的稳定性而设计的。用户空间驱动程序会经常变化。它设置和控制TPU执行，将数据重新格式化为TPU命令，将API调用转换为TPU指令，并将其转换为应用程序二进制文件。用户空间驱动程序在首次evaluate的时候编译模型，缓存程序image并将权重image写入TPU的权重存储器；第二次和以后的evaluate就是全速运行。 TPU从输入到输出完全运行大多数的模型，最大化TPU计算时间与I / O时间的比率。计算通常每次进行一层，重叠执行，允许矩阵乘法单元隐藏大多数非关键路径操作。”

文章后面很大篇幅是介绍TPU和CPU，GPU性能对比的，大家可以仔细看看，有机会再讨论吧。说实话，Google在2015年就能部署这样的ASIC张量处理器，不得不佩服啊。

T.S.

参考：

1. Norman P. Jouppi, et al."In-Datacenter Performance Analysis of a Tensor Processing Unit"，accepted by ISCA 2017