谷歌TPU分析之脉动阵列

在之前的博客中，介绍了一种使用脉动阵列计算矩阵乘法的方法，在那篇博客中，脉动阵列的主要特点是：数据从左向右流动，而权重则从上向下流动。而在谷歌第一代的TPU中，其脉动阵列却并非是这种形式的。
在谷歌的TPU中，权重是预先存储在$P_x\ast P_y$个PE上的，并且整个计算过程权重都保持不动，即weight stationary，而数据自左向右流动，同时，每个PE单元的部分和则自上而下流动，下面是一个具体的计算过程演示：

c++模拟代码：

#include<iostream>
#include<cstdlib>
#include<ctime>
#include<iomanip>
#define M 4
#define N 4
using namespace std;
typedef int dtype;
struct PE{
   dtype weight;
   dtype feature;
   dtype psum;
};

class Systolic{
public:
    PE pe_arr[M][N];
public:
    void Init(){                  //初始化，psum置0
        for(int i=0;i<M;i++)
            for(int j=0;j<N;j++){
                pe_arr[i][j].psum=0;
                pe_arr[i][j].feature=0;
            }
    }
    void Read_W(dtype W[M][N]){   //读取权重
        for(int i=0;i<M;i++)
            for(int j=0;j<N;j++)
               pe_arr[i][j].weight=W[i][j];
    }
    void Calc(){                  //计算
        for(int i=0;i<M;i++)
            for(int j=0;j<N;j++)
                pe_arr[i][j].psum+=pe_arr[i][j].weight*pe_arr[i][j].feature;
    }
    void Read_F(dtype In[M]){     //feature向右流动
        for(int i=0;i<M;i++)
            for(int j=N-1;j>0;j--)
                pe_arr[i][j].feature=pe_arr[i][j-1].feature;
        for(int i=0;i<M;i++)
            pe_arr[i][0].feature=In[i];
    }
    void Shift_psum(){            //psum向下流动
        for(int j=0;j<N;j++)
            for(int i=M-1;i>0;i--)
                pe_arr[i][j].psum=pe_arr[i-1][j].psum;
        for(int j=0;j<N;j++)
            pe_arr[0][j].psum=(dtype)0;
    }
    void Print(){
        cout<<"feature:"<<endl;
        for(int i=0;i<N;i++){
            for(int j=0;j<N;j++)
                cout<<fixed<<setw(4)<<pe_arr[i][j].feature<<",";
            cout<<endl;
        }

        cout<<"psum:"<<endl;
        for(int i=0;i<N;i++){
            for(int j=0;j<N;j++)
                cout<<fixed<<setw(4)<<pe_arr[i][j].psum<<",";
            cout<<endl;
        }

    }
};

void Matrix_Multiply(dtype A[N][N],dtype B[N][N]){             //C=B*A
    Systolic S;
    int clock=0;
    S.Init();
    dtype B_copy[N][N];
    //B需要转置后才能存入S中
    for(int i=0;i<N;i++)
        for(int j=0;j<N;j++)
            B_copy[i][j]=B[j][i];
    //打印B_copy
    cout<<"B"<<endl;
    for(int i=0;i<N;i++){
        for(int j=0;j<N;j++)
            cout<<B_copy[i][j]<<",";
        cout<<endl;
    }
    S.Read_W(B_copy);
    //开始计算
    while(clock<3*N-2){
        dtype In[N];
        for(int i=0;i<N;i++)
            if(clock>=i&&clock<N+i)
                In[i]=A[i][clock-i];
            else
                In[i]=(dtype)0;

        S.Read_F(In);
        S.Calc();
        cout<<"clock="<<clock<<endl;

        S.Print();
        S.Shift_psum();

        clock++;
    }
}
int main(){
    srand(1);
    int A[N][N];
    int B[N][N];
    for(int i=0;i<N;i++)
        for(int j=0;j<N;j++){
            A[i][j]=rand()%10;
            B[i][j]=rand()%10;
    }
    cout<<"A"<<endl;
    for(int i=0;i<N;i++){
        for(int j=0;j<N;j++)
            cout<<A[i][j]<<",";
        cout<<endl;
    }

    Matrix_Multiply(A,B);
    return 0;
}