動手寫一個簡單版的谷歌TPU
谷歌TPU是一個設計良好的矩陣計算加速單元,可以很好的加速神經網路的計算。本系列文章將利用公開的TPU V1(後簡稱TPU)相關資料,對其進行一定的簡化、推測和修改,來實際編寫一個簡單版本的谷歌TPU。計劃實現到行為模擬為止,僅為更確切的瞭解TPU的優勢和侷限性,暫無在FPGA等硬體上進一步實現的計劃。
系列目錄
谷歌TPU概述和簡化
基本單元-矩陣乘法陣列
基本單元-歸一化和池化(待發布)
TPU中的Instruction (待完成)
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的設計關鍵在於充分利用這一乘加陣列,使其利用率儘可能高。
結構圖中其他的部分基本都是為儘可能跑滿這個矩陣計算陣列服務的,據此有以下設計
- 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內部計算的實現,即下圖紅框所示。
由於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進行開發是更為合適的選擇。
按規劃之後將給出兩個關鍵計算單元的實現,以及控制邏輯和指令的設計方法;最後將給出一個實際的神經網路及其模擬結果和分析。