TPU中的指令並行和數據並行

深度學習飛速發展過程當中，人們發現原有的處理器沒法知足神經網絡這種特定的大量計算，大量的開始針對這一應用進行專用芯片的設計。谷歌的張量處理單元（Tensor Processing Unit，後文簡稱TPU）是完成較早，具備表明性的一類設計，基於脈動陣列設計的矩陣計算加速單元，能夠很好的加速神經網絡的計算。本系列文章將利用公開的TPU V1相關資料，對其進行必定的簡化、推測和修改，來實際編寫一個簡單版本的谷歌TPU，以更確切的瞭解TPU的優點和侷限性。

動手寫一個簡單版的谷歌TPU系列目錄

    谷歌TPU概述和簡化

    TPU中的脈動陣列及其實現

    神經網絡中的歸一化和池化的硬件實現

    TPU中的指令並行和數據並行

    Simple 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.