動手寫一個簡單版的谷歌TPU-指令集

暗海風發表於2019-07-03

系列目錄

    谷歌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的思想進行設計,如下所示(示意圖)

clip_image002

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

  • 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為例,對於一般的標量處理器和向量處理器而言,執行的指令序列如下所示

image

    最大的不同在於向量處理器大幅的減小了指令的數目,縮減了指令頻寬。同時,簡單的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.

相關文章