动手写一个简单版的谷歌TPU-矩阵乘法和卷积

谷歌TPU是一个设计良好的矩阵计算加速单元，可以很好的加速神经网络的计算。本系列文章将利用公开的TPU V1相关资料，对其进行一定的简化、推测和修改，来实际编写一个简单版本的谷歌TPU。计划实现到行为仿真为止，仅为更确切的了解TPU的优势和局限性，暂无在FPGA等硬件上进一步实现的计划。

系列目录

    谷歌TPU概述和简化

    基本单元-矩阵乘法阵列

    基本单元-归一化和池化（待发布）

    TPU中的Instruction （待完成）

    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构建一个脉动阵列并进行仿真。关键代码包括，

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内部，而两个输入都是时刻在变化的。这几种方式是类似的，就不再展开描述了。