TPU1.0理解

根据论文“数据中心的 TPU 性能分析”初略了解一下TPU实现

In-Datacenter Performance Analysis of a Tensor Processing Unit

结构

TPU框图

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占了几乎的晶圆面积。

Floor Plan

控制指令

TPU指令设计遵循的是CISC复杂指令集，每条指令的平均时钟周期CPI是10到20，指令有很多，下面选5个最重要的讲一下：

1. Read_Host_Memory
   从CPU memory读取数据到UB（Unified Buffer）
2. Read_Weights
   从Weight Memory（片外DRAM）读取权重到Weight FIFO作为矩阵运算单元的输入
3. MatrixMultiply/Convolve
   让Matrix Unit执行矩阵乘法或者卷积计算，从Unified Buffer到Accumulator。一次矩阵操作，输入B256大小的输入变量，于256x256的权重常数相乘，得到B256的输出，花费B个流水周期来完成计算。
4. Activate
   执行非线性函数，如ReLU、Sigmoid等。输入来自Accumulators，输出到Unified Buffer。
5. 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的读写，进而减少能量消耗。下图表示从左边输入的数据流和从上方加载的权重。

Systolic data flow

给定的256元乘法累加运算通过矩阵作为对角波前移动。此过程中权重是预先加载的，并且与数据同步。通过对控制逻辑和数据流水线操作，使得对于MMU而言，256个输入一次读取的，并且它们立即更新到256个累加器中的一个个具体位置。

具体实现

matrix单元就是一个典型的脉动阵列。weight由上向下流动，activation数据从左向右流动。在最下方有一些累加单元，

cell内部

cell内部结构如上，几组寄存器分别存储weight，activation和来自上方cell的部分和。weight从上向下传播，可以在weight path register中存储，保持不动或者传输给weight register进行运算。weight register可以把数据发到乘法器进行处理，也可以直接传递给下方的cell；同样activation register也可以把数据发到乘法器进行处理，或者直接传递给右侧的cell。乘法器的输出和Sum in register的数值求和并传递给下方的cell。所有运算和传递都由控制寄存器控制（通过指令决定）。