Google TPU的发展历程与思考（一）

df12138

已于 2022-09-25 23:28:54 修改

阅读量3.9k

点赞数 2

分类专栏： AI Accelerator

于 2021-12-19 15:28:51 首次发布

本文链接：https://blog.csdn.net/df12138/article/details/122018914

版权

AI Accelerator 专栏收录该内容

2 篇文章 0 订阅

订阅专栏

Google TPU的架构与思考—— TPU v1

2016年5月的开发者大会上，Google推出了自行研制的人工智能芯片Tensor Processing Unit, TPU。五年后的2021年5月19日，Google 又推出了 TPU v4，五年时间 Google 已经相继推出了若干代 TPU 和 TPU Edge，本文从古至今，力图呈现五年来Google在设计 TPU 时的原则与思路的变化，探究未来 AI 加速器的设计方向。

如今已经有越来越多的架构师认为，DSA（domain-specific Architecture）领域专用架构可以有效地提高硬件的能耗比。TPU v1 处理器自 2015 年起部署于数据中心，用于加速神经网络的推理阶段，其性能指标如下：TPU v1 有 65536 个 8-bit MAC(矩阵乘单元)，峰值性能为 92 TOPS，以及 28 MiB 的片上内存空间。
经过实验，Google 的科学家认为相比于 CPU 和 GPU，TPU 的运行架构可以更好地满足神经网络程序对响应时间的要求。相比于 CPU内部复杂的电路逻辑（缓存、乱序、多线程、预取等）用于提高程序指令的吞吐率，TPU 内部的结构更加简单暴力，因为其目标是降低延迟，而非获得高吞吐率。

设计目标

根据参考文献[2]，TPU v1貌似只支持深度神经网络的推理阶段运算，并且更加关注神经网络在 TPU v1 中的响应时间，因为这是面向用户的程序，Google 不希望它反应太慢。另一个开发TPU v1的理由是，随着深度神经网络的推广，在云端（数据中心）使用神经网络的需求高涨，而使用传统CPU的成本过于高昂。

TPU v1 的主要设计目标是，能在 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 v1主要针对2015年最火的深度神经网络进行优化，主要可以分为以下三类：

MLP 多层感知机
CNN 卷积神经网络
RNN LSTM 递归神经网络长短期记忆

benchmark

在这里插入图片描述

本论文亮点

神经网络推理程序通常强调响应时间而非吞吐量，因为他们经常是面向用户的
因为延迟的限制，K80 GPU 在推理过程中无法被完全利用，因此仅稍快于 Haswell CPU
尽管 TPU 更小且更节能，但相比 K80 GPU，TPU 有 25 倍多的 MACs 和 3.5 倍大的片上存储空间
TPU 比 K80 GPU 和 Haswell CPU 快大概 15-30 倍
在六个神经网络 benchmark 中，有四个是因为内存带宽而受限；如果TPU被修改为与K80 GPU相同的内存系统，它将比GPU和CPU快大约 30 - 50 倍
TPU的性能/瓦特是当代产品的 30 - 80 倍；K80内存的改版的 TPU 性能将提高 70 - 200倍。

架构

TPU v1 的架构框图如下：
TPU v1 Block Diagram
从左到右看，主机 CPU 通过 PCIe 总线与 TPU 进行连接。因此，TPU 不是与 CPU 紧密集成来减少部署延迟，而是被设计成 PCIe I/O 总线上的类似协处理器，允许它插入现有的服务器，就像 GPU 那样。
因时代限制，TPU 插口还是 SATA
但从这可以看出，Google 在当时就已经意识到，TPU 在云服务器系统中的地位。此外，TPU 的指令（下文会讲到）也是 CPU 发送给它的，目的是简化硬件设计和调试。

忽略掉红色的控制线路部分，可以看到指令和数据从主机接口（host interface）发出后，送往右处。而TPU v1 的主要计算部分是右边的黄色矩阵乘单元（Matrix Multiply Unit, MXU），这里会完成矩阵的乘加运算（Multiply-Add-Computation, MAC），其输入数据来自于蓝色的权重队列块（Weight FIFO）和蓝色的统一缓冲区（Unified Buffer, UB），计算完成后，沿着输出箭头，数据会发往蓝色的累加单元（Accumulator, Acc）和黄色的激活单元，激活单元收到数据后，使用非线性激活函数，计算完成后将结果送回 UB。

根据论文中的数据，

矩阵乘单元（MXU）有 256x256 个MACs，可用于完成 8-bit 的乘加运算，当然也可以完成 16-bit 的运算，只不过需要的时间会根据位长翻倍。每个 cycle 下，MXU 在每个时钟周期能产生256 个元素的 8-bit 部分和。除此之外，MXU 还有 64KiB*2（双缓冲隐藏延迟）的权重块，
Acc 模块是大小为 4MiB 的 256个元素的 32-bit 宽的累加器，同样是为了双缓冲机制而扩展的空间
Weight FIFO 的深度为4，MXU 使用的权值通过 Weight FIFO 暂存，Weight FIFO 从权值片外内存（8 GiB DRAM）中读取权重（推理阶段的权值是只读的）。
UB 用于保存中间计算结果，一共有 24MiB 大小，UB可通过 DMA 直接与Host CPU 交互

TPU 微架构的设计理念是尽量让计算模块保持繁忙 。反向理解一下，计算模块很可能不会一直繁忙，意味着数据无法及时喂给计算模块，因此 TPU 可能会遇到内存瓶颈。事实上论文后半部分也提到了内存瓶颈的点。

对于 MXU，TPU 采用了脉动阵列结构（见下图），不仅减少了对内存带宽的需求，还提高了能效。脉动阵列只适合用于非常规整简单的运算，但正巧矩阵乘和卷积就是这种规整又简单的运算，来自两个方向的数据以一定的间隔到达阵列中的 PE，并在那里进行运算。
脉动阵列示意图

TPU中，输入数据从左侧流入，权重从顶部流入，并且是预加载的，若给定的一256元的乘积运算操作，它会以对角波的形式通过阵列。由于权重是预先加载的，因此随着波的前进，可以立即算出乘积，通过控制部分和的通路，再往下累加。最终产生了256个输入立即被读取的错觉，并且它们立即更新256个累加器中的每个位置。从正确性的角度来看，软件不知道MXU的脉动实现，但从性能角度看，它确实需要考虑其中的延迟，因为数据一级级的传导，意味着延迟的步步累加。
在这里插入图片描述

ISA 相关

TPU v1的相关论文中，也介绍了关于 ISA 的话题，特别注意到，Google 发现 PCIe 速度较慢，因此将 ISA 设计为 CISC，从而让平均传输速度达到 10-20 CPI。ISA 中，有五条关键的指令：

Read_Host_Memory 从 CPU 内存中读取数据到 UB
Read_Weights 从 Weight 内存中读取权重到 Weight FIFO
MatrixMultiply/Convolve 让 MXU 计算矩阵乘或者卷积运算。一个矩阵操作中，大小为 B x 256 的变量，要乘以一个 256*256 的常量权重矩阵，最终得到一个 B x 256 的输出，共需要 B 个流水线时钟周期完成。
Activate 计算一个非线性激活函数，有 ReLU，Sigmoid 等等函数可选，也可以利用硬件上的资源完成池化操作。
Write_Host_Memory 把数据从 UB 传回 CPU 内存

其他指令是备用主机内存读/写、设置配置、两个版本的同步、中断主机、调试标记、nop和halt。CISC MatrixMultiply指令为12字节，其中3个是统一缓冲区地址；2为累加器地址；4是长度(有时是卷积的2维)；剩下的是操作码和标志。

这些 CISC 指令使用四级流水线，其中每条指令在单独的阶段执行，Google 计划用其他指令的执行与 MatrixMultiply 指令重叠，进而隐藏延迟。为此，Read_Weights 指令遵循去耦合访问/执行，即 Read_Weights 指令在发出要读取的权值的地址后就释放了，无需等到权值数据被取回后释放。通常，它在上一个计算繁忙时，从 Weight Memory 中获取权重至 Weight FIFO。而如果输入或权重数据没有准备好，矩阵单元将停止。
但 Google 的设计也有没考虑到的地方，他们发现，下一轮矩阵计算必须等到上一轮的激活函数计算完成后才能开始，因此会出现一个很明显的 delay slot，就是因为结果在被 UB 读取前会等待一个显式同步：

Interesting cases occur when the activations for one network layer must complete before the matrix multiplications of the next layer can begin; we see a “delay slot,” where the matrix unit waits for explicit synchronization before safely reading from the Unified Buffer

软件栈

所有 NPU 的软件栈设计与其硬件架构同样重要，因为只有 TPU 能与 CPU、GPU 兼容，才能更方便地被使用。在这方面，Google 的科学家开发了兼容性的 TPU 软件栈。TPU 的软件栈分为用户空间驱动和内核驱动。其中，内核驱动是稳定且轻量的，主要负责内存管理和中断。而用户空间驱动会频繁变化，主要管控 TPU 的启动，控制 TPU 执行流，生成 ISA 指令，转换数据维度等。当某个神经网络模型被第一次执行时，TPU 会对其进行编译，生成 TPU 的指令并缓存，将权值数据也缓存进入 Weight Memory 中。而后的几次运行时，TPU 就可以峰值速度运行神经网络。