深度學習飛速發展過程中,人們發現原有的處理器無法滿足神經網絡這種特定的大量計算,大量的開始針對這一應用進行專用芯片的設計。谷歌的張量處理單元(Tensor Processing Unit,后文簡稱TPU)是完成較早,具有代表性的一類設計,基於脈動陣列設計的矩陣計算加速單元,可以很好的加速神經網絡的計算。本系列文章將利用公開的TPU V1相關資料,對其進行一定的簡化、推測和修改,來實際編寫一個簡單版本的谷歌TPU,以更確切的了解TPU的優勢和局限性。
動手寫一個簡單版的谷歌TPU系列目錄
拓展
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的設計關鍵在於充分利用這一乘加陣列,使其利用率盡可能高。
結構圖中其他的部分基本都是為盡可能跑滿這個矩陣計算陣列服務的,據此有以下設計
- Local Unified Buffer 提供了256×8b@700MHz的帶寬(即167GiB/s,0.25Kib×700/1024/1024=167GiB/s),以保證計算單元不會因為缺少Data in而閑置;
- Local Unified Buffer 的空間高達24MiB,這意味着計算過程的中間結果幾乎無需和外界進行交互,也就不存在因為數據帶寬而限制計算能力的情況;
- Matrix Multiply Unit中每個MAC內置兩個寄存器存儲Weight,當一個進行計算時另一個進行新Weight的載入,以掩蓋載入Weight的時間;
- 30GiB/s的帶寬完成256×256Weight的載入需要大約1430個Cycles,也就意味着一組Weight至少需要計算1430Cycles,因此Accumulators的深度需要為2K(1430取2的冪次,論文中給出的數值是1350,差異未知);
- 由於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進行開發是更為合適的選擇。
按規划之后將給出兩個關鍵計算單元的實現,以及控制邏輯和指令的設計方法;最后將給出一個實際的神經網絡及其仿真結果和分析。