根据论文“数据中心的 TPU 性能分析”初略了解一下TPU实现
In-Datacenter Performance Analysis of a Tensor Processing Unit
结构
PCIe3 Gen3 x16
通过PCIe3.0x16与主机通信,内部模块一般用带宽256-byte的路径连接。
Matrix Multiply Unit
Matrix Multiply Unit (矩阵计算单元)是TPU的核心部分,包含256x256个MAC计算单元,进行8-bit的乘加运算。计算结果是16-bit,到累加器部分进行累加。
Weight FIFO
权重由片上缓存Weight FIFO输给矩阵计算单元,是从片外的8 GiB的DRAM里读取的。
Unified Buffer
大小为24 MiB,用于保存中间结果,也是矩阵计算单元的输入。通过一个可编程的DMA来负责在CPU memory和Unified Buffer间传输数据。下图可以看到24 MiB的Unified Buffer占了几乎的晶圆面积。
控制指令
TPU指令设计遵循的是CISC复杂指令集,每条指令的平均时钟周期CPI是10到20,指令有很多,下面选5个最重要的讲一下:
- Read_Host_Memory
从CPU memory读取数据到UB(Unified Buffer) - Read_Weights
从Weight Memory(片外DRAM)读取权重到Weight FIFO作为矩阵运算单元的输入 - MatrixMultiply/Convolve
让Matrix Unit执行矩阵乘法或者卷积计算,从Unified Buffer到Accumulator。一次矩阵操作,输入B256大小的输入变量,于256x256的权重常数相乘,得到B256的输出,花费B个流水周期来完成计算。 - Activate
执行非线性函数,如ReLU、Sigmoid等。输入来自Accumulators,输出到Unified Buffer。 - Write_Host_Memory
把数据写回到CPU Memory。
CISC MatrixMultiply指令为12个字节,其中3个为Unified Buffer地址; 2是累加器地址; 4是长度(有时是卷积的2个维度); 其余的是操作码和标志。
TPU加速的思路是让矩阵计算单元一直被占用,因此为CISC指令设置了4级流水,每个指令在独立的一级执行。 其目标是通过将其执行与MatrixMultiply指令重叠来隐藏其他指令(如Read_Weights等)。但是,当激活的输入或权重数据尚未就绪,矩阵单元将进入等待模式。
Systolic data flow
由于读取大的SRAM的能耗比算术运算高的多,矩阵计算单元通过脉冲执行来减少对 Unified Buffer的读写,进而减少能量消耗。下图表示从左边输入的数据流和从上方加载的权重。
给定的256元乘法累加运算通过矩阵作为对角波前移动。此过程中权重是预先加载的,并且与数据同步。通过对控制逻辑和数据流水线操作,使得对于MMU而言,256个输入一次读取的,并且它们立即更新到256个累加器中的一个个具体位置。
matrix单元就是一个典型的脉动阵列。weight由上向下流动,activation数据从左向右流动。在最下方有一些累加单元,
cell内部结构如上,几组寄存器分别存储weight,activation和来自上方cell的部分和。weight从上向下传播,可以在weight path register中存储,保持不动或者传输给weight register进行运算。weight register可以把数据发到乘法器进行处理,也可以直接传递给下方的cell;同样activation register也可以把数据发到乘法器进行处理,或者直接传递给右侧的cell。乘法器的输出和Sum in register的数值求和并传递给下方的cell。所有运算和传递都由控制寄存器控制(通过指令决定)。