Simple TPU的設計和性能評估

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

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

 

    谷歌TPU概述和簡化

    TPU中的脈動陣列及其實現

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

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

    Simple TPU的設計和性能評估

    SimpleTPU實例：圖像分類

    拓展

    TPU的邊界（規劃中）

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

 

1. 完成SimpleTPU的設計

    在 谷歌TPU概述和簡化中給出過SimpleTPU的框圖，以下圖所示。

    在TPU中的脈動陣列及其實現中介紹了矩陣/卷積計算中的主要計算單元——乘加陣列（上圖4），完成了該部分的硬件代碼並進行了簡單的驗證；在 神經網絡中的歸一化和池化的硬件實現中介紹了卷積神經網絡中的歸一化和池化的實現方式（上圖6），同時論述了浮點網絡定點化的過程，並給出了Simple TPU中重量化的實現方式，完成了該部分的硬件代碼並進行了驗證。

    在 TPU中的指令並行和數據並行中對整個處理單元的體系結構進行了分析和論述，包括指令並行和數據並行兩個方面。那麼，如何在TPU中的指令並行和數據並行中提到的設計思路下，將TPU中的脈動陣列及其實現和神經網絡中的歸一化和池化的硬件實現中提到的計算單元充分的利用，是完成Simple TPU設計的最後一部。根據SimpleTPU的框圖可知，須要實現的功能包括

• 指令的取指和譯碼（上圖4）
• Weight的讀取（上圖2）
• 各個執行單元的控制和調度（上圖1）
• 讀取圖像和寫回結果（上圖5）

    在SimpleTPU的設計中，指令的取指和譯碼和Weight的讀取功能都較爲簡單，可直接參照代碼。

    在對各個執行單元進行控制和調度時須要確保各個單元能夠共同執行，沒有相互之間的數據以來關係。

    除此以外，還須要單獨實現讀取圖像和寫回結果的功能。SimpleTPU中只關注核心的計算功能，該部分功能並未進行優化，後續對實現效果進行分析時，也會將該部分運行時間排除在外。

    至此，Simple TPU的設計基本完成了，代碼可參見https://github.com/cea-wind/SimpleTPU。

2. SimpleTPU的特性

    SimpleTPU的主要特性包括

•     支持INT8乘法，支持INT32的累加操做
•     採用VLIW進行指令並行
•     採用向量體系結構進行數據並行

    SimpleTPU依照Google TPU V1的設計思路，能夠完成神經網絡推理過程當中的大部分運算。依據設計，支持的運算包括（理論）

運算	說明
Conv3d	in_channels:資源受限 out_channels:資源受限 kerner_size:幾乎無限制 stride:幾乎無限制 padding:幾乎無限制 dilation:幾乎無限制 groups:極有限支持，架構限制 bias:支持
ConvTranspose3d	同上
Maxpool2d	kernel_size:幾乎無限制 stride:幾乎無限制 padding:幾乎無限制
Avgpool2d	同上
Relu	僅支持RELU做爲非線性函數
BatchNorm2d	推理中BatchNorm2d被融合到Conv或者Pool完成
Linear	資源受限
UpscalingNearest2D	屢次調用pool完成
UpscalingBilinear2D	屢次調用avgpool完成

     

    其中，資源受限表明該參數的取值範圍有限，主要受限於SimpleTPU的存儲設計等；幾乎無限制表示其取值僅受到寄存器位寬等限制。因爲架構設計上的問題，SimpleTPU對groupconv支持極爲有限，在不合適的參數下效率可能遠低於普通卷積；相似的，Google TPU也不能很好支持groupconv，並明確告知不制止depthwise conv（極度稀疏化的group conv）。

    BatchNorm2d在推理過程當中實際上時進行逐點的乘法和加法，其中加法計算能夠融合到下一層或者上一層的卷積計算中進行，乘法計算能夠和pooling計算融合。在SimpleTPU設計，Pooling模塊實際上一直在工做，即便網絡中沒有pooling層，SimipleTPU增長了一個1*1，stride=1的pooling層進行等價。

    Upscaling操做經過pooling完成計算。這是由於在SimpleTPU中，reshape操做（支持的）是沒有代價的。pooling操做能夠完成雙線性插值的計算，所以能夠完成upscaling中的全部數值的計算。能夠理解爲經過pooling+reshape完成了upscaling的計算。

   

3. SimpleTPU的性能

    Simple TPU設計了一個32×32的int8乘加陣列計算矩陣乘法和卷積，和一個1×32的int32乘法陣列進行池化和歸一化的計算。根據Xilinx HLS工具的綜合結果，在UltraScale+系列的FPGA器件上，工做頻率可達500MHz。所以SimpleTPU的算力約爲

32×32×500MHz×2 = 1Tops

    做爲對比，GoogleTPU V1的算力約爲91Tops（int8），差別主要在SimpleTPU的規模爲其1/64，同時在FPGA上的工做頻率會低於ASIC的工做頻率。

    依據設計，SimpleTPU在適合的任務下會有很高的運行效率，TPU中的指令並行和數據並行中針對這一點又更爲具體的描述。從宏觀上看，SimpleTPU的各個運行單元能夠流水並行的，即

    而針對網絡中計算量最大的全鏈接層和卷積層，針對性設計的乘法整列和向量計算的設計方法可讓其在每一個時鐘週期都完成有效的乘加計算；這意味着和CPU相比，SimpleTPU能夠達到極高的效率。