在之前的博客中,介绍了一种使用脉动阵列计算矩阵乘法的方法,在那篇博客中,脉动阵列的主要特点是:数据从左向右流动,而权重则从上向下流动。而在谷歌第一代的TPU中,其脉动阵列却并非是这种形式的。
在谷歌的TPU中,权重是预先存储在
P
x
∗
P
y
P_x*P_y
Px∗Py个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;
}