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

以往我分析了一些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. 總結

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

實際情況往往會復雜很多，硬件架構方面就會復雜很多。同時優化的目標往往有多個，而優化的限制條件也會有很多。

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