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