矩陣乘法加速器的設計框架
- 2020 年 3 月 10 日
- 筆記
以往我分析了一些AI加速器的設計，包括TPU，FSD，華為達芬奇等，無一例外都是從已經給出的設計出發，去分析其優缺點和應用範圍。在之前的文章中，關於這些設計是如何完成的，其背後是否有一定設計原則和理念的內容均沒有進行探討。而這兩點，實則是設計一個優秀的，可持續迭代的加速器的基礎。本文將從矩陣加速器出發，通過一些簡化的模型，給出簡單的設計框架。
1. 矩陣乘法和硬體模型
一般來說，矩陣乘法加速器中需要加速的計算可表示為
[ C = Atimes B + C ]
其中(Ain R^{mtimes k}), (Bin R^{ktimes n}), (Cin R^{mtimes 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^{ptimes s},B_{sub}in R^{stimes q},C_{sub}in R^{ptimes q}]
據此，為了完成矩陣計算，從DDR到SRAM的總數據讀寫為
[D_{size} = n/q times S_A + m/p times S_B + 2times S_C]
據此可以給出優化目標為
[ mathbf{min} : mnk/q + mnk/p +2mn \ mathbf{sub.to }: ptimes s + stimes q + ptimes q leqslant M\ p>0,s>0,q>0 ]
簡化為
[ mathbf{min} : 1/q + 1/p \ mathbf{sub.to }: ptimes s + stimes q + ptimes q leqslant M\ p>0,s>0,q>0 ]
求解得當(s=1)，(p=q=sqrt{M+1}-1)時得到最優解。即若要設計一個頻寬優化的乘法器，應該儘可能的將快取用於存儲(C_{sub})，每次計算的子矩陣為
[C_{sub}^{ptimes q} += A_{sub}^{ptimes 1} + B_{sub}^{1times q} ]
Telsa的FSD的設計和上述討論結果是一致的（只不過FSD的SRAM對應了上述的DDR，Register對應了上述的SRAM），FSD計算過程中(A_{sub}in R^{96times 1},B_{sub}in R^{96times 96},C_{sub}in R^{96times 96})。對應的FSD的設計實際上是以降低SRAM-Register之間的讀寫為目的進行優化的。
3. 計算優化的矩陣乘法加速器設計
依據第二節的結果，每次計算的子矩陣為
[C_{sub}^{ptimes q} += A_{sub}^{ptimes 1} + B_{sub}^{1times q} ]
整個計算過程中，其並行度最高為({ptimes q})（即每個周期完成({ptimes q})個乘法）。而為了完成一次計算，需要從快取里讀取(p+q+qtimes q)個數據送入到計算陣列中。因此一次讀/寫的數據位寬寬度極高，隨著並行度的增長，數據位寬線性增長。
數據位寬的問題主要存在(C_{sub})上。為了解決這一問題，Telsa FSD採用了移位的方式，在計算完成後，將計算結果依次寫回到SRAM中。
如果設計目的在於計算陣列和快取之間的優化，參考第二節的設計思路，在一定並行度上，希望儘可能降低快取的讀寫頻寬，優化目標可以表示為
[ mathbf{min}:xtimes y+ytimes z+xtimes z \ mathbf{sub.to }:xtimes ytimes z=P \ x>0,y>0,z>0 ]
其中(P)代表計算陣列的並行度，求解得當(x=y=z=sqrt[3]{P})時，此時設計的計算陣列對快取的訪問可以儘可能的低。
華為的達芬奇架構中計算陣列的設計和上述討論是一致的，達芬奇中的CUBE Core是一個(16times16times16)的MAC陣列（以Davinci Max為例），可以完成
[C_{sub}^{16times 16} += A_{sub}^{16times 16} + B_{sub}^{16times 16} ]
的矩陣計算。
4. 總結
上述的所有討論都基於一個最簡單的硬體模型，從兩個角度分別求解了理論上最優的設計應該是怎麼樣的。
實際情況往往會複雜很多，硬體架構方面就會複雜很多。同時優化的目標往往有多個，而優化的限制條件也會有很多。
但是在我看來，只有採用這樣的設計方法，即將問題建模，求解，才能造就一個好的設計。也只有採用這樣的設計方法，才能再已有的基礎上，進一步增加優化目標和優化條件，進一步的優化架構設計。
