TPU只完成推理过程,训练过程在GPU上完成。TPU可以像GPU一样通过PCIe总线接口挂载到现有的服务器上。设计目标是为了在TPU上完成所有的推理模型,从而减少和主机CPU的交互,进而满足2015年及今后的神经网络需求。下图是TPU的整体结构框图。
主机通过PCIe Gen3x16的总线发送TPU的指令到其中的指令buffer内,内部模块之间通过典型的256位宽通路连接。右上角的矩阵乘法单元是TPU的核心,它由256x256个MACs组成,每个MACs可以对有符号和无符号整形的数据完成8bit的乘加运算。在矩阵单元的下面有一个4MiB大小32bit位宽的Accumulator,收集上面16bit的计算结果。4MiB代表4096个, 256-element, 32-bit累加器。矩阵单元每个时钟周期产生一个256-element的局部和(partial sum)。当使用8 bit权重(weight)和16 bit激活(activation)(或反之亦然)的混合时,矩阵乘法单元以半速计算,并且当它们都是16bit时以四分之一速度计算。它每个时钟周期读取和写入256个值,可以执行矩阵乘法或卷积。矩阵单元保存一个64KiB的权重块( tile of weights),并使用了双缓冲(这样可以在处理时同时读入权重)。该单元设计用于密集矩阵,而没有考虑稀疏的架构支持(部署时间的原因)。稀疏性将在未来的设计中占有优先地位。
矩阵乘法单元的权重通过片上的权重FIFO(Weight FIFO)进行缓存,该FIFO从片外8 GiB DRAM读取。由于是用于推论,权重为只读。8 GiB可以同时支持许多模型。权重FIFO的深度是4个ti