芯片战争早已打响！谷歌15个月打造首个TPU

在计算的历史上，曾被丢弃或过时的想法仍然很有趣，有时甚至非常有用。

在这方面，我们这辈子能经历的最重要的例子莫过于神经网络了。

多数人可能都了解，在神经网络近70年的历史中，寒冬和泡沫交替出现，——事实上，藏在神经网络背后的专用硬件加速器（ASIC）也是如此。

神经网络和ASIC属于是「先有鸡还是先有蛋」的关系，大家都知道神经网络在CPU上效率不高，但是，如果没有证据证明神经网络是有用的，公司凭什么要为神经网络投资开发专门的硬件？

然而，正因为没有合适的硬件，神经网络也work不起来……

在历史的周期演进中，构建专用硬件的项目一次又一次地被放弃，最终，作为外来者的GPU打破了僵局。

打开网易新闻 查看更多图片

毫无疑问，我们现在正处于泡沫周期，在大模型的激励下，从初创公司到超大规模企业，都在构建建立自己的硬件加速器。

他们之中的大多数都可能会失败，但也肯定会有成功的设计成为未来计算领域的重要组成部分。

——而谷歌的张量处理单元（TPU）作为第一批大规模部署的加速器，肯定会成为幸存者之一。

当我们在谈论TPU时，我们在谈论什么

为什么谷歌的TPU既有趣又重要？因为这是谷歌，TPU被切实应用于谷歌庞大的服务（搜索、Android、Chrome、Gmail、地图、Play 商店、YouTube、照片），用户超过10亿。此外，谷歌还拥有第三大公共云。

在英伟达的GPU主导了这个AI时代硬件供应的当下，谷歌的TPU是一个真正经历了时间和规模考验的竞品。

以下的内容，小编分成两部分：第一部分讲故事，关于第一个谷歌TPU的前世今生；第二部分讲技术，揭秘TPU的架构细节和性能。

起源

机器学习对谷歌来说是一件大事。毕竟谷歌的既定使命是「组织世界信息，使其普遍可用和有用（to organize the world’s information and make it universally accessible and useful）」。

机器学习帮助谷歌发掘信息的价值，从图像和语音识别到语言翻译，再到大型语言模型，当然也少不了谷歌的「摇钱树」——年入数十亿美元的广告业务。

2010年代初，谷歌的注意力开始转向深度学习：

2011年：Jeff Dean、Greg Corrado和Andrew Ng发起了关于深度学习的研究项目——Google Brain。 2013年：继AlexNet图像识别项目取得成功后，谷歌以4400万美元的价格收购了由Geoffrey Hinton、Alex Krizhevsky和Ilya Sutskever组成的初创公司。 2014年：谷歌收购了由Demis Hassabis、Shane Legg和Mustafa Suleyman创立的DeepMind，价格高达65000万美元

2013年，当AlexNet的开发者Alex Krizhevsky来到谷歌时，他发现谷歌现有的模型都在CPU上运行。

于是，在公司工作的第一天，他出去从当地的一家电子商店买了一台GPU机器，把它插入网络，然后开始在GPU上训练神经网络。

最终，大家意识到需要自己需要更多的GPU，2014年，谷歌决定以约13000万美元的价格购买40,000个英伟达GPU。

深度学习所提供的能力可以大规模应用于谷歌的各种业务，然而，无论是在技术上还是在战略上，依赖英伟达GPU都不一定是最佳解决方案。

GPU不是ASIC，它不是为神经网络量身打造的，对于需要大规模应用的谷歌来说，相当于要付出很大的额外代价；同时，依赖单一供应商显然也是战略上的重大风险。

谷歌当时有两种选择：现场可编程门阵列（FPGA）和专用集成电路（ASIC）。

当时的实验证明，FPGA的性能打不过GPU，而定制ASIC可能在推理方面产生10倍的性价比优势。

快速交付

开发ASIC的时间成本是一个重要的考量，如果落地周期太长，那么硬件本身也就失去了意义。为此谷歌动用了一切可以快速访问的资源和现有知识。

谷歌迅速招募了一支才华横溢、经验丰富的团队，其中包括David Patterson，——伯克利RISC原始设计的开发者，也是RISC-V指令集架构开发的关键人物。

为了赶时间，谷歌没有去从头开发新的架构。幸运的是，在35年前，就有人为他们准备好了合适的架构。

1978年，卡内基梅隆大学的H.T.Kung和Charles E. Leiserson发表了《Systolic Arrays (for VLSI)》，提出了「systolic system」。

论文地址：https://www.eecs.harvard.edu/htk/static/files/1978-cmu-cs-report-kung-leiserson.pdf

A systolic system is a network of processors which rhythmically compute and pass data through the system….In a systolic computer system, the function of a processor is analogous to that of the heart. Every processor regularly pumps data in and out, each time performing some short computation so that a regular flow of data is kept up in the network. systolic system是一个处理器网络，它有节奏地计算并通过系统传递数据……处理器的功能类似于心脏，每个处理器都会定期将数据泵入和泵出，每次都执行一些简短的计算，以便在网络中保持常规的数据流。

打开网易新闻 查看更多图片

作者同时提出了架构的一种应用：矩阵计算。

Many basic matrix computations can be pipelined elegantly and efficiently on systolic networks having an array structure. As an example, hexagonally connected processors can optimally perform matrix computation……These systolic arrays enjoy simple and regular communication paths, and almost all processors used in the network are identical. As a result, special purpose hardware devices based on systolic arrays can be built inexpensively using the VLSI technology. 许多基本的矩阵计算可以在具有数组结构的脉动网络上优雅而有效地执行流水线。例如，六边形连接的处理器可以最佳地执行矩阵计算……这些脉动阵列享有简单而规则的通信路径，并且网络中使用的几乎所有处理器都是相同的。因此，使用VLSI技术可以廉价地构建基于脉动阵列的专用硬件设备。

——这不就来了吗！简单、便宜、还适合算矩阵！真是老天爷喂饭吃。

有了工程师和架构之后，谷歌还与当时的LSI公司（现在是Broadcom的一部分）合作，帮助芯片落地。

另外，全新的芯片意味着需要全新的指令集架构、全新的编译器以及上层软件，这是一个庞大的工程。

2015年初，第一批TPU正式部署在谷歌的数据中心，此时，距离项目启动仅仅过去了15个月，看一下这个庞大的研发团队：

论文地址：https://arxiv.org/ftp/arxiv/papers/1704/1704.04760.pdf

为了这个进度，谷歌也不得不做了很多取舍，包括使用较老的工艺（28nm），以及较低的时钟频率（700MHz）.

2016年3月，使用了TPU的AlphaGo Zero击败了当时的世界围棋冠军李世石，举世震惊。

TPU架构

在神经网络的运算中，矩阵乘法是关键，下面是一个简单的例子：

那么，TPU是如何计算矩阵乘法的呢？

在上面的故事中，我们提到了TPU的原理：systolic arrays（脉动阵列），假设有如下的2*2矩阵乘法：

打开网易新闻 查看更多图片

计算结果如下图所示。每个角上的方块代表一个乘法/累加单位 （MAC），可以执行乘法和加法运算。

在此图中，黄色值是从顶部和左侧输入矩阵的输入。浅蓝色值是存储的部分总和。深蓝色值是最终结果。

第一步，a11和b11加载到左上角的MAC中，相乘并存储结果。 第二步，a12和b21加载到左上角的MAC中，相乘并添加到先前计算的结果中。这一步得到了结果矩阵的左上角值。同时，b11被传输到右上角的MAC，乘以新加载的a21，并存储结果；a11被传输到左下角的MAC，乘以新加载的b12，并存储结果； 第三步，b21被传输到右上角的MAC，乘以新加载的值a22，结果被添加到以前存储的结果中；a12被传输到左下角的MAC，乘以新加载的b22，并将结果添加到先前存储的结果中。此时得到了结果矩阵的右上角和左下角值。同时，a12和b21被传输到右下角的MAC，相乘并存储结果。 第四步，将a22和b22传输到右下角的MAC，相乘并将结果添加到先前存储的值中，从而得到结果矩阵的右下角值。

至此，2*2矩阵乘法完成。

完整计算系统的最简单表示如下：

首先要注意的是，TPUv1依赖于通过PCIe（高速串行总线）接口与主机进行通信。它还可以直接访问自己的DDR3存储，

更详细的设计可以扩展成下图这样子：

DDR3 DRAM / Weight FIFO：权重存储在通过DDR3-2133接口连接到TPU v1的DDR3 RAM芯片中。权重通过PCIe从主机的内存预加载，然后可以传输到权重FIFO存储器中，供矩阵乘法单元使用。 Matrix Multiply Unit：256 x 256大小的矩阵乘法单元，顶部输入256个权重值，左侧是256个input值。 Accumulators：运算结果从脉动阵列的底部汇总到累加器中（内存）。 Activation：激活函数。 Unified Buffer / Systolic Data Setup：应用激活函数的结果存储在统一缓冲区存储器中，然后可以作为输入反馈到矩阵乘法单元，以计算下一层所需的值。

TPU v1执行8位整数乘法，利用量化来避免消耗更大的浮点计算。

TPU v1使用CISC（复杂指令集）设计，只有大约20条指令。这些指令是由主机通过PCIe接口发送给TPU的，而不是从内存中获取的。

矩阵运算采用可变大小的B*256输入，将其乘以256×256恒定权重输入，生成B*256输出，所以需要B次流水线循环才能完成。

TPU的整个执行过程大概像这样子：

Read_Host_Memory Read_Weights Loop_Start Matrix_Multiply Activate Loop_End Write_Host_Memory

由于TPU v1的指令集非常简单，可以保证芯片用于解码和相关活动的开销非常低，只有2%的芯片区域专门用于控制。

而24%的芯片面积专用于矩阵乘法单元，29%用于存储输入和中间结果的统一缓冲区存储器。

2013年，TPU v1与英特尔的Haswell CPU和英伟达的K80 GPU进行了比较：

TPU v1的MAC数量是K80 GPU的25倍，片上内存是K80 GPU的3.5倍。 TPU v1的推理速度比K80 GPU和Haswell CPU快15到30倍。 TPU v1的相对计算效率是GPU的25到29倍。

时至今日，这个「临时赶工」的ASIC，已经不断完善和壮大，在英伟达的统治之下，开辟了另一条道路。

参考资料：

https://thechipletter.substack.com/p/googles-first-tpu-architecture