脉动阵列(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的整体结构如下图,
通过上图的性能描述,可以看出