脈動陣列在二維矩陣乘法及卷積運算中的應用

脈動陣列（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.