AI芯片学习小结4-谷歌TPU与脉动阵列

AI芯片学习小结4-谷歌TPU与脉动阵列

文章:In-Datacenter Performance Analysis of a Tensor Processing Unit
时间:2017
Reference:深入理解Google TPU的脉动阵列架构,炼数成金_小数

终于迎来了鼎鼎大名的谷歌TPU。由于谷歌在其论文中没有对脉动阵列进行详细讲述,而有关脉动阵列及其具体实现的论文也五花八门,所以我主要参考了Reference中的大神博客并进行以下学习记录。

脉动阵列

脉动阵列的核心思想:

  1. 进行大量的数据复用,包括activation以及weight。其复用依赖于片上网络的性质,即PE单元可以向周围的PEs传递数据。
  2. 部署大规模的PE阵列。每个PE单元结构简单,一般只包含乘法器、加法器以及寄存器三部分,适合大量堆砌。

AI芯片学习小结4-谷歌TPU与脉动阵列
上图为脉动阵列的简单实现。可以看到,weight已经被预先固定在了PE的寄存器中,中间和不断向右累加便可得到最终的结果。
AI芯片学习小结4-谷歌TPU与脉动阵列
上图展示了如何使用脉动阵列来计算两个3*3矩阵的乘法。在经过一系列这样的过程后,最终得到的形式如下图所示。因为这些数据在阵列中传输的形式类似于波阵面在不断推进的过程,因此也与“脉动阵列”的名字相呼应。
AI芯片学习小结4-谷歌TPU与脉动阵列

谷歌TPU

AI芯片学习小结4-谷歌TPU与脉动阵列
上图是google论文中展示的TPU实现模块图。位于图片右上方的Matrix Multiply Unit为整个实现的核心计算部分。阵列中一共含有惊人的256*256,也就是64k个PE单元;每一个计算单元可以实现两个8bit有符号或无符号数的乘加操作。
AI芯片学习小结4-谷歌TPU与脉动阵列
上图为Matrix的具体结构图。由于TPU是基于脉动阵列架构实现的,因此其基本工作流程也与前面提到的脉动阵列的相同。weight被固定在了PE中,psum从上到下逐步累积,activation从左到右在PE之间进行传递。
由于二维卷积操作与矩阵乘法的工作方式不一样,因此在进行卷积前需要对ifmap和filter进行适当的变换使得对应位置的数值能够正确相乘。在TPU中这一变换操作具体由User Space drive负责实现,下面援引文章中的一段话来印证:

The User Space driver changes frequently. It sets up and controls TPU execution, reformats data into TPU order, translates API calls into TPU instructions, and turns them into an application binary. The

其具体的变换方式在参考博客的最后使用了一张图片来进行描述,但由于其变换方式较为复杂,且与我们所要学习记录的核心内容关系不大,所以我不再详细记录。