系列目錄

    谷歌TPU概述和簡化

    基本單元-矩陣乘法陣列

    基本單元-歸一化和池化（待發布）

    TPU中的指令集

    SimpleTPU例項： （計劃中）

    拓展

    TPU的邊界（規劃中）

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

  

     TPU V1定義了一套自己的指令集，雖然在介紹處理器時，往往會先談指令集架構，但此處卻把它放到了最後，這主要基於兩個原因；其一在於個人的對處理器不太瞭解，這也是主要原因，其二在於公開資料中並沒有TPU指令集的細節和TPU微架構的描述。從資料流和計算單元出發對TPU進行分析固然容易很多，但如果想理解TPU的設計思想，依舊需要回到其架構設計上進行分析。這一部分內容有些超出了我現有的能力，不當之處還請多多指正。

    本文主要探討從架構設計上看，TPU時如何做高效能和高效能的設計。高效能的多來自於並行，因此本文分別討論了指令並行和資料並行的設計方法。由於論文中並未描述TPU指令集的具體設計，除特別說明外，本文關於TPU指令集的探討均為推測；另外，SimpleTPU的指令設計並不系統/完整，此處僅闡明設計中的幾種基本思想。

1. TPU的指令集

     TPU的指令集採用CISC設計，共計有十多條指令，主要的五條指令包括

1. Read_Host_Memory 將資料從CPU的記憶體中讀取到TPU的Unified Buffer上
2. Read_Weights 將weight從記憶體中讀取到TPU的 Weight FIFO 上.
3. MatrixMultiply/Convolve 執行卷積或矩陣乘法操作.
4. Activate 執行人工神經網路中的非線性操作和Pooling操作（如有）
5. Write_Host_Memory 將結果從Unified Buffer寫回CPU記憶體.

    從給出的五條指令可以看出，TPU的指令集設計和通用處理器有很大的不同。指令需要顯示指定資料在記憶體和片上buffer之間搬移的過程。而執行指令（MatrixMultiply）直接指定了Buffer的地址，指令上並不能看到一系列通用暫存器。這是因為TPU本質上還是一個專用的處理晶片，其高效能和高效能都是建立在失去一定靈活性的前提下的。為了獲得更高的效能，可以採用一系列的常規方法進行設計，包括

• 指令並行，即一次性處理更多指令，讓所有執行單元高效執行
• 資料並行，即一次性處理多組資料，提高效能

     後文會針對這兩點做進一步描述，並簡單討論TPU設計中的更多其他的優化方法和方向。

2. 指令並行

2.1 Simple TPU中的流水線

    為了提高吞吐率和時鐘頻率，處理器通常使用流水線設計，經典的五級流水線設計一般如下所示

	clk0	clk1	clk2	clk3	clk4	clk5	clk6	clk7
instruction 0	IF	ID	EX	MEM	WB
instruction 1		IF	ID	EX	MEM	WB
instruction 2			IF	ID	EX	MEM	WB
instruction 3				IF	ID	EX	MEM	WB

    其中，IF指取指(insturction fetch)，ID指指令譯碼（instruction decode），EX指執行（Execute），MEM指記憶體讀寫（Memory Access），WB指寫回暫存器(Write back)。採用流水線設計可以提高效能，如果不採用流水線設計，那麼instruction1需要在clk5才能開始進行IF，嚴重影響其效能；如果在同一週期完成IF/ID/EX/MEM/WB的功能，由於邏輯極其複雜，會嚴重影響工作頻率。

    TPU論文中介紹其採用四級流水線設計，Simple TPU中採用了兩級流水線，來完成控制過程。

	clk0	clk1	clk2	clk3	clk4	clk5	clk6	clk7
instruction 0	IF&ID	EX
instruction 1		IF&ID	EX
instruction 2			IF&ID	EX
instruction 3				IF&ID	EX

    也認為Simple TPU內部有四級流水線，這是因為在實際執行過程中，包括了讀取暫存器，執行和寫回三個部分，這三個部分是流水設計的。

2.2 超長指令字（VLIW）

    如前文所述，Simple TPU中有兩個基本的計算單元——矩陣乘法陣列和池化計算單元。除此之外，還有一些沒有顯式描述的執行單元，譬如載入和儲存。在這一前提下，即使TPU的指令流水線做得再好，每條指令佔有的週期數也不可能小於1。如果其他執行單元的執行週期數很小，此時總會有一些執行單元處於閒置狀態，處理器的瓶頸會出現在指令上。為了解決這一問題，很直接的想法時每個週期發射多條指令（另一個方法時讓執行單元的執行時間變長，Simple TPU通過向量體系結構設計也有這一處理）。

    由於TPU的專用性，以及計算過程中不存在跳轉和控制的原因，採用VLIW設計多發射處理器似乎是一個很適合的方式。在這一設計下，指令發射結構時固定的，而且所有的冒險可以由編譯器事先檢測並處理，這很大程度可以降低硬體實現的複雜度。在Simple TPU中借鑑了VLIW的思想進行設計，如下所示(示意圖)

    其中各個欄位具體描述如下

• model mask 指定了當前指令執行的模組
• load weight 指定了從記憶體將weight讀取到SRAM的指令
• load act. & mac & store result 指定了將運算元（act.）讀取到暫存器，乘加陣列計算以及將結果寫回到儲存器的過程
• set weight 指定了將運算元（weight）讀取到計算陣列暫存器的過程
• load act. & pooling& store result field指定了將運算元（act.）讀取到暫存器，完成pooling和歸一化計算以及將結果寫回到儲存器的過程

    VLIW的設計放棄了很多的靈活性和相容性，同時將很多工作放到軟體完成，但依舊適合在TPU這樣的偏專用的處理器中使用。Simple TPU中沒有對資料衝突、依賴進行任何處理，軟體需要事先完成分析並進行規避。在這一設計下一條指令可以排程最多四個模組同時工作，效率得到了提升。

3. 卷積計算中的資料並行

3.1 單指令多資料（SIMD）

    單指令多資料，故名思意是指在一條指令控制多組資料的計算。顯然，TPU core的設計中採用了這樣一種資料並行的方式——一條instruction控制了256*256個乘加計算單元（MatirxMultiply/Convolve）。根據指令流和資料流之間的對應關係，可以將處理器分為以下幾個類別

• SISD，單指令流單資料流，順序執行指令，處理資料，可以應用指令並行方法
• SIMD，單指令流多資料流，同一指令啟動多組資料運算，可以用於開發資料級並行
• MISD，多指令流單資料流，暫無商業實現
• MIMD，多指令流多資料流，每個處理器用各種的指令對各自的資料進行操作，可以用在任務級並行上，也可用於資料級並行，比SIMD更靈活

    由於TPU應用在規則的矩陣/卷積計算中，在單個處理器內部的設計上，SIMD是資料並行的最優選擇。SIMD有多種實現方式，根據給出的描述（MatirxMultiply/Convolve指令接受B*256輸入，輸出B*256個結果），TPU中應該採用了類似向量體系結構的設計方法。

3.2 向量體系結構

    如基本單元-矩陣乘法陣列所述，計算單元完成矩陣乘法計算，即向量計算。以《計算機體系結構 : 量化研究方法》給出的例子為例，如需計算

for(int i=0;i<N;i++)
    y[i] += a*x[i];

    以MIPS為例，對於一般的標量處理器和向量處理器而言，執行的指令序列如下所示

    最大的不同在於向量處理器大幅的減小了指令的數目，縮減了指令頻寬。同時，簡單的MIPS指令中可能存在互鎖的情況，會降低效能，而這一現象在向量處理器中則不存在。

    對於卷積神經網路中的卷積操作而言，計算可以表示為（已忽略bias）

for(int i=0;i<M;i++){
    for(int j=0;j<N;j++){
        for(int k=0;k<K;k++){
            for(int c=0;c<C;c++){
                for(int kw=0;kw<KW;kw++){
                    for(int kh=0;kh<KH;kh++){
                        result(i,j,k) += feature(i+kw,j+kh,c)*w(k,kw,kh,c);
                    }
                }
            }
        }
    }
}

    由於KW和KH可能為1（即卷積核的寬度和高度），而weight在計算過程中認為是固定在計算陣列內部的，因此調整迴圈順序後有

for(int kw=0;kw<KW;kw++){
    for(int kh=0;kh<KH;kh++){
        for(int k=0;k<K;k++){
            for(int i=0;i<M;i++){
                for(int j=0;j<N;j++){
                    for(int c=0;c<C;c++){
                        result(i,j,k) += feature(i+kw,j+kh,c)*w(k,kw,kh,c);
                    }
                }
            }
        }
    }
}

    其中第一二層迴圈通過指令進行控制，第三層迴圈在計算陣列中以256並行度進行計算，指令排程；第4-6層迴圈按向量處理器的設計思路進行設計，通過一條指令完成三層迴圈的計算。為了完成迴圈的計算，需要設定三個向量長度暫存器，另外，由於向量在SRAM中的地址並不連續，還需要設定三個不同的步幅暫存器。參考 基本單元-矩陣乘法陣列的程式碼，具體為

    short ubuf_raddr_step1;
    short ubuf_raddr_step2;
    short ubuf_raddr_step3;
    short ubuf_raddr_end1;
    short ubuf_raddr_end2;
    short ubuf_raddr_end3

    採用這樣的設計，SimpleTPU中一條指令可以完成大量資料的計算，提高了資料並行度。這些資料會並行的進入到計算陣列中完成計算（可以認為是多條車道）。由於SimpleTPU中資料的讀取延時是固定的（指從SRAM），因此向量化的設計較一般處理器還更為簡單。

    根據谷歌論文中的描述，TPU中有repeat fileld，但MatrixMultiply/Convolve指令長度有限，因此可能只有一組或兩組向量長度暫存器和步幅暫存器，但設計思路應該類似。

4. 其他

    從谷歌論文中的引數來看，TPU具有極高單位功耗下效能。這一部分來自於其核心設計，正如之前的文章中所描述的

• 採用了INT8資料型別進行計算
• 採用了脈動陣列優化計算
• 沒有采用快取，沒有分支跳轉，預測和資料衝突處理（編譯器完成）

    而從本文的內容可以看出，TPU還採用了簡單的指令集設計+SIMD+向量體系結構+VLIW來進一步優化單位功耗下效能；除此之外，在V2/V3中google更進一步，還利用多核和多處理器設計進一步提高了效能。

 

參考

Jouppi, Norman P. , et al. "In-Datacenter Performance Analysis of a Tensor Processing Unit." the 44th Annual International Symposium IEEE Computer Society, 2017.

JohnL.Hennessy, and DavidA.Patterson. Computer architecture : a quantitative approach = 計算機體系結構 : 量化研究方法/ 5th ed