自Google在2016年推出第一代自研人工智能芯片Tensor Processing Unit, TPU，经过数年发展已升级到了第4代 TPU v4 (截止到2022年底)。TPU架构设计也是通过高效地并行运算大量乘法和累加操作来实现深度学习卷积层，全连接层等网络层的高效计算。

脉动阵列

脉动阵列(systolic array)是由大量紧密耦合的PE组成阵列网络，是一种数据流的架构设计。脉动阵列中每一个PE都会与周围一个或者多个的PE进行数据通信，PE的功能相对简单，系统通过大量PE的并行计算来实现高效的计算，通过数据在不同PE间的流动最大化的实现数据的复用和减少芯片在运算过程中的内存访问的次数，不仅可以节省内存带宽还降低了因内存访问带来的功耗损失。

在传统计算系统的运算过程中，处理单元PE先从主存中读取数据，然后再进行运算操作，运算完成后结果再写回到存储器中，因此访存的速度称为了整个系统的处理性能的瓶颈。相比于CPU通过多级缓存的方式来解决访存速度瓶颈的问题，脉动架构则是让数据在不同处理单元中流动来减少访问主存带来的性能瓶颈。如图所示，在一个一维脉动阵列的运算过程中，数据首先从主存中进入第一个PE，经过处理后被传递到下一个PE，同时下一个数据进入第一个PE，以此类推，数据以流动的方式在不同PE之间传递，直到完成全部计算后才返回主内存。所以脉动阵列架构实现了输入数据的多次复用，省去了数据被写回主存后再读取的过程，减少主存的访问次数，因此脉动阵列可以在较小的内存带宽的情况下实现高吞吐率。

脉动阵列

脉动阵列架构有如下几个特点：

1. PE功能结构简单，实现成本低，可以集成大量PE来提高并行计算能力。

2. 由大量同构的PE组成的一维、二维或树的阵列结构，可灵活拓展。

3. 不同PE之间采用流水线的方式进行数据通讯，实现高效的数据复用。

4. 数据只能在相邻PE间流动，只适合特定的算法，比如矩阵运算，卷积等。

TPU架构设计

TPU使用了脉动阵列架构设计，数据在不同方向上以固定的时间间隔送入阵列中的 PE进行运算，在经过多步计算后最后将结果汇总输出来。脉动阵列只适合用于非常简单且规整的运算，而矩阵乘和卷积刚好符合脉动阵列的运算特性。

第一代TPU的脉动阵列结构实现如图所示，与GPU类似，TPU通过PCI-E总线与主机CPU连接，TPU的指令由CPU完成编码以简化硬件设计和调试。矩阵乘法单元(Matrix Multiply Unit, MXU)是主要的计算单元，其主要功能为完成矩阵的乘法运算。矩阵乘法单元周围有三个功能不同的数据缓存区以及专用的Activitation, Normalize和Pool单元，其中三个数据缓存区分别用于缓存权重队列(Weight FIFO)和输入数据缓冲区(Unified Buffer, UB)和输出数据累加单元(Accumulator, Acc)。

TPU架构

具体执行时，指令和数据从主机接口(host interface)进入TPU，复用率更高的权重参数会被预先加载到Weight FIFO，输入数据会被加载到统一缓冲区UB。输入数据和权重参数在矩阵乘法单元里完成矩阵乘法运算后结果再送入累加单元Acc，在Acc完成部分和累加后根据模型设计需要，选择是否进行送入Activitation, Normalize和Pool单元执行对应操作，最后再将结果送回统一缓冲区UB。

在第一代TPU硬件的布局中，矩阵乘单元和激活单元一共占30%的面积，矩阵乘单元(MXU)有256x256个MACs，每个时钟周期能完成256x256个8bit数据的乘加运算，Acc是大小为4MiB的32bit累加器。UB大小为24MiB，占29%的面积，其可以通过DMA直接与Host CPU交互，用于缓存输入数据或保存中间计算结果。Weight FIFO的深度为4，权值参数从片外内存中读取。

TPU布局

矩阵乘法单元(Matrix Multiply Unit, MXU)是一个典型的脉动阵列，如图所示，weight由上向下流动，数据从左向右流动。输入的像素数据从乘法矩阵左侧进入，并从左向右传播以重复利用中间数据，由于权重是预先加载的，因此随着输入数据的前进，可以快速算出乘积结果，在通过控制通路进入下一步累加计算。

矩阵乘单元数据流

TPU架构演进

TPU v1是谷歌第一代AI专用芯片，主要专注与处理推理任务，Google在推出用于推理的TPU v1之后开始开发和设计针对训练的第二代TPU。与TPU v1相比，TPU v2的有如下的改进点：

1. 1.TPU v2每个芯片有两个Tensor Core。

2. 2.将固定的激活函数(Activation Pipeline)改为可编程性更高的向量单元(Vector Unit)。

3. 3.使用一个向量存储器(Vector Memory)代替Accumulator和Activation Storage中的缓存。

4. 4.矩阵乘法单元作为向量单元的协处理器直接与向量单元连接，增加其可编程性。

5. 5.使用HBM代替DDR3，并改为与向量存储区相连，可以提供更高的带宽和读写速度。

6. 6.在HBM和向量存储区之间增加互连模块(Interconnect)，提供更强的扩展能力用于TPU之间的连接。

7. 7.增加了Scalar Unit, Transpose/Permute Unit等单元， 对Transpose等操作做硬件加速。
   

TPU v2架构

TPUv3在TPUv2的基础上进一步的提升了性能，其在时钟频率、内存带宽和芯片间带宽上提升了30%，矩阵乘法单元MXU的数量增加了1倍，HBM容量提升了1倍，可连接节点数增加了4倍。

TPU v1 v2 v3架构差异

TPU v4出于成本的考虑，Google对用于训练的和用于推动TPU进行了分开设计，用于训练的TPU v4有两个Tensor Core，而用于推理的PU v4i只有一个Tensor Core，在通用性、性能和成本上进行了取舍和平衡。在TPU v4i中，单个Tensor Core包含4个矩阵乘法单元MXU，是TPU v3的两倍。Google在TPU v4i的设计中还增加了性能计数器(Performance counter)以便于协助编译器更好地掌握芯片的运行情况。

郑重声明：文中仅代表作者观点，不代表本公司观点，本文版权归原作者所有，转载文章仅为传播更多信息之目的，如作者信息标记有误，请第一时间联系我们修改或者删除，感谢您的关注！