【芯片原理】NPU矩陣乘法加速詳解

概述

深度學習模型中費時的算子是往往是卷積操作,而卷積的本質是矩陣的乘加計算,所以通過硬件加速矩陣的乘加運算就能達到加速深度學習訓練與推理的目標;本文旨在形象化的解釋Google TPU與Nvidia TensorCore對矩陣乘加的加速過程

 

Nvidia TensorCore

NV TensorCore 加速矩陣乘加運算示意圖

Nvidia官方給出的TensorCore 加速矩陣乘加運算示意圖如上,雖然很直觀但卻不容易理解圖中灰色3維的立方體計算單元是如何完成矩陣乘加運算的,本文試圖解釋如下

4x4矩陣乘法與3維計算單元

如上圖所示由於矩陣乘法計算天然存在3維計算的特性(3個for循環),所以設計一個3維的硬件計算單元就可以並行計算所有矩陣元素的乘法再做累加;上圖展示了一個4x4x4的三維計算單元如何加速4x4矩陣乘加的過程

8x8矩陣乘法轉換爲4x4矩陣乘加

那麼當矩陣大小超過4x4的時候,上圖展示瞭如何將8x8的矩陣分解爲多個4x4的矩陣乘加,從而可以繼續利用4x4x4三維計算單元進行乘加硬件加速,所以這也解釋了爲什麼NV要求用戶的模型中矩陣的大小最好是4的倍數

 

Google TPU

Google TPU MAC陣列加速矩陣乘法

Google TPU中通過MAC模塊進行矩陣乘法的加速,官方給出的加速矩陣乘法運算示意圖如上,本文試圖解釋如下

MAC計算矩陣乘法的示意圖

如上圖所示,兩個矩陣的元素按照特定的排列方式逐步流入灰色的MAC計算陣列,每一步就可以在MAC計算單元中得到一箇中間結果,當兩個矩陣的元素全部流過MAC計算陣列後,MAC計算單元中就是完整的矩陣乘法結果

 

脈動陣列

以上的矩陣乘加硬件加速在芯片設計上的原理都離不開脈動陣列理論:Systolic Array源於1982年H. T. Kung發表的論文(Why Systolic Architectures?),這篇論文闡述了並行處理芯片採用脈動陣列的三條理由。今天所有的*PU設計都源自這一理念,TPU MAC、NV TensorCore 、含光800 Tensor Array 等……

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章