根據(jù)論文“數(shù)據(jù)中心的 TPU 性能分析”初略了解一下TPU實(shí)現(xiàn)
In-Datacenter Performance Analysis of a Tensor Processing Unit
結(jié)構(gòu)
PCIe3 Gen3 x16
通過PCIe3.0x16與主機(jī)通信，內(nèi)部模塊一般用帶寬256-byte的路徑連接。
Matrix Multiply Unit
Matrix Multiply Unit (矩陣計(jì)算單元)是TPU的核心部分歼捏，包含256x256個(gè)MAC計(jì)算單元届囚，進(jìn)行8-bit的乘加運(yùn)算竿屹。計(jì)算結(jié)果是16-bit各淀，到累加器部分進(jìn)行累加盾鳞。
?Weight FIFO
權(quán)重由片上緩存Weight FIFO輸給矩陣計(jì)算單元潘酗，是從片外的8 GiB的DRAM里讀取的。
?Unified Buffer
大小為24 MiB雁仲，用于保存中間結(jié)果仔夺，也是矩陣計(jì)算單元的輸入。通過一個(gè)可編程的DMA來負(fù)責(zé)在CPU memory和Unified Buffer間傳輸數(shù)據(jù)攒砖。下圖可以看到24 MiB的Unified Buffer占了幾乎的晶圓面積缸兔。
控制指令
TPU指令設(shè)計(jì)遵循的是CISC復(fù)雜指令集，每條指令的平均時(shí)鐘周期CPI是10到20吹艇，指令有很多惰蜜，下面選5個(gè)最重要的講一下：
- Read_Host_Memory
從CPU memory讀取數(shù)據(jù)到UB（Unified Buffer） - Read_Weights
從Weight Memory（片外DRAM）讀取權(quán)重到Weight FIFO作為矩陣運(yùn)算單元的輸入 - MatrixMultiply/Convolve
讓Matrix Unit執(zhí)行矩陣乘法或者卷積計(jì)算，從Unified Buffer到Accumulator受神。一次矩陣操作抛猖，輸入B256大小的輸入變量，于256x256的權(quán)重常數(shù)相乘鼻听，得到B256的輸出财著，花費(fèi)B個(gè)流水周期來完成計(jì)算。 - Activate
執(zhí)行非線性函數(shù)撑碴，如ReLU撑教、Sigmoid等。輸入來自Accumulators醉拓，輸出到Unified Buffer伟姐。 - Write_Host_Memory
把數(shù)據(jù)寫回到CPU Memory病线。
CISC MatrixMultiply指令為12個(gè)字節(jié)裕坊，其中3個(gè)為Unified Buffer地址; 2是累加器地址; 4是長(zhǎng)度（有時(shí)是卷積的2個(gè)維度）; 其余的是操作碼和標(biāo)志烟馅。
TPU加速的思路是讓矩陣計(jì)算單元一直被占用晰洒，因此為CISC指令設(shè)置了4級(jí)流水，每個(gè)指令在獨(dú)立的一級(jí)執(zhí)行蚀浆。 其目標(biāo)是通過將其執(zhí)行與MatrixMultiply指令重疊來隱藏其他指令（如Read_Weights等）木缝。但是涣旨，當(dāng)激活的輸入或權(quán)重?cái)?shù)據(jù)尚未就緒傍念，矩陣單元將進(jìn)入等待模式矫夷。
Systolic data flow
由于讀取大的SRAM的能耗比算術(shù)運(yùn)算高的多，矩陣計(jì)算單元通過脈沖執(zhí)行來減少對(duì) Unified Buffer的讀寫憋槐，進(jìn)而減少能量消耗双藕。下圖表示從左邊輸入的數(shù)據(jù)流和從上方加載的權(quán)重。
給定的256元乘法累加運(yùn)算通過矩陣作為對(duì)角波前移動(dòng)阳仔。此過程中權(quán)重是預(yù)先加載的忧陪，并且與數(shù)據(jù)同步。通過對(duì)控制邏輯和數(shù)據(jù)流水線操作近范，使得對(duì)于MMU而言嘶摊，256個(gè)輸入一次讀取的，并且它們立即更新到256個(gè)累加器中的一個(gè)個(gè)具體位置评矩。
matrix單元就是一個(gè)典型的脈動(dòng)陣列叶堆。weight由上向下流動(dòng)，activation數(shù)據(jù)從左向右流動(dòng)斥杜。在最下方有一些累加單元虱颗，
cell內(nèi)部結(jié)構(gòu)如上，幾組寄存器分別存儲(chǔ)weight蔗喂，activation和來自上方cell的部分和忘渔。weight從上向下傳播，可以在weight path register中存儲(chǔ)缰儿，保持不動(dòng)或者傳輸給weight register進(jìn)行運(yùn)算畦粮。weight register可以把數(shù)據(jù)發(fā)到乘法器進(jìn)行處理，也可以直接傳遞給下方的cell乖阵；同樣activation register也可以把數(shù)據(jù)發(fā)到乘法器進(jìn)行處理宣赔，或者直接傳遞給右側(cè)的cell。乘法器的輸出和Sum in register的數(shù)值求和并傳遞給下方的cell瞪浸。所有運(yùn)算和傳遞都由控制寄存器控制（通過指令決定）拉背。
