摘要

TPUv2是Google在TPUv1的Systolic Array基础上使用VLIW指令架构进行的大规模架构更新，TPUv3则在TPUv2基础上的PPA优化。本文基于公开信息讨论Google TPU的架构演进，供大家参考，如有错误疏漏还请批评指正。

 

Table 1TPU v1-v4 chip架构指标[2]

Table 1 TPU v1-v4 chip架构指标[2]

注：向量单元数量为推导数值，TPUv4数值源自MLperf training 0.7[3]

 

TPUv1到TPUv2的设计思路和架构演化

 

TPUv1对于google来说是其进入芯片设计领域的第一步试水的作品，主打inference。而TPUv2的战略意义在于补齐了google在training方面的能力，在这之前training一直是GPU的天下，而其它架构例如CPU/DSP的架构和GPU相比各方面都无法抗衡，TPU在推出之后，google在training领域也能和GPU竞争了。

关于ML training的一些挑战，google认为主要有以下几个方面：

1 计算：训练需要反向计算，包括转置，batch-norm统计，这些是前向所没有的需求。

2 存储：训练需要存储前向的activation，用于计算反向的gradient。所以需要更大的存储空间。

3 精度：训练需要更大的动态范围，所以int8已经不够用了，必须上浮点。

4 编程性：新型神经网络的研发非常快，需要支持不同的调参优化器，需要有能力构造新型的神经网络。

5 并行性：对于训练而言，如何将多个device组成一个超大系统成为设计目标之一，多卡和多节点的支持变成硬性需求。

Google在文献2中介绍了一下的自己团队，看上去人可能并不是很多，并不像GPU/手机芯片这种动则几百人的前端团队，用中国的话说，google的team 非常精干。但是，我认为，也正是这种人员配置，使其只能选择一定的复杂度的架构和schedule，是google选择偏向VLIW架构的原因之一。

中间google花了详细的几页PPT介绍了一下自己的TPUv1到TPUv2的演化路线。

 

Figure 1TPUv1 架构图[1]

 

Figure 2TPUv2架构图[1]

 

我提取精要，总结如下：

• Dedicate的activation storage和accumulators设计被取消。
• 增加了vector memory。
• activation pipe换成一个vector unit，其中替代了acc和activation pipe处理的功能。
• matrix multiply unit不再直接访存，而是通过vector unit进行数据的读取和写回。
• 将DDR替换乘HBM，latency更小，带宽更大。
• 同时为了支持多卡多节点互联，设计了interconnect模块。

这里的改动基本上都是google针对training的特点而进行的架构升级，这里提一下第4点，这样的好处我认为有两点，第一，也就是op fuse，矩阵乘卷积操作后不用写回memory就能进行element wise操作。第二，DDR3之间和matrix multiply unit进行load store 操作变成了matrix multiply unit和vector memory的操作，这样效率上可能会有一定损失，但是访存灵活度提高。可以实现更灵活的数据复用策略。估计这里可能也参考了GPU的tensor core的一些访存设计。

在我看来，TPUv2实现了从fixed pipeline结构加速器到可编程处理器结构的思路转变，同样的存储单元和运算单元，变成了程序显式可编程单元，而不是简单的配置寄存器，硬件调度，这样可编程性和灵活度大大提高，而这些正是训练时所需要的。

 

TPUv2模块详细介绍

 

TPUv2的ISA 是VLIW形式, 一个指令bundle 322 bit，里面包括2个scalar slot，4个vector slot (2 for load/store) 2个matrix slot（push,pop）和1个misc slot，也就是一条指令里面包含9个指令slot, 指令包里面有6个立即数用于计算或者寻址。

对于VLIW而言，硬件设计相对简单，但是指令之间fetch和issue是紧耦合的，一旦有一个slot的指令有dependency，导致流水线hazard，那么整个指令包将被stall住，因此complier编译设计的时候十分小心，不然slot之间互相牵制，产生很多空nop指令slot，导致流水线气泡，performance惨不忍睹。

TPUv2整个chip结构里面包括core x 2、pcie x 2、interconnect router、link x 4、HBM x 2。

 

Figure 3TPUv2 chip结构[1]

 

其中Core里面主要分为4块：

·标量单元

·向量单元

·矩阵乘法单元

·转置和重排单元

标量单元相当简洁清楚，标量单元负责从instruction buffer中fetch整个VLIW指令bundle，然后issue到ALU或者mem，进行load/store/mov/alu/dma操作，这里标量配有2路ALU，32x32b register file，4kx32b的memory，这里还有sync flags模块用于进行一些外部的同步控制操作。另一方面标量部分还控制着往vector和matrix unit进行issue指令的过程。

 

Figure 4 TPUv2标量单元[1]

 

向量处理部分主要有vector memory，vector reg, vector alu, vector memory的大小为32k x 32b=128KB，register file是32个entry，32bit，一共8个lane，两路vector ALU，本人猜测这里有一个是激活函数，一个进行乘加mac，这里我们可以知道访存的宽度是32x128=256b，一共有128组这样的vector unit，所以一个core的memory是16MB, mac数是1024。而matrix unit的访存操作都通过vector reg进行，注意这里从矩阵输出的数据还有一些绿色的storage被标出用于缓冲数据。

Google提到，TPUv2的向量单元在设计之初由于batch-norm的加入导致vector算力成为瓶颈，因此在TPUv2中大幅强化了vector能力[2]。

 

Figure 5 Tpuv2向量单元[1]

 

矩阵单元脉动阵列是TPUv1留给TPUv2的核心财富之一。包括128x128的脉动阵列，bf16格式的乘法，fp32的累加单元，其中左矩阵和结果是流动的，右矩阵是stationary在矩阵中的。且只有右矩阵能进行转置操作，这里我们从反向算法角度，可以推测出，在计算正向和反向activation gradient的时候，右矩阵应该是输入weight，而在计算反向weight gradient的时候，input activation应该在右矩阵。Google详细说明了其选择128x128这种脉动阵列原因，红线表示每个操作数会进行多少次乘加操作，蓝线表示矩阵乘法的mac利用率。对于design来说，1个256x256单元和16个64x64单元和4个128x128单元所用的mac总的数量是相同的，都是65536个mac。

 

Figure 6Tpuv2矩阵单元尺寸设计[1]

 

首先是operations per operand ratio这条红线，对于256x256的脉动阵列，对应的逻辑操作是矩阵乘法m256xn256xk256, 每个operand in A会在算256个结果元素时用到，同理，64x64，每个operand in A会被64个结果元素用到，在数据复用性上来说，前者是后者的4倍，但是要找到一个应用有很大的数据复用性并不容易。另外一条蓝线是mac利用率，对于尺寸更大的矩阵乘来说，某些应用就是一些小的矩阵乘法，其算力利用率不高。而数据复用性和利用率对于mac array来说其实是两个相矛盾的指标，因此2者权衡之下，Google选择了4个128x128的matrix multiply单元。

另外，同样mac数量的矩阵，面积占比不一样，大矩阵乘法器运算单元占比大，而小的矩阵单元，IO和control面积会占比较大。

这里强调一下，TPUv1的矩阵单元只能进行int8操作，TPUv2/TPUv3的矩阵单元只能进行bf16操作，并使用bf16进行inference运算，这种单一的数据格式使得整个芯片功耗和面积上面TPU对NV tensor core都具有一定优势，但是在inference应用的TOPS/W会不如Nvidia T4。

另外，TPU中还提供了矩阵的transpose，reduction，permutation等操作，能够实现不同lane之间数据的互换重排。

Memory系统是芯片重要设计部分，对于访存操作, load/store对象都是SRAM scratchpads，同时Google提到了striding over vector，能从一块地址连续的数据中取出某些数据段，这些数据的首地址按stride跳跃的，数据的搬移通过异步DMA操作进行，软件需要参与管理。但是好处是系统的访存调度是可预测的，对应performance分析很容易得到准确的数值。Google 在load store操作中通过sync flags来进行stall或者indicate操作完成，防止out of order造成的程序语义错误。

 

Figure 7 Tpuv2 2D torus结构[1]

 

 

Figure 8 TPUv2 chip互联结构[1]

 

Google的TPU board 互联设计中包含一个router，和2个core相连，同时将2个HBM也连入router，router上面有4个link，每个link 400Gbps，这些link和其它TPU board相链接，对应2D torus结构中虚线框中的8个输入输出。2D torus的ring的特点在于每次reduce操作都在充分利用所有core的算力，过程是先是在水平方向做scatter-reduce，然后在垂直方向做all reduce，最后在水平方向做all-gather进行数据复制。从软件看来，DMA操作就像访问HBM，用不同chip id来区分不同的core。

 

Figure 9TPUv2 floorplan[2]

 

TPUv2的Floorplan的设计看起来中规中矩，值得注意的是，这里vector unit/memory分别到了两个physical partition。当中间的matrix 单元在读写数据时，其实它面对的是一个物理上分布式的存储结构，而在程序看来是一块完整的L1 memory，这里和含光的分布式L1设计有类似的地方。

 

TPUv3 架构演化

TPUv3 架构其实就是在TPUv2上面做的PPA优化，具体包括：

l矩阵乘法单元x2

lCore频率提升30%

l提升30% HBM带宽

l提升2倍的HBM存储容量

l提升4倍的节点能力

l提升30%的link带宽

 

Figure 10 TPUv3 架构更新[1]

 

实际应用部署落地

 

任何处理器，比起benchmark，实际应用产品的经验反馈更能体现现实情况，这些反馈对于设计的迭代也是非常非常重要的一环。Google天生就有大量的应用需求、数据和模型，在这一点上有很多一般的芯片公司所不具备的实战能力。

这里Google主要列出了6大应用，3种网络类型

这里Google主要列出了6大应用，3种网络类型

 

Figure 11 TPU实际应用部署[1]

 

总结

 

Google TPUv1可以说是Google试水芯片的作品，其脉动阵列在5年后看依然是经典设计，很多大学将其放进教科书，TPUv2完美的继承了脉动阵列的设计，而可编程性的VLIW指令集的加入让TPUv2重新define了整个架构，这些转变其实是翻天覆地的变化，从硬件设计、compiler、driver这些角度来说，相当于重新设计了一款芯片。而TPUv3在TPUv2的基础上PPA方便做了进一步提升，可以说TPUv3使其架构真正稳定下来。要知道google在现在依然采购在Nvidia T4等产品，说明Google现在还不能完全替代GPU的解决方案！“your beautiful DSA can fail if best-practice algorithms change”[2]，随着算法的快速演进，这让我们更加期待现在测试中的TPUv4中会引入哪些新feature。是引入SIMT？还是引入新的内存模型？新的任务分配调度模型？让我们拭目以待！

 

参考资料

 

[1] Thomas Norrie, Nishant Patil, Doe Hyun Yoon, George Kurian, Sheng Li, James Laudon, Cliff Young, Norman P. Jouppi, and David Patterson, "Google’s Training Chips Revealed: TPUv2 and TPUv3", HOTCHIPS 2020

[2] NORMAN P. JOUPPI, DOE HYUN YOON, GEORGE KURIAN, SHENG LI, NISHANT PATIL, JAMES LAUDON, CLIFF YOUNG, AND DAVID PATTERSON, "A Domain-Specific Supercomputer for Training Deep Neural Networks", communications of THE ACM JULY 2020 VOL.64 NO.7

[3] MLPerf Training v0.7 Results https://mlperf.org/training-results-0-7/