在 Google 发表 TPU 一年后，这款机器学习定制芯片的神秘面纱终于被揭开了。
5 日，Google 资深硬件工程师 Norman Jouppi 刊文表示，Google 的专用机器学习芯片 TPU 处理速度要比 GPU 和 CPU 快 15-30 倍（和 TPU 对比的是英特尔 Haswell CPU 以及 Nvidia Tesla K80 GPU），而在能效上，TPU 更是提升了 30 到 80 倍。
从这次发表的测试结果来看，TPU 似乎已经超出业界的预期，但是藏在这芯片背后的内部架构究竟有什么秘密？我们从 Jouppi 之前发表的论文当中，可以找到答案。
据了解，早在 4 年前，Google 内部就开始使用消耗大量计算资源的深度学习模型，这对 CPU、GPU 组合而言是一个巨大的挑战，Google 深知如果利用现有硬件，他们将不得不将资料中心数量翻倍来支援这些复杂的运算任务。
所以 Google 开始研发一种新的架构，Jouppi 称之为“下一个平台”。Jouppi 曾是 MIPS 处理器的首席架构师之一，他开创了内存系统中的新技术。3 年前他加入 Google 的时候，公司上下正在用 CPU、GPU 混合架构上来进行深度学习的训练。
Jouppi 表示，Google 的硬件工程团队在转向定制 ASIC 之前，早期还曾用 FPGA 来解决廉价、高效和高性能推理的问题。但他指出，FPGA 的性能和每瓦性能相比 ASIC 都有很大的差距。他解释，“TPU 可以像 CPU 或 GPU 一样可编译，它可以在不同的网络（卷积神经网络、LSTM 模型和大规模完全连接的模型）上执行 CISC 指令，而不是为某个专用的神经网络模型设计的。”一言以蔽之，TPU 兼具了 CPU 和 ASIC 的优点，它不仅是可编译的，而且比 CPU、GPU 和 FPGA 拥有更高的效率和更低的能耗。
TPU 的内部架构
▲ 该图显示了 TPU 上的内部结构，除了外挂的 DDR3 内存，左侧是主机界面。指令从主机发送到队列中（没有循环）。这些触发控制逻辑可以根据指令多次运行相同的指令。
TPU 并非一款复杂的硬件，它看起来像是雷达应用的讯号处理引擎，而不是标准的 X86 衍生架构。Jouppi 说，尽管它有众多的矩阵乘法单元，但是 GPU 更精于浮点单元的协处理。另外，需要注意的是，TPU 没有任何储存的程序，它可以直接从主机发送指令。
TPU 上的 DRAM 做为一个单元并行运行，因为需要获取更多的权重以馈送到矩阵乘法单元（算下来，吞吐量达到了 64,000）。Jouppi 并没有提到是他们是如何缩放（systolic）数据流的，但他表示，使用主机体件加速器都将成为瓶颈。
▲ 256×256 阵列缩放数据流引擎，经过矩阵乘法积累后实现非线性输出。
从第二张图片可以看出，TPU 有两个内存单元，以及一个用于模型中参数的外部 DDR3 DRAM。参数进来后，可从顶部加载到矩阵乘法单元中。同时，可以从左边加载触发（或从“神经元”输出）。那些以收缩的方式进入矩阵单元以产生矩阵乘法，它可以在每个周期中进行 64,000 次累加。
毋庸置疑，Google 可能使用了一些新的技巧和技术来加快 TPU 的性能和效率。例如，使用高带宽内存或混合 3D 内存。然而，Google 的问题在于保持分布式硬件的一致性。
TPU 对比 Haswell 处理器
在和英特尔“Haswell”Xeon E5 v3 处理器来的对比中，我们可以看到，TPU 各方面的表现都要强于前者。
在 Google 的测试中，使用 64 位浮点数学运算器的 18 核心运行在 2.3 GHz 的 Haswell Xeon E5-2699 v3 处理器能够处理每秒 1.3 TOPS 的运算，并提供 51GB／秒的内存带宽；Haswell 芯片功耗为 145 瓦，其系统（拥有 256GB 内存）满载时消耗 455 瓦特。
相比之下，TPU 使用 8 位元整数数学运算器，拥有 256GB 的主机内存以及 32GB 的内存，能够达成 34G B／秒的内存带宽，处理速度高达 92 TOPS ，这比 Haswell 提升了 71 倍，此外，TPU 服务器的热功率只有 384 瓦。
除此之外，Google 还测试了 CPU、GPU 和 TPU 处理不同批量大小的每秒推断的吞吐量。
如上图所示，在小批量任务中（16），Haswell CPU 的响应时间接近 7 毫秒，其每秒提供 5,482 次推断（IPS），其可以实现的最大批量任务（64）每秒则可以完成 13,194 次推断，但其响应时间为 21.3 毫秒。相比之下，TPU 可以做到批量大小为 200，而响应时间低于 7 毫秒，并提供 225,000 个 IPS 运行推理基准，是其峰值性能的 80%，当批量大小为 250，响应时间为 10 毫秒。
不过需要注意的是，Google 所测试的 Haswell Xeon 处理器似乎也不能完全说明问题，英特尔 Broadwell Xeon E5 v4 处理器和最新的“Skylake”Xeon E5，每核心时钟（IPC）的指令比这款处理器提升了约 5%。在 Skylake 是 28 核，而 Haswell 为 18 核，所以 Xeon 的总体吞吐量可能会上升 80%。当然，这样的提升与 TPU 相比仍有差距。
最后需要强调的是，TPU 是一个推理芯片，它并非是要取代 GPU，可以确定的是，TPU 与 CPU 一起使用对训练分析更加有益。但对于 CPU 制造商而言，如何研发出像 ASIC 一样兼顾性能和能效的芯片是现在以及未来要做的。
可以确定的是，Google 已经逐渐在自己的数据中心部署 TPU，但是上述测试数据只是理论结果，实际应用表现如何？我们还不得而知。
- First In-Depth Look at Google’s TPU Architecture
（本文由 雷锋网 授权转载；首图来源：Google）
延伸阅读：
- Google 研发了一块芯片，省下建资料中心的钱还推动机器学习的发展
