7.4 DSA设计_TPU

7.4TPU

Tensor处理单元（TPU）是谷歌首款用于WSC的自定义ASIC DSA,一种推理数据中心加速器。它的领域是DNN的推理阶段，并使用为DNN设计的TensorFlow框架进行编程。部署了第一个TPU2015年在谷歌数据中心。TPU的核心是65536（256x256）8位ALU矩阵乘法单元以及一个大型软件管理的片上存储器。TPU的单线程、确定性执行模型与第99百分位的响应时间非常匹配典型DNN推理应用程序的要求。

（1）TPU起源

早在2006年，谷歌的工程师就开始讨论部署其数据中心中的GPU、FPGA或自定义ASIC。他们得出的结论是几乎没有几个可以在特殊硬件上运行的应用程序可以免费运行利用大型数据中心的过剩容量，很难在免费的基础上进行改进。2013年，对话发生了变化，当时预测如果人们每天使用三分钟语音搜索,通过语音识别DNN，为了满足这样的计算需求，谷歌的数据中心计算资源需要翻一番。
对于传统的CPU来说，这将是非常昂贵的。然后谷歌启动了一个高优先级项目，以快速生产用于推理的定制ASIC（并购买了现成的GPU用于训练）。目标是将成本效益提高10个GPU以上。鉴于这一授权，TPU在短短15个月内就被设计、验证（Steinberg，2015）、构建并部署在数据中心。

（2）TPU架构

为了减少延迟部署的机会，TPU被设计为PCIe I/O总线上的协处理器，可将其插入现有的服务器。此外，为了简化硬件设计和调试，主机服务器通过PCIe总线直接向TPU发送指令以供其执行而不是让TPU获取指令。因此，TPU在精神上更接近FPU（浮点单元）协处理器，而不是从内存中获取指令的GPU。
下图显示了TPU的框图。主机CPU通过PCIe总线将TPU指令发送到指令缓冲器中。内部块通常通过256字节宽（2048位）的路径连接在一起。从右上角开始，矩阵乘法单元是TPU的核心。它包含256x256个ALU，可以对有符号或无符号执行8位乘法和加法整数。16位乘积被收集在矩阵单元下方的32位累加器的4 MiB中。当使用8位权重和16位激活的混合时（反之亦然），矩阵单元以一半的速度计算，当两者都是16位时，它以四分之一的速度计算。它每个时钟周期读取和写入256个值，并且可以执行矩阵乘法或卷积。非线性函数由激活硬件计算。矩阵单元的权重通过芯片上的Weight FIFO进行分级，该FIFO从称为Weight Memory的芯片外8GiB DRAM读取（为了推断，权重是只读的；8GiB支持许多同时活动的模型）。
中间结果保存在24 MiB片上统一缓冲区中，该缓冲区可以作为矩阵乘法单元的输入。可编程DMA控制器向CPU主机存储器和统一缓冲区传输数据或从CPU主机存储器传输数据。

（3）TPU指令集体系结构

由于指令是通过相对较慢的PCIe总线发送的，因此TPU指令遵循CISC传统，包括重复字段。TPU没有程序计数器，也没有分支指令；从主机CPU发送指令。这些CISC指令的每条指令的时钟周期（CPI）通常为10-20。它总共有十几条指令，但这五条是关键
其中：
1.Read_Host_Memory将数据从CPU主机存储器读取到统一缓冲区中。
2.Read_Weights将权重从权重存储器读取到权重FIFO中作为矩阵单元的输入。
3.MatrixMultiply/Convolve使矩阵乘法单元执行矩阵矩阵乘法、向量矩阵乘法、逐元素矩阵乘法、按元素向量乘法或从统一缓冲器到累加器的卷积。矩阵运算采用可变大小的B256输入，将其乘以256x256输入常量，并产生B256输出，需要B个流水线循环才能完成。例如，如果输入是256的4个矢量元素，B将是4，因此需要4个时钟周期才能完成。
4.Activate执行人工神经元的非线性功能，有ReLU、Sigmoid、tanh等选项。它的输入是Accumulators，输出是Unified Buffer。
5.Write_Host_Memory将数据从统一缓冲区写入CPU主机存储器。
其他指令是备用主机内存读/写、设置配置、两个版本的同步、中断主机、调试标记、nop和暂停。CISC MatrixMultiply指令为12个字节，其中3个是统一缓冲区地址；2是累加器地址；4是长度（有时对于卷积是2维）；其余的是操作码和标志。
目标是在TPU中运行整个推理模型，以减少与主机CPU的交互，并具有足够的灵活性，以满足2015年及以后的DNN需求，而不仅仅是2013年DNN所需的需求。

（4）TPU微体系架构

TPU的微体系结构原理是让矩阵乘法单元保持忙碌。计划是通过将其他指令的执行与MatrixMultiply指令重叠来隐藏这些指令的执行。因此，前面四类通用指令中的每一类都有单独的执行硬件（读写主机存储器合并到同一单元中）。为了进一步提高指令并行性，Read_Weights指令遵循去耦访问，因为它们可以在发送地址之后但在从权重存储器提取权重之前完成。矩阵单元有来自统一缓冲器和权重FIFO的未准备就绪信号，如果它们的数据还不可用，则会导致矩阵单元停滞。
请注意，TPU指令可以执行多个时钟周期，这与每级一个时钟周期的传统RISC流水线不同。
由于读取大型SRAM比算术要昂贵得多，矩阵乘法单元使用收缩执行来通过减少统一缓冲区的读取和写入来节省能量。收缩数组是算术单元的二维集合，每个算术单元独立地计算部分结果，作为来自被认为是每个单元上游的其他算术单元的输入的函数。它依赖于来自不同方向的数据，这些数据以规则的间隔到达阵列中的单元，并在那里进行组合。因为数据以前进的波前的形式流经阵列，所以它类似于心脏通过人体循环系统泵送血液，这就是收缩名称的起源。
图7.13展示了收缩阵列的工作原理。底部的六个圆圈是用权重wi初始化的乘法累加单元。交错输入数据xi显示为从上面进入数组。图中的10个步骤表示从页面顶部向下移动的10个时钟周期。收缩阵列向下传递输入，向右传递乘积和总和。当数据完成其通过收缩阵列的路径时，所需的乘积总和出现。请注意，在收缩数组中，输入数据只从内存中读取一次，输出数据只写入内存一次。

在TPU中，收缩阵列是旋转的。图7.14显示了从顶部加载权重，输入数据从左侧流入数组。给定的256元素乘累加运算作为对角波前在矩阵中移动。权重被预加载，并随着前进的波与新块的第一数据一起生效。控制和数据被流水线化，以产生一种错觉，即256个输入被一次读取，并且在馈送延迟之后，它们更新256个累加器存储器中的每个存储器的一个位置。从正确性的角度来看，软件不知道矩阵单元的收缩性质，但为了性能，它确实担心单元的延迟。

（5）TPU软件

TPU软件堆栈必须与为CPU和GPU开发的软件堆栈兼容，以便快速移植应用程序。在TPU上运行的应用程序部分通常使用TensorFlow编写，并编译为可以在GPU或TPU上执行的API（Larabel，2016）。图7.17显示了MLP的一部分的TensorFlow代码。
与GPU一样，TPU堆栈分为用户空间驱动程序和内核驱动程序。内核驱动程序是轻量级的，只处理内存管理和中断。它是为长期稳定性而设计的。用户空间驱动程序经常更改。它设置并控制TPU的执行，将数据重新格式化为TPU顺序，并将API调用转换为TPU指令并将其转换为应用程序二进制文件。用户空间驱动程序在第一次评估模型时编译模型，缓存程序图像并将权重图像写入TPU权重存储器；第二次和随后的评估以全速运行。TPU完全从输入到输出运行大多数模型，最大限度地提高了TPU计算时间与I/O时间的比率。计算通常一次只进行一层，重叠执行允许矩阵单元隐藏大多数非关键路径操作。

（6）TPU改进

TPU架构师研究了微体系结构的变体，看看他们是否可以改进TPU。
与FPU一样，TPU协处理器有一个相对容易评估的微体系结构，因此TPU架构师创建了一个性能模型，并随着内存带宽、矩阵单元大小、时钟速率和累加器数量的变化来估计性能。使用TPU硬件计数器进行的测量发现，建模的性能平均在硬件的8%以内。
图7.18显示了TPU的性能灵敏度，因为这些参数在0.25到4的范围内缩放。除了评估仅提高时钟速率（图7.18中的时钟）的影响外，图7.18还绘制了一个设计（clock+），该设计可以提高时钟速率并相应地缩放累加器的数量，以便编译器可以保持更多的内存引用。同样，图7.18绘制了矩阵单位展开图，如果累加器的数量随着一维（矩阵+）的平方增加而增加，因为矩阵在两个维度上都在增长，并且只增加了矩阵单位（矩阵）。

首先，增加内存带宽（内存）的影响最大：当内存带宽增加4时，性能平均提高3，因为这减少了等待权重内存的时间。其次，时钟速率在有或没有更多累加器的情况下平均没有什么好处。第三，当所有应用程序的矩阵单元从256x256扩展到512x512时，图7.18中的平均性能略有下降，无论它们是否有更多的累加器。这个问题类似于大页面的内部碎片，只是更糟，因为它是二维的。
考虑LSTM1中使用的600x600矩阵。对于256x256矩阵单元，平铺600x600需要九个步骤，总共需要18μs的时间。更大的512x512单元只需要四个步骤，但每个步骤需要四倍的时间，即32μs的时间。TPU的CISC指令很长，因此解码是无关紧要的，并且不会隐藏DRAM加载的开销。
考虑到性能模型的这些见解，TPU架构师接下来评估了一种替代和假设的TPU，如果他们有超过15个月的时间来设计，他们可能会用相同的工艺技术来设计。更积极的逻辑合成和块设计可能会将时钟速率提高50%。架构师发现，为K80使用的GDDR5存储器设计接口电路，可以将Weight memory带宽提高五倍以上。如图7.18所示，将时钟速率提高到1050 MHz，但对内存没有帮助，性能几乎没有变化。如果时钟保持在700 MHz，但它使用GDDR5代替Weight Memory，则性能将提高3.2，甚至考虑到在修改后的TPU上调用DNN的主机CPU开销。两者兼而有之并不能进一步提高平均性能。

（7）TPU设计总结

尽管I/O总线具有相对较小的内存带宽，这限制了TPU的充分利用。但TPU实现了在运行DNN推理应用程序时，其性价比比GPU提高十倍的目标。此外，重新设计的TPU，采用与GPU相同的内存技术时，计算速度将快三倍。解释TPU成功设计遵循了以下指导方针：
1.使用专用存储器来最小化数据移动的距离。
TPU具有24 MiB统一缓冲区，该缓冲区保存MLP和LSTM的中间矩阵和向量以及CNN的特征图。它针对一次256字节的访问进行了优化。它还有4个MiB累加器，每个32位宽，用于收集矩阵单元的输出，并作为计算非线性函数的硬件的输入。8位权重存储在单独的片外权重存储器DRAM中，并通过片内权重FIFO进行访问。
2.将放弃高级微体系结构优化所节省的资源投入到更多的运算单元或更大的内存中。
TPU提供28 MiB的专用内存和65536个8位ALU，这意味着它的内存和ALU数量是服务器级CPU的60%和250倍，尽管它的大小和功率是服务器级GPU的一半（见第7.9节）。与服务器级GPU相比，TPU的片上内存和ALUs数量是服务器类GPU的3.5倍和25倍。
3.使用与域匹配的最简单的并行形式。
TPU通过其256x256矩阵乘法单元的二维SIMD并行性提供其性能，该单元与收缩组织内部管道连接。GPU依赖于多处理、多线程和一维SIMD，而CPU则依赖于多处理器、无序执行和一维SIMD。
4.将数据大小和类型减少到域所需的最简单的大小和类型。
TPU主要基于8位整数进行计算，尽管它支持16位整数并以32位整数进行累加。CPU和GPU还支持64位整数以及32位和64位浮点。
5.使用特定于域的编程语言将代码移植到DSA。
TPU是使用TensorFlow编程框架编程的，而GPU依赖于CUDA和OpenCL，CPU必须运行几乎所有的东西。