深入理解Google TPU的脉动阵列架构

Google TPU让“脉动阵列”(systolic array)这项“古老”的技术又回到大家的视野当中。短短几天,各种争论不绝于耳。其中有一个评论我比较喜欢,“这次google tpu让像我这样的年轻后辈又有机会能重新认识systolic,也很有意思”。是啊,还是让我们静下心来,重新认识一下“脉动阵列”这位老朋友吧。

Why systolic architectures?

要真正理解脉动阵列,首先要问的就是发明者的初衷。这正好也是1982年H. T. Kung论文的题目[1]。对于为什么要设计这样的架构,作者给出了三个理由:

1. Simple and regular design:可以说,简单和规则是脉动阵列的一个重要原则。而这样的设计主要是从“成本”的角度来考虑问题的。“Cost-effectiveness has always been a chief concern in designing special-purpose systems; their cost must be low enough to justify their limited applicability.”换言之,由于一个专用系统往往是功能有限的,因此它的成本必须足够低才能克服这一劣势。这也是设计专用处理器的一个基本考虑。更进一步,作者把成本分成两部分:nonrecurring cost(设计 design) 和recurring cost(器件 parts)。器件成本对于专用设计和通用设计基本是一样的,因此决定性的是设计成本。而设计一个合理的架构(appropriate architectures)是降低设计成本的一个重要方法。Google的TPU可以说很好的证明了这一点。通过采用脉动阵列这个简单而规则的硬件架构,Google在很短的时间内完成了芯片的设计和实现。从另一个角度来说,硬件设计相对简单,尽量发挥软件的能力,也是非常适合Google的一种策略。

2. Concurrency and communication:这一点主要强调并行性和通信的重要。由于这方面大家应该比较熟悉,这里不再赘述。

3. Balancing computation with I/O:平衡运算和I/O,应该说是脉动阵列最重要的设计目标。对于这一点,作者用下面这幅图给出了很好的说明。

首先,图中上半部分是传统的计算系统的模型。一个处理单元(PE)从存储器(memory)读取数据,进行处理,然后再写回到存储器。这个系统的最大问题是:数据存取的速度往往大大低于数据处理的速度。因此,整个系统的处理能力(MOPS,每秒完成的操作)很大程度受限于访存的能力。这个问题也是多年来计算机体系结构研究的重要课题之一,可以说是推动处理器和存储器设计的一大动力。而脉动架构用了一个很简单的方法:让数据尽量在处理单元中多流动一会儿。

正如上图的下半部分所描述的,第一个数据首先进入第一个PE,经过处理以后被传递到下一个PE,同时第二个数据进入第一个PE。以此类推,当第一个数据到达最后一个PE,它已经被处理了多次。所以,脉动架构实际上是多次重用了输入数据。因此,它可以在消耗较小的memory带宽的情况下实现较高的运算吞吐率。当然,脉动架构还有其它一些好处,比如模块化的设计容易扩展,简单和规则的数据和控制流程,使用简单并且均匀的单元(cell),避免了全局广播和扇入(fan-in),以及快速的响应时间(可能?)等等。

总结起来,脉动架构有几个特征:1. 由多个同构的PE构成,可以是一维或二维,串行、阵列或树的结构(现在我们看到的更多的是阵列形式);2. PE功能相对简单,系统通过实现大量PE并行来提高运算的效率;3. PE只能向相邻的PE发送数据(在一些二维结构中,也可能有对角线方向的数据通道)。数据采用流水线的方式向“下游”流动,直到流出最后的PE。

到这里不难看出,脉动架构是一种很特殊的设计,结构简单,实现成本低。但它灵活性较差,只适合特定运算。而作者认为,卷积运算是展示脉动架构特点的理想应用,因此在文章中用了很大篇幅介绍了脉动架构实现卷积运算的方法。首先看看作者对卷积运算的定义:

下面是作者给出的一种方法(broadcast inputs, move results, weights stay):

在这个例子中,X值广播到各个运算单元,W值预先存储在PE中并保持不动,而部分结果Y采用脉动的方式在PE阵列间向右传递(初始值为零)。不难看出,经过三个时刻,最右边的PE的输出就是X和W两个序列的卷积运算的第一个结果,这之后就会不断输出Y值。文中还给出了其它几种实现卷积的方式,如broadcast inputs, move weights, results stay;fan-in results, move inputs, weights stay。这几种都算作“(Semi-) systolic convolution arrays with global data communication”。另外还有一大类是“(Pure-) systolic convolution arrays without global data communication”,包括results stay, inputs and weights move in opposite directions;results stay, inputs and weights move in the same direction but at different speeds;weights stay, inputs and results move in opposite direction;weights stay, inputs and results move in the same direction but at different speeds。建议大家好好看看,还是挺有启发的。

那么二维的脉动阵列如何做矩阵相乘呢?下面这个例子比较清楚。A和B都是3x3矩阵,T表示时刻。

依此类推,中间省略几步,可以得到最后结果。

这里值得注意的是,要实现正确的矩阵运算,数据进入脉动阵列需要调整好形式,并且按照一定的顺序。这就需要对原始的矩阵进行一些reformat,这也增加了额外的操作。在后面对Google TPU的讨论中,我们可以更清楚的看到这一问题。

除了卷积运算,作者在文章中还说明的脉动架构适合的其它运算,主要包括信号和图像处理(signal and image processing),矩阵算术(matrix arithmetic)和一些非数值型应用(non-numeric application)。不过,由于脉动架构的灵活性较差,加之当时半导体技术的限制,它在被发明之后并没有得到广泛的应用。

时至今日,在我们设计Deep Learning处理器的时候,面对的一个主要矛盾仍然是I/O和处理不平衡的问题。这种不平衡既体现在速度上,也体现在功耗上。而一种减少I/O操作的思路 “…discuss how dataflows can increase data reuse from low cost memories in the memory hierarchy”[2],和脉动架构的方法是类似的。而另一方面,深度神经网络中大量使用了卷积运算和矩阵运算,这正好也是是脉动架构的优势。所以,在很多DL处理器中我们都可以看到脉动架构的影子(虽然有或多或少的变化和改进)。而Google TPU的设计更是让脉动阵列成了大家关注的焦点。

Google TPU Implementation

下面我们看看Google在TPU设计中是怎么实现脉动阵列的。作为一个强迫症患者,我把Google TPU的框图重新画了一下,又结合他们的专利细化了一下Matrix Multiply Unit的结构以及其中的cell的结构,算是小小的福利吧。 当然,很多细节也是我的猜想,可能不准确,如果大家发现什么问题也请指教。

首先,从整体来看,整个TPU的核心就是Matrix Multiply Unit,它正是一个256X256的脉动阵列。而整个芯片的其它部分,都是围绕这个脉动阵列来运转的,目的就是能够让脉动阵列尽量高效的运作(不过从TPU论文来看,实际效率并不算很高)。

这里有两个值得注意的地方。第一,是确实实现了相对较低的I/O带宽30GiB/s;第二是64K个MAC(可能是目前公开文献中看到的最大的MAC数量)。这两点很大程度上归功于脉动阵列的特点:平衡I/O和运算,以及相对简单的PE(cell)。当然,还有另外一个原因,TPU采用8bit定点的MAC(可以支持16bit运算)。这也是Google做出的最重要的设计选择之一。还是那句话,专用处理器的设计基于对目标应用的理解。在这一点上,估计没有人比Google更了解data center的DNN inference需求了。也正因如此,Google的TPU才会引起这么多关注。


根据Google的专利,matrix单元就是一个典型的脉动阵列。weight由上向下流动,activation数据从左向右流动。在最下方有一些累加单元,主要用于weight矩阵或者activation矩阵超出matrix单元范围的时候保存部分结果。控制单元实际上就是把指令翻译成控制信号,控制weight和activation如何传入脉动阵列以及如何在脉动阵列中进行处理和流动。由于指令比较简单,相应的控制也是比较简单的。

脉动阵列中的一个cell也非常简单。几组寄存器分别存储weight,activation和来自上方cell的部分和。weight从上向下传播,可以在weight path register中存储,保持不动或者传输给weight  register进行运算。weight register可以把数据发到乘法器进行处理,也可以直接传递给下方的cell;同样activation register也可以把数据发到乘法器进行处理,或者直接传递给右侧的cell。乘法器的输出和Sum in register的数值求和并传递给下方的cell。所有运算和传递都由控制寄存器控制(通过指令决定)。That’s it。

那么,在这样一个大规模的脉动阵列中怎么实现卷积操作呢?在Google的专利中有这样一个例子,可以看出一些端倪。如下图所示,activation输入被转换成一定的向量形式,作为脉动阵列的行输入。而3x3卷积核(Kernel)经过旋转,变成9个矩阵分别输入到脉动阵列的9个列。第一列计算的是卷积核和activation矩阵左上角的9个数据的卷积,得到输出矩阵的左上角的第一个数据,依此类推。

不难看出,为了要体现脉动阵列的运算效率(“keep the matrix unit busy”),需要对weight和activation进行很多形式上的转换。从TPU论文来看,似乎这项工作是由software stack中的“User Space driver”来完成,“It sets up and controls TPU execution, reformats data into TPU order, translates API calls into TPU instructions, and turns them into an application binary.”这项工作的运算量也不小,具体是怎么优化的就不得而知了。

差不多就这些内容了,欢迎大家和我讨论。

T.S.

一  点  感  想

Google TPU论文公开以来,引起各种争论。相信随着时间推移,它带来的影响还会不断显现。我甚至还看到这样的观点,“ASIC systolic arrays are going to flood the market with likely Chinese government subsidized hardware manufacturers cranking this stuff out. It just takes someone to come up with an open source reference design. Just like what happened in the Bitcoin world. This stuff is ancient design stuff, and I’m sure someone has the old iWarp designs and compilers in their basement somewhere! ”

(“ASIC脉动阵列架构只要有人搞出个开源的参考设计,就会席卷整个市场。就像当年比特币挖矿机一样。这个东西有年头了,而且我非常确定有些人在硬盘的某个角落还保留着当年的iWarp设计和编译器!”)

即使抛开这些,对我们来说,有机会结合Google的TPU深入了解一下脉动架构,确实是一件颇为有趣的事情。

参考:1. H. T. Kung, "Why systolic architectures?", IEEE Computer, Vol. 15, N°1, pp 37-46, 1982.2. Vivienne Sze, Yu-Hsin Chen, Tien-Ju Yang, Joel Emer, "Efficient Processing of Deep Neural Networks: A Tutorial and Survey",arxiv.org/pdf/1703.09039.pdf3.  Norman P. Jouppi, et al."In-Datacenter Performance Analysis of a Tensor Processing Unit",accepted by ISCA 2017


产业工程谷歌TPU
相关数据
神经网络技术
Neural Network

(人工)神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型,那时候研究者构想了「感知器(perceptron)」的想法。这一领域的研究者通常被称为「联结主义者(Connectionist)」,因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型,它们都是前馈神经网络:卷积神经网络(CNN)和循环神经网络(RNN),其中 RNN 又包含长短期记忆(LSTM)、门控循环单元(GRU)等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习,但也有一些为无监督学习设计的变体,比如自动编码器和生成对抗网络(GAN)。

比特币技术
Bitcoin

比特币是一种用去中心化、全球通用、不需第三方机构或个人,基于区块链作为支付技术的电子加密货币。比特币由中本聪于2009年1月3日,基于无国界的对等网络,用共识主动性开源软件发明创立。比特币也是目前知名度与市场总值最高的加密货币。 任何人皆可参与比特币活动,可以通过称为挖矿的电脑运算来发行。

深度神经网络技术
Deep neural network

深度神经网络(DNN)是深度学习的一种框架,它是一种具备至少一个隐层的神经网络。与浅层神经网络类似,深度神经网络也能够为复杂非线性系统提供建模,但多出的层次为模型提供了更高的抽象层次,因而提高了模型的能力。

图像处理技术
Image processing

图像处理是指对图像进行分析、加工、和处理,使其满足视觉、心理或其他要求的技术。 图像处理是信号处理在图像领域上的一个应用。 目前大多数的图像均是以数字形式存储,因而图像处理很多情况下指数字图像处理。

张量技术
Tensor

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

矽说
矽说

机器之心编辑

推荐文章
如何对比评价各种深度神经网络硬件?不妨给它们跑个分如何对比评价各种深度神经网络硬件?不妨给它们跑个分
机器之心机器之心
研学社 · 入门组 | 第六期:初入贝叶斯网络研学社 · 入门组 | 第六期:初入贝叶斯网络
Synced 深度研学社Synced 深度研学社
2
AI会给芯片设计带来什么?AI会给芯片设计带来什么?
矽说矽说
返回顶部