TPU中的脈動陣列及其實現

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

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

    谷歌TPU概述和簡化

    TPU中的脈動陣列及其實現

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

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

    Simple TPU的設計和性能評估

    SimpleTPU實例：圖像分類

    拓展

    TPU的邊界（規劃中）

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

 

    本文將對TPU中的矩陣計算單元進行分析，並給出了SimpleTPU中32×32的脈動陣列的實現方式和採用該陣列進行卷積計算的方法，以及一個卷積的設計實例，驗證了其正確性。代碼地址https://github.com/cea-wind/SimpleTPU/tree/master/lab1

1. 脈動陣列和矩陣計算

    脈動陣列是一種複用輸入數據的設計，對於TPU中的二維脈動陣列，不少文章中構造了脈動陣列的寄存器模型，致使閱讀較爲困難，而實際上TPU中的二維脈動陣列設計思路十分直接。譬如當使用4×4的脈動陣列計算4×4的矩陣乘法時，有

    如上圖所示，右側是一個乘加單元的內部結構，其內部有一個寄存器，在TPU內對應存儲Weight，此處存儲矩陣B。左圖是一個4×4的乘加陣列，假設矩陣B已經被加載到乘加陣列內部；顯然，乘加陣列中每一列計算四個數的乘法並將其加在一塊兒，即獲得矩陣乘法的一個輸出結果。依次輸入矩陣A的四行，能夠獲得矩陣乘法的結果。

    因爲硬件上的限制，須要對傳播路徑上添加寄存器，而添加寄存器相對於在第i個時刻處理的內容變成了i+1時刻處理；這一過程能夠進行計算結果上的等效。以下圖所示，採用z-1表明添加一個延時爲1的寄存器，若是在縱向的psum傳遞路徑上添加寄存器，爲了保證結果正確，須要在橫向的輸入端也添加一個寄存器（即本來在i進行乘加計算的兩個數均在i+1時刻進行計算）。給縱向每一個psum路徑添加寄存器後，輸入端處理如右圖所示。（下圖僅考慮第一列的處理）

    當在橫向的數據路徑上添加寄存器時，只要每一列都添加相同延時，那麼計算結果會是正確的，可是結果會在後一個週期輸出，以下圖所示

    上述分析能夠，一個4×4的乘加陣列能夠計算一組4×4的乘加陣列完成計算，而對於其餘維度的乘法，則能夠經過屢次調用的方式完成計算。譬如（4×4）×（4×8），能夠將（4×8）的乘法拆分乘兩個4×4的矩陣乘；而對於（4×8）×（8×4），兩個矩陣計算完成後還須要將其結果累加起來，這也是爲什麼TPU在乘加陣列後須要添加Accumulators的緣由。最終脈動陣列設計以下所示（以4×4爲例）

2. 脈動陣列的實現

    如第一節所述，可經過HLS構建一個脈動陣列並進行仿真。相似TPU中的設計，採用INT8做爲計算陣列的輸入數據類型，爲防止計算過程當中的溢出，中間累加結果採用INT32存儲。因爲INT32的表示範圍遠高於INT8，認爲計算過程當中不存在上溢的可能性，所以沒有對溢出進行處理。脈動陣列的計算結果數據類型爲INT32，會在後文進行下一步處理。

    脈動陣列實現的關鍵代碼包括

1. Feature向右側移動

for(int j=0;j<MXU_ROWNUM;j++){
    for(int k=MXU_ROWNUM+MXU_COLNUM-2;k>=0;k--){
        if(k>0)
            featreg[j][k] = featreg[j][k-1];
        else
            if(i<mxuparam.ubuf_raddr_num)
                featreg[j][k] = ubuf[ubuf_raddr][j];
            else
                featreg[j][k] = 0;
    }
}

2. 乘法計算以及向下方移動的psum

for(int j=MXU_ROWNUM-1;j>=0;j--){
    for(int k=0;k<MXU_COLNUM;k++){
        ap_int<32> biasreg;
        biasreg(31,24)=weightreg[MXU_ROWNUM+0][k];
        biasreg(23,16)=weightreg[MXU_ROWNUM+1][k];
        biasreg(15, 8)=weightreg[MXU_ROWNUM+2][k];
        biasreg( 7, 0)=weightreg[MXU_ROWNUM+3][k];
        if(j==0)
            psumreg[j][k] = featreg[j][k+j]*weightreg[j][k] + biasreg;
        else
            psumreg[j][k] = featreg[j][k+j]*weightreg[j][k] + psumreg[j-1][k];
    }
}

    完成代碼編寫後可進行行爲級仿真，能夠看出整個計算陣列的時延關係

1. 對於同一列而言，下一行的輸入比上一行晚一個週期

2. 對於同一行而言，下一列的輸入比上一列晚一個週期（注意同一行輸入數據是同樣的）

3. 下一列的輸出結果比上一列晚一個週期

 

3. 從矩陣乘法到三維卷積

    卷積神經網絡計算過程當中，利用kh×kw×C的卷積核和H×W×C的featuremap進行乘加計算。以3×3卷積爲例，以下圖所示，省略Channel方向，拆分kh和kw方向分別和featuremap進行卷積，能夠獲得9個輸出結果，這9個輸出結果按照必定規律加在一塊兒，就能夠獲得追後的卷積計算結果。下圖給出了3×3卷積，padding=2時的計算示意圖。按F1-F9給9個矩陣乘法結果編號，輸出featuremap中點（2，1）——指第二行第一個點——是F1（1，1），F2（1，2），F3（1，3），F4（2，1），F5（2，2），F6（2，3），F7（3，1），F8（3，2），F9（3，3）的和。

 

    下面的MATLAB代碼闡明瞭這種計算三維卷積的方式，9個結果錯位相加的MATLAB代碼以下所示

output = out1;
output(2:end,2:end,:) = output(2:end,2:end,:) + out2(1:end-1,1:end-1,:);
output(2:end,:,:) = output(2:end,:,:) + out3(1:end-1,:,:);
output(2:end,1:end-1,:) = output(2:end,1:end-1,:) + out4(1:end-1,2:end,:);
output(:,2:end,:) = output(:,2:end,:) + out5(:,1:end-1,:);
output(:,1:end-1,:) = output(:,1:end-1,:) + out6(:,2:end,:);
output(1:end-1,2:end,:) = output(1:end-1,2:end,:) + out7(2:end,1:end-1,:);
output(1:end-1,:,:) = output(1:end-1,:,:) + out8(2:end,:,:);
output(1:end-1,1:end-1,:) = output(1:end-1,1:end-1,:) + out9(2:end,2:end,:);

    而在實際的HLS代碼以及硬件實現上，部分未使用的值並未計算，所以實際計算的index和上述示意圖並不相同，具體可參考testbench中的配置方法。

4. 其餘

    GPU的volta架構中引入了Tensor Core來計算4×4的矩陣乘法，因爲4×4的陣列規模較小，其內部可能並無寄存器，設計可能相似第一節圖1所示。因爲其平均一個週期就能完成4×4矩陣計算，猜想採用第一節中陣列進行堆疊，以下圖所示。

    一些FPGA加速庫中利用脈動陣列實現了矩陣乘法，不過不一樣與TPU中將一個輸入固定在MAC內部，還能夠選擇將psum固定在MAC內部，而兩個輸入都是時刻在變化的。這幾種方式是相似的，就再也不展開描述了。