谷歌TPU是一个设计良好的矩阵计算加速单元,可以很好的加速神经网络的计算。本系列文章将利用公开的TPU V1(后简称TPU)相关资料,对其进行一定的简化、推测和修改,来实际编写一个简单版本的谷歌TPU。计划实现到行为仿真为止,仅为更确切的了解TPU的优势和局限性,暂无在FPGA等硬件上进一步实现的计划。
系列目录
1. TPU设计分析
In-Datacenter Performance Analysis of a Tensor Processing Unit,TPU V1的结构框图如下所示
基本单元-矩阵乘法阵列进行更进一步的阐述。TPU的设计关键在于充分利用这一乘加阵列,使其利用率尽可能高。
- Local Unified Buffer 提供了的带宽(即167GiB/s,0.25Kib×700/1024/1024=167GiB/s),以保证计算单元不会因为缺少Data in而闲置;
- Local Unified Buffer 的空间高达24MiB,这意味着计算过程的中间结果几乎无需和外界进行交互,也就不存在因为数据带宽而限制计算能力的情况;
- Matrix Multiply Unit中每个MAC内置两个寄存器存储Weight,当一个进行计算时另一个进行新Weight的载入,以掩盖载入Weight的时间;
- 30GiB/s的带宽完成256×256Weight的载入需要大约1430个Cycles,也就意味着一组Weight至少需要计算1430Cycles,因此Accumulators的深度需要为2K(1430取2的幂次,论文中给出的数值是1350,差异未知);
- 由于MAC和Activation模块之间需要同时进行计算,因此Accumulators需要用两倍存储来进行pingpang设计,因此Accumulators中存储的深度设计为4k;
- Read_Host_Memory
- Read_Weights
- MatrixMultiply/Convolve
- Activation
- Write_Host_Memory
需要注意的是,一条指令可以执行数千个周期,因此TPU设计过程中没有对流水线之间的空闲周期进行掩盖,这是因为由于Pipline带来的数十个周期的浪费对最终性能的影响不到1%。
2. TPU的简化
Resource | TPU | SimpleTPU |
Matrix Multiply Unit | 256*256 | 32*32 |
Accumulators RAM | 4K*256*32b | 4K*32*32b |
Unified Buffer | 96K*256*8b | 16K*32*8b |
3. 基于Xilinx HLS的实现方案
高层次综合(HLS)-简介,以及一个简单的开发实例利用Xilinx HLS实现LDPC译码器。
HLS易写、易读、易验证,此处依旧选择了HLS,并通过一些手段规避掉了部分问题。真实应用中,采用HDL或者HDL结合HLS进行开发是更为合适的选择。