Loading

学习论文:In-Datacenter Performance Analysis of a Tensor Processing Unit​TM (TPU)

TPU--Tensor Processing Unit​TM-张量处理单元

 

TPU的核心

是一个65,536的8位MAC矩阵相乘单元,

提供了92 TeraOps/秒(top)的峰值吞吐量和

一个大的(28 MiB)软件管理的片上存储器。

 

量子化

将浮点数转换为狭窄的整数(通常只有8位)(这对于推断来说已经足够了)

8位整数乘 比ieee754 16位浮点乘能量小6倍,面积小6倍,

整数加法的优点是能量小13X,面积小38X [Dal16]。

 

为了简化硬件设计和调试,主机服务器发送TPU指令让它执行,而不是自己获取指令。

因此,TPU在本质上更接近于FPU(浮点单元)协处理器,而不是GPU。

 

TPU框图

主要计算部分是黄色矩阵乘单元。它的输入是蓝色权重FIFO和蓝色统一缓冲区(UB)。

它的输出是蓝色累加器,黄色激活单元(Activation)执行非线性函数,并发送到UB上。

 

TPU指令通过PCIe Gen3 x16总线从主机发送到指令缓冲区。内部块通常由256字节宽

的路径连接在一起。从右上角开始,矩阵乘法单位是TPU的核心。它包含256x256个mac,

可以对有符号或无符号整数执行8位乘加运算。16位积在矩阵单元下面32位累加器的4个

MiB中收集。4MiB代表4096,256个32位的累加器。矩阵单元在每个时钟周期产生一个256

个元素的部分和。我们选择了4096,因为我们首先注意到每个字节的操作需要达到峰值

性能(第4节中的roofline knee),即~1350,所以我们将其四舍五入到2048,然后复制它,

以便编译器可以在运行峰值性能时使用双缓冲。

 

TPU模具平面布置图

 

指令

由于指令通过相对缓慢的PCIe总线发送,TPU指令遵循CISC传统,包括重复字段。

这些CISC指令的每条指令的平均时钟周期(CPI)通常是10到20。它总共有十几条指令,

但这五条是最关键的:

1. Read_Host_Memory从CPU主机内存中读取数据到UB中。

2. Read_Weights从权重存储器读取权重到权重FIFO作为矩阵单位的输入。

3.MatrixMultiply/Convolve (矩阵乘/卷积) 使矩阵单元执行矩阵乘或从统一缓冲区到累加

器的卷积。一个矩阵运算需要一个可变大小的B*256输入,乘以一个256x256等权值输

入,并产生一个B*256输出,需要B流水线周期来完成。

4. Activate执行人工神经元的非线性功能,有ReLU、Sigmoid等选项。它的输入是累加

器,输出是统一的缓冲区。它还可以使用die上的专用硬件执行卷积所需的池操作,因为

它连接到非线性函数逻辑。

5. Write_Host_Memory将统一缓冲区的数据写入CPU主机内存。

其他指令包括备用主机内存读写、设置配置、两个版本的同步、中断主机、调试标记、nop和halt。

(PS: CISC矩阵乘法指令为12字节,其中3个为统一的缓冲地址;2为累加器地址;

4是长度(有时是二维的卷积);剩下的是操作码和标志。)   

 

矩阵乘法单位的收缩压数据流

 

由于读取一个大的SRAM比算术使用更多的功率,矩阵单元使用收缩执行,通过减少对统一缓

冲区的读写来节省能量。图4显示数据从左侧流入,权重从顶部加载。给定的256个元素的乘积

运算以对角波前的形式通过矩阵。这些重量是预先加载的,并随着新区块的第一批数据的推进

波而起作用。控制和数据被流水线化,从而产生一种假象,即256个输入被一次读取,并且它

们会立即更新256个累加器中的每个位置。从正确性的角度来看,软件不知道矩阵单元的收缩

特性,但就性能而言,它确实担心单元的延迟。

posted @ 2021-01-21 18:19  kyshan  阅读(642)  评论(0编辑  收藏  举报