在之前的博客中，介绍了一种使用脉动阵列计算矩阵乘法的方法，在那篇博客中，脉动阵列的主要特点是：数据从左向右流动，而权重则从上向下流动。而在谷歌第一代的TPU中，其脉动阵列却并非是这种形式的。
在谷歌的TPU中，权重是预先存储在 P
x
∗
P
y
P_x*P_y Px∗Py个PE上的，并且整个计算过程权重都保持不动，即weight stationary，而数据自左向右流动，同时，每个PE单元的部分和则自上而下流动，下面是一个具体的计算过程演示：
c++模拟代码：
#include #include #include #include #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; }
