自Google在2016年推出第一代自研人工智能芯片Tensor Processing Unit, TPU，經過數年發展已升級到了第4代 TPU v4 (截止到2022年底)。TPU架構設計也是通過高效地并行運算大量乘法和累加操作來實現深度學習卷積層，全連接層等網絡層的高效計算。
脈動陣列
脈動陣列(systolic array)是由大量緊密耦合的PE組成陣列網絡，是一種數據流的架構設計。脈動陣列中每一個PE都會與周圍一個或者多個的PE進行數據通信，PE的功能相對簡單，系統通過大量PE的并行計算來實現高效的計算，通過數據在不同PE間的流動最大化的實現數據的復用和減少芯片在運算過程中的內存訪問的次數，不僅可以節省內存帶寬還降低了因內存訪問帶來的功耗損失。
在傳統計算系統的運算過程中，處理單元PE先從主存中讀取數據，然后再進行運算操作，運算完成后結果再寫回到存儲器中，因此訪存的速度稱為了整個系統的處理性能的瓶頸。相比于CPU通過多級緩存的方式來解決訪存速度瓶頸的問題，脈動架構則是讓數據在不同處理單元中流動來減少訪問主存帶來的性能瓶頸。如圖所示，在一個一維脈動陣列的運算過程中，數據首先從主存中進入第一個PE，經過處理后被傳遞到下一個PE，同時下一個數據進入第一個PE，以此類推，數據以流動的方式在不同PE之間傳遞，直到完成全部計算后才返回主內存。所以脈動陣列架構實現了輸入數據的多次復用，省去了數據被寫回主存后再讀取的過程，減少主存的訪問次數，因此脈動陣列可以在較小的內存帶寬的情況下實現高吞吐率。
脈動陣列
脈動陣列架構有如下幾個特點：
PE功能結構簡單，實現成本低，可以集成大量PE來提高并行計算能力。
由大量同構的PE組成的一維、二維或樹的陣列結構，可靈活拓展。
不同PE之間采用流水線的方式進行數據通訊，實現高效的數據復用。
數據只能在相鄰PE間流動，只適合特定的算法，比如矩陣運算，卷積等。
TPU架構設計
TPU使用了脈動陣列架構設計，數據在不同方向上以固定的時間間隔送入陣列中的 PE進行運算，在經過多步計算后最后將結果匯總輸出來。脈動陣列只適合用于非常簡單且規整的運算，而矩陣乘和卷積剛好符合脈動陣列的運算特性。
第一代TPU的脈動陣列結構實現如圖所示，與GPU類似，TPU通過PCI-E總線與主機CPU連接，TPU的指令由CPU完成編碼以簡化硬件設計和調試。矩陣乘法單元(Matrix Multiply Unit, MXU)是主要的計算單元，其主要功能為完成矩陣的乘法運算。矩陣乘法單元周圍有三個功能不同的數據緩存區以及專用的Activitation, Normalize和Pool單元，其中三個數據緩存區分別用于緩存權重隊列(Weight FIFO)和輸入數據緩沖區(Unified Buffer, UB)和輸出數據累加單元(Accumulator, Acc)。
TPU架構
具體執行時，指令和數據從主機接口(host interface)進入TPU，復用率更高的權重參數會被預先加載到Weight FIFO，輸入數據會被加載到統一緩沖區UB。輸入數據和權重參數在矩陣乘法單元里完成矩陣乘法運算后結果再送入累加單元Acc，在Acc完成部分和累加后根據模型設計需要，選擇是否進行送入Activitation, Normalize和Pool單元執行對應操作，最后再將結果送回統一緩沖區UB。
在第一代TPU硬件的布局中，矩陣乘單元和激活單元一共占30%的面積，矩陣乘單元(MXU)有256x256個MACs，每個時鐘周期能完成256x256個8bit數據的乘加運算，Acc是大小為4MiB的32bit累加器。UB大小為24MiB，占29%的面積，其可以通過DMA直接與Host CPU交互，用于緩存輸入數據或保存中間計算結果。Weight FIFO的深度為4，權值參數從片外內存中讀取。
TPU布局
矩陣乘法單元(Matrix Multiply Unit, MXU)是一個典型的脈動陣列，如圖所示，weight由上向下流動，數據從左向右流動。輸入的像素數據從乘法矩陣左側進入，并從左向右傳播以重復利用中間數據，由于權重是預先加載的，因此隨著輸入數據的前進，可以快速算出乘積結果，在通過控制通路進入下一步累加計算。
矩陣乘單元數據流
TPU架構演進
TPU v1是谷歌第一代AI專用芯片，主要專注與處理推理任務，Google在推出用于推理的TPU v1之后開始開發和設計針對訓練的第二代TPU。與TPU v1相比，TPU v2的有如下的改進點：
1.TPU v2每個芯片有兩個Tensor Core。
2.將固定的激活函數(Activation Pipeline)改為可編程性更高的向量單元(Vector Unit)。
3.使用一個向量存儲器(Vector Memory)代替Accumulator和Activation Storage中的緩存。
4.矩陣乘法單元作為向量單元的協處理器直接與向量單元連接，增加其可編程性。
5.使用HBM代替DDR3，并改為與向量存儲區相連，可以提供更高的帶寬和讀寫速度。
6.在HBM和向量存儲區之間增加互連模塊(Interconnect)，提供更強的擴展能力用于TPU之間的連接。
7.增加了Scalar Unit, Transpose/Permute Unit等單元， 對Transpose等操作做硬件加速。
TPU v2架構
TPUv3在TPUv2的基礎上進一步的提升了性能，其在時鐘頻率、內存帶寬和芯片間帶寬上提升了30%，矩陣乘法單元MXU的數量增加了1倍，HBM容量提升了1倍，可連接節點數增加了4倍。
TPU v1 v2 v3架構差異
TPU v4出于成本的考慮，Google對用于訓練的和用于推動TPU進行了分開設計，用于訓練的TPU v4有兩個Tensor Core，而用于推理的PU v4i只有一個Tensor Core，在通用性、性能和成本上進行了取舍和平衡。在TPU v4i中，單個Tensor Core包含4個矩陣乘法單元MXU，是TPU v3的兩倍。Google在TPU v4i的設計中還增加了性能計數器(Performance counter)以便于協助編譯器更好地掌握芯片的運行情況。
鄭重聲明：文中僅代表作者觀點，不代表本公司觀點，本文版權歸原作者所有，轉載文章僅為傳播更多信息之目的，如作者信息標記有誤，請第一時間聯系我們修改或者刪除，感謝您的關注！