脉动阵列在二维矩阵乘法及卷积运算中的应用

脉动阵列（Systolic Array）本身是一个“古老”的概念，在1982年就已经提出了，可是，最近由于Google的TPU采用了这个结构作为计算的核心结构，脉动阵列又一次地火了起来。我也是因为关注TPU才开始去了解脉动阵列的，但是由于目前脉动的阵列比较零散，在搞明白脉动阵列以及TPU怎么使用脉动阵列进行卷积的过程中走了很多“弯路”，所以就用这篇文章作为我学习脉动阵列的笔记，也记录一下我自己对脉动阵列的理解，如果理解有误不到位的话，欢迎指出。

一、什么是脉动阵列？

脉动阵列，本身的核心概念就是让数据在运算单元的阵列中进行流动，减少访存的次数，并且使得结构更加规整，布线更加统一，提高频率。

可以看下面这个图，传统计算结构和脉动阵列结构的对比。左边是传统的计算架构，可用于各种形式的计算。CPU、GPU就是这种架构，用寄存器存储数据，一个程序告诉ALU从寄存器中取数，然后执行一种操作（例如加法、乘法或者逻辑操作），然后再把结果写会指定的寄存器中。脉动阵列，第一个ALU取数，经过处理后被传递到下一个ALU同时第二个数据进入第一个ALU。依次类推。在第一个数据到最后一个ALU之后，每个周期都能得到一个结果。这样，脉动阵列可以平衡IO读写与运算，在消耗较小memory带宽的情况下提高吞吐率，有效解决数据存取速度远远大于数据处理速度的结构。

脉动阵列本身只是一个有数据流动的结构，根据不同的应用会有不同的数据进行流动以及不同的流动方向。

 

二、脉动阵列做矩阵乘法

使用脉动阵列进行两个二维矩阵的乘法我看到的有两种做法：

为了方便解释，举个栗子

 

1. 在脉动阵列中流动的是X,W，而每一步的Y保存在每个PE或者cell中

每一个PE就是一个乘加器MAC。

值得注意的是，在使用脉动矩阵进行矩阵计算的时候需要对数据调整好形式，按照一定顺序，分时进入脉动阵列。

 

2. 在脉动阵列中流动的数据是X和Y的中间结果，W存储在每个cell或者PE中

 

 

这种传递参数的方式在Google Cloud的Blog[1]中也有提及，估计是TPU中脉动阵列传递数据的模式也是这样计算，在PE中固定Weight，传播部分和。

这两种方式大同小异，只是选择固定和流动的数据不一样。基本上，就是在W, X, Y三个变量中固定其中一个在PE单元，其余的两个数据在脉动阵列中进行传递。

三、脉动阵列在卷积神经网络inference中的应用

我是在调研卷积神经网络加速器的时候接触到脉动阵列，继而去进一步了解它的计算过程的。而采用脉动阵列作为卷积神经网络的加速器架构一个典型的例子就是开篇提到的TPU。

TPU的整体结构如下图，

通过上图的性能描述，可以看出TPU的特点就是计算单元很多（256x256），而带宽却没有那么大（Weight端30GiB/s,在700MHZ的情况下，一个cycle传输的weight只有352位左右）。PE多，带宽小就是由Matrix Multiply Unit中的脉动阵列带来的优势。下面一张图说明了是TPU中的脉动阵列采用的数据流动方式，Input Data向右传播，PartialSum向下传播。控制单元产生的控制信号控制数据的流动。

 

具体，TPU怎么计算卷积的过程细节在TPU的论文[4]中没有提及，我也只能从Google几个专利中窥见一斑。这里提一句，Google卷积神经网络处理器的几个专利分别是：

[Ros15b] Computing Convolutions Using a Neural NetworkProcessor. Patent Application NO.62/164,902, US20160342889A1, WO2016186811A1

[Ros15d] Rotating Data for Neural NetworkComputation.

Patent Application NO.62/164,908,US20160342893, WO2016186826A1

其中[Ros15d]讲的更仔细一点。我下面说的也是基于[Ros15d]中的内容的个人理解

在专利中，给出了一个cell的结构。cell中包含一个activation（input）寄存器，放activationinput的值。activation寄存器可以从左边的cell单元或者unifiedbuffer中获得数据，这个取决于这个cell在脉动阵列中的位置。cell中还有一个用来放weight的寄存器。同样，根据cell在脉动阵列中位置，weight数据的来源可以是上方的cell单元或者weightfetch interface。剩余的就是一个乘法电路和求和电路。求和电路将乘法得到的乘积和sum in register中的数加起来。在专利中，还提到了一点，这个在图中没有画出来，cell还包含一个controlregister。这个寄存器存的是一个控制信号，决定是否把当前的weight或者activation的是否传递到临近的cell中。

一个cell的结构

虽然，从cell的结构看，weight、input、sum都可以传播，但是在真正计算中，这三者的传播并不是同时流动。根据专利中所述，weight是提前先被放在weight register上的，并且会保持多个周期不进行传递，直至需要更新weight的时候，weight才会进行传递。在计算中，主要传递的是activation和partial sum。

“The weight register 502can statically store the weight input such that as activation inputs aretransferred to the cell, e.g., through the activation register, over multipleclock cycles, the weight input remains within the cell and is transferred to anadjacent cell. Therefore, the weight input can be applied to multipleactivation inputs, e.g., using the multiplication circuitry, and respective accumulatedvalues can be transferred to an adjacent cell.”

卷积与矩阵乘法还是有些许差别的，那怎么使用脉动阵列进行卷积运算呢？

Google的专利中提出，为了进行卷积计算，系统会将卷积计算转为二维矩阵乘法。Google专利中给出了下右图的例子来说明如下作图的矩阵结构怎么用脉动阵列计算卷积。一般来说，神经网络处理器分别将activation inputs和weight inputs放到矩阵的行和列中。activation和weight在脉动阵列中右传、下传，直到到达指定的cell中的指定寄存器。一旦input放好了，通过控制信号，处理器就开始用存在cell中的inputs计算得到output。处理器在把activation矩阵送进systolic array之前会先把矩阵“压扁”。如下图，矩阵在深度上有三个通道602, 604, 606。不同通道上的二维矩阵式送到脉动阵列中的不同行。在那张脉动阵列的图中，可以看到，第一行的输入是左边矩阵602这一层，第二行的输入是右边矩阵604这一层。weight那边，也有好多kernel，比方说下面这个例子中有Kernel A-D，不同kernel送到不同列。当把一个矩阵数据送到一个阵列中，矩阵的第一个元素在一个clock cycle中被送到cell中，在下一个的clock cycle，下一个元素被送到cell中，第一个元素被传递到相邻的cell中。

 

下面这是专利中的一个例子，解释了weight在3x3脉动阵列的传递方式。这个方法和之前解释的脉动阵列计算乘法基本是一致的。

Google专利中还提到了以个旋转数据的例子。讲真，这个例子我看的很模糊，感觉和之前的硬件和解释没法对应上。这里就贴一张图，和简单翻译几句。

activation inputs是5x5的是170x170原图中的一部，5x5这个尺寸的选取也是根据kernel大小选择的。把activation inputs变成vector inputs ，作为行输入。这四个向量是1102这个矩阵的四个象限按从上到下，再从左到右的一个排列。kernel呢，则经过重新的排列组合，变成9个矩阵，分别输入到脉动矩阵的9个列上。按照专利中的说法，1114就是脉动阵列，但是我这里就挺费解的，感觉和之前那张systolic array图中的列输入不同。（还有疑惑的一点是，如果按照这种方法，activation input是怎么在同一行之中脉动起来的，一行中的数据有交错的现象，不是一个顺序的。如果要硬是要使用脉动的话，我认为会降低计算效率呀。）

Google 没有公开TPU的细节，在我找到的这个专利上也比较模糊不清，感觉专利上只是列出目前的结构能够实现的功能，用到的大概思路，比较零散还是没有很好的理解怎么做卷积数据怎么重排，当然，也有可能是我才疏学浅，没能参悟。如果有人能告诉我到底TPU是具体怎么在脉动阵列上传递数据的，欢迎。

 

另外，在2017DAC会议上也有一篇利用脉动阵列来加速卷积神经网络的Inference [5]。大致浏览了一下，硬件的结构基本上与其他的脉动阵列式类似的。在下面这张图中，WB（Weight Buffer），IB（Input Buffer），OB（Output Buffer），在进行计算时，脉动阵列中Weight和Input分别向右、向下传播。在Buffer之间WB之间可由上向下传播参数，在IB之间可由左至右传播参数。最后计算的结果由PE一级级向上传至OB（图中的实线路径）。在PE中，一个乘加器，其他的都是为了缓存数据的寄存器或者进行数据选择的MUX。在Buffer的设计中，论文的作者使用了double buffering，因为每个buffer都会进行两个方向的传播，double buffering的方法是其能pipeline起来。

四、总结

上边就是我对我看过一些脉动阵列的总结。总体来说，个人认为，脉动阵列没有传说中的那么“神奇”，我更多地认为它就是一个用计算单元换I/O带宽的方法。在计算单元很多，带宽不足的情况下，是一个很好的备选结构。脉动阵列还有一个优点就是布线比较简洁，可以很好提高频率。脉动阵列也会带来一些问题，就是怎么去分配数据，向TPU中，如果采用rotatingdata的话，就会需要额外的资源完成这项任务，具体这个的代价在论文和专利中都没有细说。

 

参考资料：

[1] " An in-depth look at Google’s firstTensor Processing Unit (TPU) "

https://cloud.google.com/blog/big-data/2017/05/an-in-depth-look-at-googles-first-tensor-processing-unit-tpu

[2] "脉动阵列 - 因Google TPU获得新生" https://zhuanlan.zhihu.com/p/26522315

[3] " Should We All Embrace SystolicArrays?"

https://www.linkedin.com/pulse/should-we-all-embrace-systolic-arrays-chien-ping-lu

[4] Norman P. Jouppi, etal."In-Datacenter Performance Analysis of a Tensor Processing Unit"，accepted byISCA 2017

[5] Wei, Xuechao, et al. "AutomatedSystolic Array Architecture Synthesis for High Throughput CNN Inference onFPGAs." Proceedings of the 54th Annual Design Automation Conference 2017.ACM, 2017.