動手寫一個簡單版的谷歌TPU

深度學習飛速發展過程當中，人們發現原有的處理器沒法知足神經網絡這種特定的大量計算，大量的開始針對這一應用進行專用芯片的設計。谷歌的張量處理單元（Tensor Processing Unit，後文簡稱TPU）是完成較早，具備表明性的一類設計，基於脈動陣列設計的矩陣計算加速單元，能夠很好的加速神經網絡的計算。本系列文章將利用公開的TPU V1相關資料，對其進行必定的簡化、推測和修改，來實際編寫一個簡單版本的谷歌TPU，以更確切的瞭解TPU的優點和侷限性。

動手寫一個簡單版的谷歌TPU系列目錄

    谷歌TPU概述和簡化

    TPU中的脈動陣列及其實現

    神經網絡中的歸一化和池化的硬件實現

    TPU中的指令並行和數據並行

    Simple TPU的設計和性能評估

    SimpleTPU實例：圖像分類

    拓展

    TPU的邊界（規劃中）

    從新審視深度神經網絡中的並行（規劃中）

1. TPU設計分析

    人工神經網絡中的大量乘加計算（譬如三維卷積計算）大多均可以概括成爲矩陣計算。而以前有的各種處理器，在其硬件底層完成的是一個（或多個）標量/向量計算，這些處理器並無充分利用矩陣計算中的數據複用；而Google TPU V1則是專門針對矩陣計算設計的功能強大的處理單元。參考Google公開的論文In-Datacenter Performance Analysis of a Tensor Processing Unit，TPU V1的結構框圖以下所示

    結構框圖中最受矚目的是巨大的Matrix Multiply Unit，共計64K的MAC能夠在700MHz的工做頻率下提供92T int8 Ops的性能。這樣一個陣列進行矩陣計算的細節將會在基本單元-矩陣乘法陣列進行更進一步的闡述。TPU的設計關鍵在於充分利用這一乘加陣列，使其利用率儘量高。

    結構圖中其餘的部分基本都是爲儘量跑滿這個矩陣計算陣列服務的，據此有如下設計

1. Local Unified Buffer 提供了256×8b@700MHz的帶寬（即167GiB/s，0.25Kib×700/1024/1024=167GiB/s），以保證計算單元不會由於缺乏Data in而閒置；
2. Local Unified Buffer 的空間高達24MiB，這意味着計算過程的中間結果幾乎無需和外界進行交互，也就不存在由於數據帶寬而限制計算能力的狀況；
3. Matrix Multiply Unit中每一個MAC內置兩個寄存器存儲Weight，當一個進行計算時另外一個進行新Weight的載入，以掩蓋載入Weight的時間；
4. 30GiB/s的帶寬完成256×256Weight的載入須要大約1430個Cycles，也就意味着一組Weight至少須要計算1430Cycles，所以Accumulators的深度須要爲2K（1430取2的冪次，論文中給出的數值是1350，差別未知）；
5. 因爲MAC和Activation模塊之間須要同時進行計算，所以Accumulators須要用兩倍存儲來進行pingpang設計，所以Accumulators中存儲的深度設計爲4k；

    所以從硬件設計上來看，只要TPU ops/Weight Byte達到1400左右，理論上TPU就能以接近100%的效率進行計算。但在實際運行過程當中，訪存和計算之間的調度，讀寫之間的依賴關係（譬如Read After Write，須要等寫完才能讀），指令之間的流水線和空閒週期的處理都會在必定程度影響實際的性能。

    爲此，TPU設計了一組指令來控制其訪問存和計算，主要的指令包括

• Read_Host_Memory
• Read_Weights
• MatrixMultiply/Convolve
• Activation
• Write_Host_Memory

    全部的設計都是爲了讓矩陣單元不閒下來，設計但願全部其餘指令能夠被MatrixMultiply指令所掩蓋，所以TPU採用了分離數據獲取和執行的設計（Decoupled-access/execute），這意味着在發出Read_Weights指令以後，MatrixMultiply就能夠開始執行，不須要等待Read_Weight指令完成；若是Weight/Activation沒有準備好，matrix unit會中止。

    須要注意的是，一條指令能夠執行數千個週期，所以TPU設計過程當中沒有對流水線之間的空閒週期進行掩蓋，這是由於因爲Pipline帶來的數十個週期的浪費對最終性能的影響不到1%。

    關於指令的細節依舊不是特別清楚，更多細節有待討論補充。

2. TPU的簡化

    實現一個完整的TPU有些過於複雜了，爲了下降工做量、提升可行性，須要對TPU進行一系列的簡化；爲作區分，後文將簡化後的TPU稱爲SimpleTPU。全部的簡化應不失TPU自己的設計理念。

    TPU中爲了進行數據交互，存在包括PCIE Interface、DDR Interface在內的各種硬件接口；此處並不考慮這些標準硬件接口的設計，各種數據交互均經過AXI接口完成；僅關心TPU內部計算的實現，更準確的來講，Simple TPU計劃實現TPU core，即下圖紅框所示。

    因爲TPU的規模太大，乘法器陣列大小爲256×256，這會給調試和綜合帶來極大的困難，所以此處將其矩陣乘法單元修改成32×32，其他數據位寬也進行相應修改，此類修改包括

Resource	TPU	SimpleTPU
Matrix Multiply Unit	256*256	32*32
Accumulators RAM	4K*256*32b	4K*32*32b
Unified Buffer	96K*256*8b	16K*32*8b

    因爲Weight FIFO實現上的困難（難以採用C語言描述）, Weight採用1K*32*8b的BRAM存放，Pingpang使用；

    因爲Matrix Multiply Unit和Accumulators之間的高度相關性，SimpleTPU將其合二爲一了；

    因爲Activation和Normalized/Pool之間的高度相關性，SimpleTPU將其合二爲一了（TPU自己可能也是這樣作的），同時只支持RELU激活函數；

    因爲並不清楚Systolic Data Setup模塊到底進行了什麼操做，SimpleTPU將其刪除了；SimpleTPU採用了另外一種靈活而又簡單的方式，即經過地址上的設計，來完成卷積計算；

    因爲中間結果和片外緩存交互會增長instruction生成的困難，此處認爲計算過程當中無需訪問片外緩存；(這也符合TPU自己的設計思路，但因爲Unified Buffer大小變成了1/24，在這一約束下只可以運行更小的模型了)

    因爲TPU V1並無提供關於ResNet中加法操做的具體實現方式，SimpleTPU也不支持ResNet相關運算，但能夠支持channel concate操做；（雖然有多種方式實現Residual Connection，但均需添加額外邏輯，彷佛都會破壞原有的結構）

    簡化後的框圖以下所示，模塊基本保持一致

3. 基於Xilinx HLS的實現方案

    通常來講，芯片開發過程當中多采用硬件描述語言（Hardware Description Language），譬如Verilog HDL或者VHDL進行開發和驗證。但爲了提升編碼的效率，同時使得代碼更爲易懂，SimpleTPU試圖採用C語言對硬件底層進行描述；並經過HLS技術將C代碼翻譯爲HDL代碼。因爲以前使用過Xilinx HLS工具，所以此處依舊採用Xilinx HLS進行開發；關於Xilinx HLS的相關信息，能夠參考高層次綜合（HLS）-簡介，以及一個簡單的開發實例利用Xilinx HLS實現LDPC譯碼器。

    雖然此處選擇了Xilinx HLS工具，但據我所瞭解，HLS可能並不適合完成這種較爲複雜的IP設計。儘管SimpleTPU已經足夠簡單，但依舊沒法在一個函數中完成全部功能，而HLS並不具備函數間相對複雜的描述能力，兩個模塊之間每每只能是調用關係或者經過FIFO Channel相連。但因爲HLS易寫、易讀、易驗證，此處依舊選擇了HLS，並經過一些手段規避掉了部分問題。真實應用中，採用HDL或者HDL結合HLS進行開發是更爲合適的選擇。

    按規劃以後將給出兩個關鍵計算單元的實現，以及控制邏輯和指令的設計方法；最後將給出一個實際的神經網絡及其仿真結果和分析。