矩陣乘法加速器的設計框架

以往我分析了一些AI加速器的設計，包括TPU，FSD，華爲達芬奇等，無一例外都是從已經給出的設計出發，去分析其優缺點和應用範圍。在以前的文章中，關於這些設計是如何完成的，其背後是否有必定設計原則和理念的內容均沒有進行探討。而這兩點，實則是設計一個優秀的，可持續迭代的加速器的基礎。本文將從矩陣加速器出發，經過一些簡化的模型，給出簡單的設計框架。

1. 矩陣乘法和硬件模型

通常來講，矩陣乘法加速器中須要加速的計算可表示爲

\[ C = A\times B + C \]

其中\(A\in R^{m\times k}\), \(B\in R^{k\times n}\), \(C\in R^{m\times n}\) 。

矩陣乘法加速器，通常至少包括計算單元，緩存（SRAM等構成）和內存（譬如DDR等）。其中緩存的讀寫速率較高，能夠和計算單元的運算速度相匹配，但容量較小；內存的容量相對緩存較大，但讀寫速率較低。

2. 帶寬優化的矩陣乘法加速器設計

和通常的處理器相比，特定的加速器能夠設計數量巨大的計算單元（譬如Google TPU V1設計了65536個乘法器）；可是DDR的帶寬的提高倒是有限的。所以，設計目標之一在於優化數據訪問，下降DDR的讀寫帶寬。

假設加速器的總緩存大小爲\(M\), 在一次計算過程當中，用於存儲矩陣\(A,B,C\)的緩存空間大小分別爲\(M_A,M_B,M_C\)。

矩陣乘法加速器的設計目的通常是爲了加速大規模的矩陣乘法計算，爲了簡化分析過程，假設矩陣\(A,B,C\)的大小\(S_A,S_B,S_C\)均遠大於\(M\)，即計算過程當中每次只能在緩存中存放一部分數據，完成子矩陣\(A_{sub},B_{sub},C_{sub}\)的計算。顯然，存放在緩存中的數據都會參與運算，否在有冗餘數據浪費存儲和帶寬。所以\(A_{sub},B_{sub},C_{sub}\)應可以完成一組矩陣計算，即

\[A_{sub}\in R^{p\times s},B_{sub}\in R^{s\times q},C_{sub}\in R^{p\times q}\]

據此，爲了完成矩陣計算，從DDR到SRAM的總數據讀寫爲

\[D_{size} = n/q \times S_A + m/p \times S_B + 2\times S_C\]

據此能夠給出優化目標爲

\[ \mathbf{min} : mnk/q + mnk/p +2mn \\ \mathbf{sub.to }: p\times s + s\times q + p\times q \leqslant M\\ p>0,s>0,q>0 \]

簡化爲

\[ \mathbf{min} : 1/q + 1/p \\ \mathbf{sub.to }: p\times s + s\times q + p\times q \leqslant M\\ p>0,s>0,q>0 \]

求解得當\(s=1\)，\(p=q=\sqrt{M+1}-1\)時獲得最優解。即若要設計一個帶寬優化的乘法器，應該儘量的將緩存用於存儲\(C_{sub}\)，每次計算的子矩陣爲

\[C_{sub}^{p\times q} += A_{sub}^{p\times 1} + B_{sub}^{1\times q} \]

Telsa的FSD的設計和上述討論結果是一致的（只不過FSD的SRAM對應了上述的DDR，Register對應了上述的SRAM），FSD計算過程當中\(A_{sub}\in R^{96\times 1},B_{sub}\in R^{96\times 96},C_{sub}\in R^{96\times 96}\)。對應的FSD的設計其實是以下降SRAM-Register之間的讀寫爲目的進行優化的。

3. 計算優化的矩陣乘法加速器設計

依據第二節的結果，每次計算的子矩陣爲

\[C_{sub}^{p\times q} += A_{sub}^{p\times 1} + B_{sub}^{1\times q} \]

整個計算過程當中，其並行度最高爲\({p\times q}\)（即每一個週期完成\({p\times q}\)個乘法）。而爲了完成一次計算，須要從緩存裏讀取\(p+q+q\times q\)個數據送入到計算陣列中。所以一次讀/寫的數據位寬寬度極高，隨着並行度的增加，數據位寬線性增加。

數據位寬的問題主要存在\(C_{sub}\)上。爲了解決這一問題，Telsa FSD採用了移位的方式，在計算完成後，將計算結果依次寫回到SRAM中。

若是設計目的在於計算陣列和緩存之間的優化，參考第二節的設計思路，在必定並行度上，但願儘量下降緩存的讀寫帶寬，優化目標能夠表示爲

\[ \mathbf{min}:x\times y+y\times z+x\times z \\ \mathbf{sub.to }:x\times y\times z=P \\ x>0,y>0,z>0 \]

其中\(P\)表明計算陣列的並行度，求解得當\(x=y=z=\sqrt[3]{P}\)時，此時設計的計算陣列對緩存的訪問能夠儘量的低。

華爲的達芬奇架構中計算陣列的設計和上述討論是一致的，達芬奇中的CUBE Core是一個\(16\times16\times16\)的MAC陣列（以Davinci Max爲例），能夠完成
\[C_{sub}^{16\times 16} += A_{sub}^{16\times 16} + B_{sub}^{16\times 16} \]
的矩陣計算。

4. 總結

上述的全部討論都基於一個最簡單的硬件模型，從兩個角度分別求解了理論上最優的設計應該是怎麼樣的。

實際狀況每每會複雜不少，硬件架構方面就會複雜不少。同時優化的目標每每有多個，而優化的限制條件也會有不少。

可是在我看來，只有採用這樣的設計方法，即將問題建模，求解，才能造就一個好的設計。也只有採用這樣的設計方法，才能再已有的基礎上，進一步增長優化目標和優化條件，進一步的優化架構設計。