深度學習飛速發展過程中，人們發現原有的處理器無法滿足神經網絡這種特定的大量計算，大量的開始針對這一應用進行專用芯片的設計。谷歌的張量處理單元（Tensor Processing Unit，后文簡稱TPU）是完成較早，具有代表性的一類設計，基於脈動陣列設計的矩陣計算加速單元，可以很好的加速神經網絡的計算。本系列文章將利用公開的TPU V1相關資料，對其進行一定的簡化、推測和修改，來實際編寫一個簡單版本的谷歌TPU，以更確切的了解TPU的優勢和局限性。
動手寫一個簡單版的谷歌TPU系列目錄
拓展
TPU的邊界（規划中）
重新審視深度神經網絡中的並行（規划中）
TPU V1定義了一套自己的指令集，雖然在介紹處理器時，往往會先談指令集架構，但此處卻把它放到了最后，這主要基於兩個原因；其一在於個人的對處理器不太了解，這也是主要原因，其二在於公開資料中並沒有TPU指令集的細節和TPU微架構的描述。從數據流和計算單元出發對TPU進行分析固然容易很多，但如果想理解TPU的設計思想，依舊需要回到其架構設計上進行分析。這一部分內容有些超出了我現有的能力，不當之處還請多多指正。
本文主要探討從架構設計上看，TPU時如何做高性能和高效能的設計。高性能的多來自於並行，因此本文分別討論了指令並行和數據並行的設計方法。由於論文中並未描述TPU指令集的具體設計，除特別說明外，本文關於TPU指令集的探討均為推測；另外，SimpleTPU的指令設計並不系統/完整，此處僅闡明設計中的幾種基本思想。
1. TPU的指令集
TPU的指令集采用CISC設計，共計有十多條指令，主要的五條指令包括
- Read_Host_Memory 將數據從CPU的內存中讀取到TPU的Unified Buffer上
- Read_Weights 將weight從內存中讀取到TPU的 Weight FIFO 上.
- MatrixMultiply/Convolve 執行卷積或矩陣乘法操作.
- Activate 執行人工神經網絡中的非線性操作和Pooling操作（如有）
- 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.