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