概述
深度学习模型中费时的算子是往往是卷积操作,而卷积的本质是矩阵的乘加计算,所以通过硬件加速矩阵的乘加运算就能达到加速深度学习训练与推理的目标;本文旨在形象化的解释Google TPU与Nvidia TensorCore对矩阵乘加的加速过程
Nvidia TensorCore
Nvidia官方给出的TensorCore 加速矩阵乘加运算示意图如上,虽然很直观但却不容易理解图中灰色3维的立方体计算单元是如何完成矩阵乘加运算的,本文试图解释如下
如上图所示由于矩阵乘法计算天然存在3维计算的特性(3个for循环),所以设计一个3维的硬件计算单元就可以并行计算所有矩阵元素的乘法再做累加;上图展示了一个4x4x4的三维计算单元如何加速4x4矩阵乘加的过程
那么当矩阵大小超过4x4的时候,上图展示了如何将8x8的矩阵分解为多个4x4的矩阵乘加,从而可以继续利用4x4x4三维计算单元进行乘加硬件加速,所以这也解释了为什么NV要求用户的模型中矩阵的大小最好是4的倍数
Google TPU
Google TPU中通过MAC模块进行矩阵乘法的加速,官方给出的加速矩阵乘法运算示意图如上,本文试图解释如下
如上图所示,两个矩阵的元素按照特定的排列方式逐步流入灰色的MAC计算阵列,每一步就可以在MAC计算单元中得到一个中间结果,当两个矩阵的元素全部流过MAC计算阵列后,MAC计算单元中就是完整的矩阵乘法结果
脉动阵列
以上的矩阵乘加硬件加速在芯片设计上的原理都离不开脉动阵列理论:Systolic Array源于1982年H. T. Kung发表的论文(Why Systolic Architectures?),这篇论文阐述了并行处理芯片采用脉动阵列的三条理由。今天所有的*PU设计都源自这一理念,TPU MAC、NV TensorCore 、含光800 Tensor Array 等……