在 CFU Playground 上加速 MLPerf™ Tiny 影像分類模型 #3:分塊矩陣乘法

Yeecy
20 min readJan 8, 2024

--

3.1 矩陣乘法

假設矩陣 A 和矩陣 B 的大小分別為 M×KK×N,令小寫字母表示矩陣元素,則 AB 的乘積 C 按矩陣乘法定義可得 cₘₙ = aₘ₁×b + aₘ₂×b + … + aₘₖ×bₖₙ,其中 1≤mM 且 1≤nN。以下的 GEMM 為通用矩陣乘法 general matrix multiplication 的縮寫。

# Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
# Naive GEMM

for m in range(M):
for n in range(N):
for k in range(K):
C[m][n] += A[m][k] * B[k][n]

上面的程式碼即為矩陣乘法,在 C/C++ 中,陣列儲存的順序為列優先(row major),意即同一列的元素會在記憶體內連續儲存,用 B 舉例,其儲存順序如下。

B[0][0], B[0][1], ..., B[0][K-1], ..., B[M-1][0], B[M-1][1], ..., B[M-1][K-1]

那這對上面的矩陣乘法實作有什麼影響呢?我們先觀察 ABC 的存取模式。

m = 0, n = 0, k = 0 to K-1
-----------------------------------------
A[0][0], A[0][1], A[0][2], ..., A[0][K-1]
B[0][0], B[1][0], B[2][0], ..., B[K-1][0]
C[0][0], C[0][0], C[0][0], ..., C[0][0]

可以看到 A 的存取步長為 1,B 的存取步長為 N(因為 B[k][0]B[k+1][0] 之間存在 N-1 個元素 ),而 C 的存取步長為 0。從 CPU 的角度來說,我們希望記憶體的存取能像 A 一樣,每次存取的位址都跟上一次存取連續,這樣除了對 CPU 的快取記憶體較為友善外,也能較容易使用之前所提過的 SIMD 指令,不過考慮到後續會使用矩陣乘法加速器而非 SIMD,因此接下來不會討論如何用 SIMD 指令加速矩陣乘法。

一般而言 CPU 暫存器存取跟記憶體存取在速度上的差異高達數十倍甚至是上百倍,因此 CPU 中引入了多層級的快取記憶體(cache)來掩蓋兩者之間的速度鴻溝,其特性為上一級快取記憶體的存取速度快於下一級快取記憶體,最後一級快取記憶體的存取速度快於主記憶體,但出於成本考量,上一級快取記憶體的容量會小於下一級的快取記憶體,讀者在下圖中能看到跟 L2 快取記憶體比起來 L1 快取記憶體的面積顯著更小。

Intel® Alder Lake Golden Cove 核心面積,右圖為更加精確的面積估計/來源:@Locuza_

因為上一級快取記憶體的速度較下一級快,我們自然會希望需要的資料都能存放在最上一級的快取記憶體裡,然而快取記憶體的運作原理是當所需資料不存在時,會一次從下一級的快取記憶體裡取一定長度的連續資料回來儲存(這個長度通常為 64 個位元組),這也是為何連續存取記憶體位址會對快取較為友善的原因。然而在原先的矩陣乘法中當 N 太大時,存取 B[k][0] 後因為其與 B[k+1][0] 之間間隔超過 64 個位元組,因此 B[k+1][0] 並不會儲存在最上一級的快取記憶體中,使得存取 B[k+1][0] 依然需要像存取 B[k][0] 一樣額外花時間等待資料從下一級快取記憶體傳回來。

那我們能改善 B 的存取模式嗎?答案是可以,而且只需要將迴圈的順序調整一下就好。

# Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
# Naive GEMM with Loop Interchange

for m in range(M):
for k in range(K):
for n in range(N):
C[m][n] += A[m][k] * B[k][n]

相應的存取模式如下,可以看到此時 A 的存取步長為 0,BC 的存取步長 1,相對於原本的實作來說,這個版本對快取記憶體更加友善。

m = 0, k = 0, n = 0 to N-1
-----------------------------------------
A[0][0], A[0][0], A[0][0], ..., A[0][0]
B[0][0], B[0][1], B[0][2], ..., B[0][N-1]
C[0][0], C[0][1], C[0][2], ..., C[0][N-1]

接著我們來看看不同迴圈順序的矩陣乘法跑在 VexRiscv_FullCfu @ 75MHz 需要的週期數,以下用英文字母的順序表示迴圈從外到內的順序。

M = 512, N = 512, K = 1024
--------------------------
MNK: 15541927148 cycles
MKN: 14750774696 cycles
NMK: 18299931059 cycles
NKM: 22353420557 cycles
KMN: 15541658567 cycles
KNM: 15541659510 cycles

從數據中可以看到 MKN 的速度較一開始的 MNK 快,符合前面的推論,至於其他情況為何較糟,就留待讀者自行分析。另外請注意如果讀者要在自己的電腦上進行實驗的話,根據經驗需要將矩陣大小設定為 2048 會比較容易觀察出差異,其原因為個人電腦 CPU 的快取記憶體容量較大,且因為與快取相關的演算法電路不同,得到結果的趨勢可能跟上面的數據不同,不過我相信 MKN 的速度應該還是最快的。

在結束本小節前,我們快速帶過兩個相關的最佳化技巧。

矩陣轉置

除了調整迴圈順序外,我們其實也能透過將矩陣轉置來做到差不多的效果。

# Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
# Naive GEMM after Transposing Matrix B

for k in range(K):
for n in range(N):
BT[n][k] = B[k][n]

for m in range(M):
for n in range(N):
for k in range(K):
C[m][n] += A[m][k] * BT[n][k]

先轉置再做矩陣乘法的存取模式如下。

m = 0, n = 0, k = 0 to K-1
-----------------------------------------
A[0][0], A[0][1], A[0][2], ..., A[0][K-1]
B[0][0], B[0][1], B[0][2], ..., B[0][K-1]
C[0][0], C[0][0], C[0][0], ..., C[0][0]

雖然如此一來 AB 的存取步長為 1,但我們需要考慮到矩陣轉置需要額外花費時間和空間(不須額外空間的矩陣轉置不太好寫),而且當 B 足夠大時矩陣轉置的寫法也可能會影響到快取記憶體的表現,所以跟前面的 MKN 相比孰優孰劣需要實際在機器上跑過一遍才能確定。

迴圈展開(Loop Unrolling)

矩陣乘法的運算除了乘積累加之外,另一個計算來源是更新迴圈變數,假如能確定最內層迴圈的執行次數的話,我們可以透過將迴圈展開減少運算,首先先用 C/C++ 改寫前面的 MKN 方法以便觀察迴圈變數計算。

// Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
// Naive GEMM with Loop Interchange in C/C++

for (uint32_t m = 0; m < M; m++) {
for (uint32_t k = 0; k < K; k++) {
for (uint32_t n = 0; n < N; n++) {
C[m][n] += A[m][k] * B[k][n];
}
}
}

我們可以觀察到每一次遞增完 n 後都需要檢查其值是否超出邊界(其實 mk 也是),假設已知 N 的是 4 的倍數,那麼我們可以將其展開如下,減少四次 n < N 判斷。

// Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
// Naive GEMM with Loop Interchange and Loop Unrolling

for (uint32_t m = 0; m < M; m++) {
for (uint32_t k = 0; k < K; k++) {
for (uint32_t n = 0; n < N; n += 4) {
C[m][n] += A[m][k] * B[k][n];
C[m][n+1] += A[m][k] * B[k][n+1];
C[m][n+2] += A[m][k] * B[k][n+2];
C[m][n+3] += A[m][k] * B[k][n+3];
}
}
}

如果懶得自己展開的話,也可以使用編譯器的指示詞做到一樣的效果,以下以 GCC 做為範例。

// Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
// Naive GEMM with Loop Interchange and Loop Unrolling

for (uint32_t m = 0; m < M; m++) {
for (uint32_t k = 0; k < K; k++) {
#pragma GCC unroll 4
for (uint32_t n = 0; n < N; n++) {
C[m][n] += A[m][k] * B[k][n];
}
}
}

最後讓我們看一下在 MKN 上應用迴圈展開的效果,以下 _ux 表示將最內層迴圈展開 x 次,可以看到隨著展開次數的增加,所需的周期數相應變少,通常來說可以預期展開越多表現越好,但是實際效果依舊需要測試後才能確定。

M = 512, N = 512, K = 1024
---------------------------
MKN: 14750774696 cycles
---------------------------
MKN_u2: 14069664701 cycles
MKN_u4: 13999195280 cycles
MKN_u8: 13868456401 cycles
MKN_u16: 13770646464 cycles
MKN_u32: 13733894203 cycles
MKN_u64: 13729167233 cycles

3.2 分塊矩陣乘法

在線性代數裡大家應該學過分塊矩陣(block matrix),並且知道兩個可以相乘的矩陣在經過分塊後相乘所得的分塊矩陣等價於原本直接相乘得到的矩陣,如果無法接受或不知道這個知識的話,我們可以考慮最簡單的分塊方法,也就是把矩陣每一個元素當成一個 1×1 的矩陣,那麼兩個矩陣分塊相乘後的結果自然而然跟沒分塊的相乘結果相同。

以下用個簡單的例子演示一遍分塊矩陣乘法,首先定義兩個矩陣。

Matrix A              Matrix B
--------------- ---------------
[a] [b] [c] [d] [q] [r] [s] [t]
[e] [f] [g] [h] [u] [v] [w] [x]
[i] [j] [k] [l] [y] [z] [α] [β]
[m] [n] [o] [p] [γ] [δ] [ϵ] [ζ]

接著將它們切成 2×2 的分塊矩陣,也就是說兩個矩陣裡各別都有四個子矩陣。

Blocked Matrix A      Blocked Matrix B
---------------- ----------------
[a b] [c d] [q r] [s t]
[e f] [g h] [u v] [w x]

[i j] [k l] [y z] [α β]
[m n] [o p] [γ δ] [ϵ ζ]

不妨用符號替代子矩陣如下。

Blocked Matrix A      Blocked Matrix B
---------------- ----------------
[S] [T] [W] [X]
[U] [V] [Y] [Z]

按照矩陣乘法將兩個分塊矩陣相乘。

Block Matrix C
---------------
[SW+TY] [SX+TZ]
[UW+VY] [UX+VZ]

分別計算每個子矩陣之值如下。

SW+TY                            SX+TZ
--------------------------- ---------------------------
[aq+bu+cy+dγ] [ar+bv+cz+dδ] [as+bw+cα+dϵ] [at+bx+cβ+dζ]
[eq+fu+gy+hγ] [er+fv+gz+hδ] [es+fw+gα+hϵ] [et+fx+gβ+hζ]

UW+VY UX+VZ
--------------------------- ---------------------------
[iq+ju+ky+lγ] [ir+jv+kz+lδ] [is+jw+kα+lϵ] [it+jx+kβ+lζ]
[mq+nu+oy+pγ] [mr+nv+oz+pδ] [ms+nw+oα+pϵ] [mt+nx+oβ+pζ]

將算出來的子矩陣代入原分塊矩陣可以得到下面的結果。

Blocked Matrix C
-----------------------------------------------------
[aq+bu+cy+dγ ar+bv+cz+dδ] [as+bw+cα+dϵ at+bx+cβ+dζ]
[eq+fu+gy+hγ er+fv+gz+hδ] [es+fw+gα+hϵ et+fx+gβ+hζ]

[iq+ju+ky+lγ ir+jv+kz+lδ] [is+jw+kα+lϵ it+jx+kβ+lζ]
[mq+nu+oy+pγ mr+nv+oz+pδ] [ms+nw+oα+pϵ mt+nx+oβ+pζ]

最後計算未分塊的矩陣乘法作為對照,兩相比對可以發現矩陣對應位置的數值皆相同。

Matrix C
-------------------------------------------------------
[aq+bu+cy+dγ] [ar+bv+cz+dδ] [as+bw+cα+dϵ] [at+bx+cβ+dζ]
[eq+fu+gy+hγ] [er+fv+gz+hδ] [es+fw+gα+hϵ] [et+fx+gβ+hζ]
[iq+ju+ky+lγ] [ir+jv+kz+lδ] [is+jw+kα+lϵ] [it+jx+kβ+lζ]
[mq+nu+oy+pγ] [mr+nv+oz+pδ] [ms+nw+oα+pϵ] [mt+nx+oβ+pζ]

那麼討論分塊矩陣乘法有什麼意義呢?分塊一詞在數學的語境下原文為 block,但在高性能計算領域領域裡分塊更多地指代一類最佳化技巧 blocking 或 tiling,意即將整體計算切割為各個部分以達成不同的最佳化目的,以 CPU 來說,應用分塊技巧主要目的為了讓計算所需資料盡可能留在快取記憶體裡,減少等待資料傳輸的時間,而在 GPU 上(多 CPU/GPU 系統也是)則是透過分塊將計算任務切割成數十甚至數百個子任務以便分配給不同硬體運算,增加平行度。

總而言之,根據出發點的不同,分塊矩陣乘法的分塊一詞可能對應到 block 也可能對應到 tiled,除此之外,部分讀者可能會發現我們不只可以對原本的矩陣分塊得到子矩陣,也能對子矩陣繼續分塊,實際上本文後續的矩陣乘法加速即包含兩次分塊,第一次分塊在軟體上將矩陣乘法切分成矩陣乘法加速器可以計算的大小,但因為實際的計算電路只能執行更小規模的矩陣乘法,需要在硬體上對傳入的矩陣進行第二次分塊,設計細節可以在後續章節看到。

此處我們引入矩陣分塊的參考程式碼,令 T 表示分塊的大小,也就是說矩陣會被切成 T×T 的數個子矩陣,雖然此處假設分塊尺寸為正方形,但根據需求的不同分塊尺寸可以是長方形,而且為了簡潔起見,此處假設 T 整除 MNK,實際程式碼需要考慮未能整除的部分。

# Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
# Tiled GEMM

for m in range(0, M, T):
for k in range(0, K, T):
for n in range(0, N, T):
for mm in range(m, m+T):
for kk in range(k, k+T):
for nn in range(n, n+T):
C[mm][nn] += A[mm][kk] * B[kk][nn]

後續我們將這個方法稱為 MKNMKN,雖然程式碼看起來貌似很複雜,但我們只要將最外層三個迴圈遮起來,就會發現裡面三層就是單純的矩陣乘法。

接下來觀察其存取模式,設 mnk 為 0,T 為 2,這次我們完整跑一遍內部三層迴圈。

mm = 0, kk = 0, nn = 0 to 1
---------------------------
A[0][0], A[0][0]
B[0][0], B[0][1]
C[0][0], C[0][1]

mm = 0, kk = 1, nn = 0 to 1
---------------------------
A[0][1], A[0][1]
B[1][0], B[1][1]
C[0][0], C[0][1]

mm = 1, kk = 0, nn = 0 to 1
---------------------------
A[1][0], A[1][0]
B[0][0], B[0][1]
C[1][0], C[1][1]

mm = 1, kk = 1, nn = 0 to 1
---------------------------
A[1][1], A[1][1]
B[1][0], B[1][1]
C[1][0], C[1][1]

假設快取在資料不存在時會連續取兩個元素長度的連續資料、快取記憶體容量只夠存放六條兩個元素長度的連續資料且快取按先進先出的順序儲存(這不太現實 xD),那麼上面的分塊策略剛好能讓計算所需元素都放在快取記憶體裡,當然以上推論省略了非常多細節(如記憶體對齊或快取對映方法等),更加詳細的討論讀者請移步計算機架構課程。

雖然這裡舉的例子非常簡單且侷限,但希望這個例子能讓讀者感覺到矩陣分塊對於最佳化快取記憶體表現的潛力,以下為簡化版快取記憶體的內容,供讀者參考。

Cache Content                 Event
------------------------ --------------
Line 0: A[0][0], A[0][1] A[0][0] access
Line 1: B[0][0], B[0][1] B[0][0] access
Line 2: C[0][0], C[0][1] C[0][0] access
Line 3: B[1][0], B[1][1] B[1][0] access
Line 4: A[1][0], A[1][1] A[1][0] access
Line 5: C[1][0], C[1][1] C[1][0] access

接著讓我們看看 MKNMKN 在不同分塊大小下的實際表現,後綴 _x 表示分塊大小為 x

M = 512, N = 512, K = 1024
-----------------------------
MKN: 14750774696 cycles
-----------------------------
MKNMKN_2: 9187992217 cycles
MKNMKN_4: 6518496316 cycles
MKNMKN_8: 6017335446 cycles
MKNMKN_16: 6462769240 cycles
MKNMKN_32: 16127959473 cycles
MKNMKN_64: 15779409116 cycles

可以發現 MKNMKN_8 的表現最佳,運行速度約為 MKN 的 2.45 倍。如同前面對 MKN 做迴圈展開一樣,我們也能對 MKNMKN 的最內層迴圈完全展開得到加速,以後綴 _u 表示將最內層迴圈完全展開,從下面的數據可以看到當分塊大小為 4 和 8 時,將最內層迴圈展開能夠減少需要的周期數,雖然不多就是了。

M = 512, N = 512, K = 1024
------------------------------
MKNMKN_4: 6518496316 cycles
MKNMKN_8: 6017335446 cycles
MKNMKN_16: 6462769240 cycles
------------------------------
MKNMKN_4_u: 6393604653 cycles
MKNMKN_8_u: 6008034654 cycles
MKNMKN_16_u: 6483322167 cycles

等等,展開一層迴圈夠嗎?有沒有可能再額外展開一層迴圈能得到收益呢?答案是會,而且加速相當明顯!以下以 _uu 後綴表示將先將最內層迴圈完全展開,再將倒數第二層迴圈完全展開。

M = 512, N = 512, K = 1024
-------------------------------
MKNMKN_4_u: 6393604653 cycles
MKNMKN_8_u: 6008034654 cycles
MKNMKN_16_u: 6483322167 cycles
-------------------------------
MKNMKN_4_uu: 6126239534 cycles
MKNMKN_8_uu: 4003743694 cycles
MKNMKN_16_uu: 5686292694 cycles

對於 MKNMKN_8 來說,MKNMKN_8_u 的加速比約為 1,但再額外多展開一層迴圈的情況下, MKNMKN_8_uu 得到的加速比約為 1.5,相當不講道理,至於額外再展開一層迴圈會不會有收益就待有興趣的讀者研究,筆者做實驗做到累了。

除了正方形分塊以外,我們也能採用不對稱的分塊策略,例如下方程式碼展示的方法,將左矩陣和右矩陣分別切成 1×T 的向量和 T×T 的矩陣。

# Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
# Tiled GEMM in Vector-Matrix Manner

for k in range(0, K, T):
for n in range(0, N, T):
for m in range(M):
for kk in range(k, k+T):
for nn in range(n, n+T):
C[m][nn] += A[m][kk] * B[kk][nn]

依照前面的慣例稱這個方法為 KNMKN,因為最內層的 m 一次只增加 1,我們實際上能將該方法視為 MKNMKN 的特例,其實驗結果展示如下。

M = 512, N = 512, K = 1024
------------------------------
MKNMKN_8_uu: 4003743694 cycles
------------------------------
KNMKN_2: 9487730297 cycles
KNMKN_4: 5583226743 cycles
KNMKN_8: 5914475719 cycles
KNMKN_16: 6378373923 cycles
KNMKN_32: 16125694545 cycles
KNMKN_64: 15779330036 cycles
------------------------------
KNMKN_4_u: 5606538441 cycles
KNMKN_8_u: 5938096362 cycles
KNMKN_16_u: 6498209409 cycles
------------------------------
KNMKN_4_uu: 5590906823 cycles
KNMKN_8_uu: 3388916297 cycles
KNMKN_16_uu: 4679968567 cycles

可以發現 KNMKN_8_uu 又比先前最快的 MKNMKN_8_uu 有著約 1.18 倍的加速。

在本章節中介紹了包含迴圈順序排列、迴圈展開和矩陣分塊等不同的最佳化技巧,從前面的實驗數據中不難看出不同最佳化技巧的使用順序和次數會顯著影響演算法的運行時間,對於牽涉多層迴圈運算的演算法來說,較好的最佳化策略通常需要人們在目標機器上進行大量人工或自動化測試才能得到(當然一言不合上 AI 也是可以的,大名鼎鼎的深度學習編譯器 TVM 便是使用 AI 模型來決定),不過考慮到我們會用矩陣乘法單元而不是單純 CPU 進行矩陣乘法的運算,倒也不需要太過煩惱這些事情。

--

--

Yeecy

A Ph.D. student at NYCU CS and a compiler engineer at ICEshell Co., Ltd. You can find more information about me on my GitHub page github.com/ADNRs.