作者：馬駿 | 曠視 MegEngine 架構師
前言
單精度矩陣乘法（SGEMM）幾乎是每一位學習 CUDA 的同學繞不開的案例，這個經典的計算密集型案例可以很好地展示 GPU 程式設計中常用的優化技巧，而能否寫出高效率的 SGEMM Kernel，也是反映一位 CUDA 程式設計師對 GPU 體系結構的理解程度的優秀考題。本文將詳細介紹 CUDA SGEMM 的優化手段，適合認真閱讀過 《CUDA C++ Programming Guide》，具備一定 CUDA 程式設計基礎的同學閱讀，希望能給追求極致效能的同學們一些啟發。
CUDA 矩陣乘法優化手段詳解
Naive 實現的分析：到底差在哪裡？
筆者面試過不少具有 CUDA 程式設計經驗的校招同學，當提問使用 CUDA 編寫一個 SGEMM Kernel 的時候，通常會獲得這麼一個答案：
__global__ void matrixMul(const float *A, const float *B, float *C,
int M, int N, int K) {
int tx = blockIdx.x * blockDim.x + threadIdx.x;
int ty = blockIdx.y * blockDim.y + threadIdx.y;
if(ty < M && tx < N) {
float c = 0;
for(int i = 0; i < K; ++i){
c += A[ty * K + i] * B[i * N + tx];
}
C[ty * N + tx] = c;
}
}
這樣一個 Naive 的 Kernel 當然不是筆者所期待的，因為這個 Kernel 的效能基本可以斷定連 cublas 的 1/10 都不到，顯然不符合我們追求高效能的需求。那麼這個 Naive 的實現究竟差在哪呢？
分析程式碼我們可以看到，計算一次 FMA（乘累加）之前需要讀一次 A 和讀一次 B，眾所周知，讀取 Global Memory 的代價很大，通常都需要幾百個 cycle（時鐘週期），而計算一次 FMA 通常只需要幾個 cycle，大量的時間被花費在了訪存上。也許有思維活絡的同學立馬想到，可以將 A 和 B 矩陣先搬運到 Shared Memory（SM 中低延遲的 on-chip memory，block 內執行緒共享，附 NVIDIA GPU 記憶體結構圖）中降低訪存的開銷，這的確是一個很好的思路，但是這隻能將訪存代價從幾百 cycle 降低到幾十 cycle，並不改變問題的本質。問題的關鍵在於主體迴圈由兩條 Load 指令與一條 FMA 指令構成，計算指令只佔總體的 1/3，計算訪存比過低，最終導致了訪存延遲不能被隱藏，從而效能不理想。
讓我們開啟思路，若一個 thread 並不只計算一個結果，而是計算 4x4 個結果，並且使用 Shared Memory 優化，Hot Loop 會是什麼樣呢，虛擬碼如下所示：
float c[4][4] = {{0}};
float a_reg[4];
float b_reg[4];
for(int i = 0; i < K; i += TILE_K){
__syncthreads();
// transfer tile from global mem to shared mem
load_gmem_tile_to_smem(A, i, smemA);
load_gmem_tile_to_smem(B, i, smemB);
__syncthreads();
#pragma unroll
for(int j = 0; j < TILE_K; ++j) {
// load tile from shared mem to register
load_smem_tile_to_reg(smemA, j, a_reg);
load_smem_tile_to_reg(smemB, j, b_reg);
// compute matrix multiply accumulate 4x4
mma4x4(a_reg, b_reg, c)；
}
}
分析可以得出從 smemA 讀取到暫存器 a_reg 中，需要進行 4 次訪存操作，B 同理，那麼主體的計算訪存指令比例變成了 16/8，相對於之前的情況，計算指令的佔比大大提高了。足夠大的計算訪存比能提升計算單元的利用率，並能起到隱藏訪存延遲的作用。我們可以進一步提升計算訪存比，從而使得 kernel 的效能接近理論峰值。
矩陣分塊與資源分配
顯然我們不能只使用一個 block 計算一個超大矩陣，這樣會造成大量 SM（Streaming Multiprocessor）的閒置浪費，這就需要對矩陣進行分塊計算，如下圖所示：
不同的分塊大小在不同 shape 的矩陣乘法應用上效能各有優劣，本文選取 128x128 的分塊舉例。
從上一小節我們可以看到，提升計算訪存比有很大的好處，那麼計算訪存比可以無限提升嗎，答案是否定的。因為要提升計算訪存比，單個 thread 就需要計算一個更大的塊，這就需要更多的暫存器，但暫存器的個數是有限的。以 Turing 架構的 GPU 為例，單個 SM 的暫存器總量為 65536，因為指令編碼的限制，單個 thread 能使用的最大暫存器個數為 255，並且暫存器個數並不是用得越多越好。這裡需要引入一個 Occupancy（佔用率）的概念，Occupancy 是指每個 SM 中活動執行緒束（Warp）數量與最大併發執行緒束數量的比值，高的 Occupancy 不一定意味著高效能，但可以通過切換執行 Warp 來起到一定隱藏延遲的作用。而每個 SM 中的 Active Warp 數量，取決於 block 使用的資源數量，具體為每個執行緒使用的暫存器個數與 Shared Memory 用量。Occupany 可通過 CUDA Toolkit 中提供的 CUDA_Occupancy_Calculator.xls 工具獲得。
考慮一個 block 計算 128x128 的分塊，若每個執行緒計算 128 個結果，需要的 block size 為 128，單個執行緒需要 128 個暫存器儲存計算結果，加上所需的 Gmem to Smem，Smem to Reg 等一些所需的暫存器，大概共需要至少 180 多個，計算 Occupany 可知此時的 Active Warp 數只有 8，Occupany 為 25%；若設定 block size 為 256，則每個執行緒僅需計算 64 個結果，調整暫存器和 Shared Memory 的使用量並觀察 Occupany，可知若每個執行緒只使用 128 個暫存器，block 內的 Shared Memory 使用量限制在 32K，Active Warp 數可以達到 16，是一個更優的選擇：
並且此時的配置計算訪存比可以達到 64/4（使用向量讀取），已經足夠隱藏訪存延遲。
極致的訪存優化
通常情況下，在選取了合適的 block 資源配置，利用 Shared Memory 降低訪存延遲，做好迴圈展開之後，SGEMM Kernel 的效能已經能達到一個不錯的水平（80% cublas），但這並不是我們旅程的終點。首先，我們可以使用向量讀取指令LDS.128優化 Shared Memory 訪問（對應 float4 資料型別），這能大幅減少訪存指令的數量，進一步提升計算訪存比，由此我們需要將 A 矩陣存入 smemA 之前做一次轉置：
同時，我們的 kernel 為 256 個執行緒計算 128x128 的分塊，為了能夠合併訪問 Shared Memory，我們將 256 個執行緒劃為二維，令：
int tx = threadIdx.x % 16;
int ty = threadIdx.x / 16;
並按照如下方式向量讀取 Shared Memory 中的資料：
最終單個執行緒計算 2x2 個 4x4 的結果，結果佈局如圖所示：
並且通過 micro benchmark 可以探測出，Turing(Tesla T4) 的 Global Memory 的訪存延遲約 300 cycle，Shared Memory 的訪存延遲在約 30 cycle，需要充分利用** Prefetch **的思想，隱藏 Global Memory 讀入中間暫存器、將來自 Global Memory 的資料塊寫入 Shared Memory、從 Shared Memory 中讀出資料塊的訪存延遲，以免計算單元因為 stall 而空閒太久，最終的虛擬碼如下所示：
#define TILE_K 16
__shared__ float4 smemA[2][TILE_K * 128 / 4];
__shared__ float4 smemB[2][TILE_K * 128 / 4];
float4 c[8][2] = {{make_float4(0.f, 0.f, 0.f, 0.f)}};
float4 ldg_a_reg[2];
float4 ldg_b_reg[2];
float4 a_reg[2][2];
float4 b_reg[2][2];
// transfer first tile from global mem to shared mem
load_gmem_tile_to_reg(A, 0, ldg_a_reg);
load_gmem_tile_to_reg(B, 0, ldg_b_reg);
store_reg_to_smem_tile_transpose(ldg_a_reg, 0, smemA[0]);
store_reg_to_smem_tile(ldg_b_reg, 0, smemB[0]);
__syncthreads();
// load first tile from shared mem to register
load_smem_tile_to_reg(smemA[0], 0, a_reg[0]);
load_smem_tile_to_reg(smemB[0], 0, b_reg[0]);
int write_stage_idx = 1; //ping pong switch
do {
i += TILE_K;
// load next tile from global mem
load_gmem_tile_to_reg(A, i, ldg_a_reg);
load_gmem_tile_to_reg(B, i, ldg_b_reg);
int load_stage_idx = write_stage_idx ^ 1;
#pragma unroll
for(int j = 0; j < TILE_K - 1; ++j) {
// load next tile from shared mem to register
load_smem_tile_to_reg(smemA[load_stage_idx], j + 1, a_reg[(j + 1) % 2]);
load_smem_tile_to_reg(smemB[load_stage_idx], j + 1, b_reg[(j + 1) % 2]);
// compute matrix multiply accumulate 8x8
mma8x8(a_reg[j % 2], b_reg[j % 2], c)；
}
if(i < K) {
// store next tile to shared mem
store_reg_to_smem_tile_transpose(ldg_a_reg, 0, smemA[write_stage_idx]);
store_reg_to_smem_tile(ldg_b_reg, 0, smemB[write_stage_idx]);
// use double buffer, only need one sync
__syncthreads();
// switch
write_stage_idx ^= 1;
}
// load first tile from shared mem to register of next iter
load_smem_tile_to_reg(smemA[load_stage_idx ^ 1], 0, a_reg[0]);
load_smem_tile_to_reg(smemB[load_stage_idx ^ 1], 0, b_reg[0]);
// compute last tile mma 8x8
mma8x8(a_reg[1], b_reg[1], c)；
} while (i < K);
store_c(c, C);
注：此處偷懶假設了 M、N、K 都是 4 的倍數，若非 4 的倍數則 Global Memory 不能使用 float4 進行讀取，結果也不能用 float4 進行寫回，而且為了合併寫回，需要通過 Shared Memory 交換 warp 內的結果，保證每個 warp 執行一條 Store 指令能夠寫回一片連續的記憶體空間。
至此我們獲得了一個充分優化的 SGEMM Kernel。另外 Ampere GPU 新增了LDGSTS指令，資料塊從 Global Memory 到 Shared Memory 的過程不需要經過中間暫存器，可以進一步的優化 SGEMM 的效能。
效能對比
為了避免 cublas 選取到 split K 的 Kernel，我們將 K 固定為 1024，取 M, N = 2048, 4096, 8192 和 16384 作為測試用例，對比了上述 SGEMM Kernel 與 cublas 的效能（測試 GPU 為 Tesla T4，鎖定核心頻率為 1100)：
可以看到所實現的 SGEMM Kernel 達到了 cublas 平均 97.5% 的效能。
超越 cublas：使用 SASS 調優 Kernel
到這裡，可能有同學依然有一個疑問，我們似乎把所有能想到的優化手段都用上了，為什麼寫出來的 CUDA C Kernel 依然離 cublas 有一定的差距，答案是 cublas 所使用的 kernel 中有一大部分並不是通過 nvcc 編譯的 CUDA Kernel，而是使用 NVIDIA GPU 的組合語言（Shader Assembly，簡稱 SASS）編寫的深度調優版本。儘管 nvcc 編譯器在不斷的進步，特別是 CUDA 11 中的 nvcc，所編譯的 Kernel 與手工彙編優化版本之間的差距已大幅縮小，但仍然無法完全避免暫存器 Bank conflict 的影響以及充分利用暫存器的 Reuse Cache（這兩個概念下面會進行詳細的介紹），使得差距仍然存在。即使 PTX 這樣的偽組合語言，也無法精確控制暫存器的分配，和 CUDA C 面臨著一樣的困境。所以為了充分挖掘 GPU 的效能極限，需要對 GPU 指令和暫存器進行精確控制，就必須交由 GPU 原生組合語言 SASS 完成。這方面已經有了很多研究，如出自 Citadel 的深入研究 NV GPU 架構的 Dissecting the NVidia XXX GPU architecture via microbenchmarking 系列論文，這一系列文章對底層架構做了系統的測試、分析和總結，雖然其中某些結論可能並不準確，但總體來講有很高的參考價值。同時催生了不少開源彙編器如 KeplerAs、maxas（最成熟，影響深遠）、turingas 和 CuAssembler 等一系列開源 SASS 彙編器，使得使用 SASS 編寫高效能 Kernel 變成了可能。
暫存器 Bank conflict
我們知道 Shared Memory 有 Bank conflict，而暫存器的 Bank conflict 也是類似的概念。NVIDIA GPU 每個 SM 有獨立的 Register File，而 Register File 被分為若干個 Bank，以 Maxwell 為例，若一條指令所需的源暫存器有 2 個以上來自於同一 Bank，則會產生 conflict，指令會相當於重發射，浪費一個 cycle。Maxwell/Pascal 的 Register File 的 Bank 數為 4，暫存器的id%4即為該暫存器的所屬 bank（如 R0 屬於 Bank 0，R5 屬於 Bank 1），FFMA R1, R0, R4, R1這樣的指令就回產生暫存器 Bank conflict。而 Turing 架構做了改進，Register File 被分為 2 個 Bank，每個 Bank 有 2 個 Port，若非三個源暫存器 id 同奇偶則不會產生衝突，大大緩解了暫存器 Bank conflict。
maxas 中的 Maxwell SGEMM SASS Kernel 為了緩解暫存器 Bank conflict，就對參與 FFMA 計算的暫存器做了精巧的分配（參考 maxas 的 SGEMM 文件），如下圖所示：
經過對 C 的巧妙排布，暫存器 Bank conflict 大大減少，但依然無法完全避免（如上圖中黑框標識的部分，A/B 所使用的暫存器會產生 Bank conflict），這部分衝突就需要用到暫存器 Reuse 來消除。
Register Reuse
暫存器 Reuse 是 NVIDIA 為了緩解暫存器 Bank conflict 的問題，在 Maxwell 開始引入的一種機制，NVIDIA 在讀取指令運算元的 Collector 單元加入了暫存器的 Reuse Cache。Reuse Cache 是隻讀的，指令獲取 Operand 是否通過此 Cache 由該指令的 control code（maxas 的 control code wiki 中有詳細的介紹）所指定，使用 cuobjdump 反彙編一些 Kernel 可以發現一些暫存器後有.reuse的 flag，即表示該暫存器從 Reuse Cache 而非 Register File 中取值，從而消除暫存器 Bank conflict：
# Maxwell GPU
FFMA R2, R64.reuse, R73, R2; # R64 進入 Reuse Cache
FFMA R3, R64.reuse, R72, R3; # R64 從 Reuse Cache 中獲取，避免與 R72 衝突
但是使用.reuse需要滿足一定條件（暫存器將被改寫前不能設定.reuse），胡亂設定 reuse flag 會有可能獲取的是歷史值，造成計算錯誤，根據筆者的理解，.reuse更像是使該暫存器的值在 Reuse Cache 中 hold 住的標識。nvcc 編譯 CUDA Kernel 也會使用 Reuse Cache 去規避一些暫存器 Bank conflict，但是因為暫存器分配及指令排布的原因，Reuse 的利用率並不高，反彙編我們剛才寫的 SGEMM Kernel，對主迴圈的所有 FFMA 指令做個統計，可以發現 Reuse Cache 僅達到 20% 左右，而 maxas 的 SASS Kernel 通過設計使得 Reuse 的利用率可以達到 49%。
最終通過 SASS 精細調優的 SGEMM Kernel 的效能可以全面超越 cublas，感興趣的同學們可以自行編譯 maxas 中的 SGEMM Kernel 在 Maxwell 或者 Pascal GPU 上進行測試。最後，雖然使用 SASS 能充分挖掘 GPU 的效能，但面臨有三大問題：1. 第三方 NV GPU 彙編器依賴於對 GPU 架構的逆向研究，可能因為沒有探究到全部的硬體底層細節而存在未知的 BUG；2. 彙編 Kernel 難於開發，更難於除錯；3. NV 每一代 GPU 的 ISA（指令集）都不盡相同，需要不斷開發對應的彙編器和彙編 Kernel。正因為這幾大問題的存在，使得使用 SASS 編寫 Kernel 是個費時費力的工作，除非有追求極致效能的需求，否則不建議輕易嘗試。
GEMM 的延伸：優化卷積運算
我們都知道優化卷積運算可以通過 im2col 將卷積對映為矩陣乘法來實現，對於上述 SGEMM Kernel，只需要將 Global Memory 的資料搬運到 Shared Memory 這一過程稍作修改，由對應位置的對映變為 im2col 對映，SGEMM Kernel 就搖身一變成為了計算 Conv 的 Kernel，這即是 cudnn 卷積運算的 Implicit Gemm 演算法。而在 im2col 過程中，若直接計算指標的偏移量的話，會引入大量的整數除法和取餘運算，這是一筆不小的開銷，所以可以將地址的偏移量在 host 端預先計算好，作為 param 傳入 kernel 中，則可以在需要時從常量記憶體中讀取，避免整數除法和取餘，實現 Implicit Precomp Gemm，詳細細節可以參看我們之前的文章 MegEngine TensorCore 卷積運算元實現原理。在這一思路的指引下，我們基於 cutlass 實現了高效的 int8/int4 卷積運算（MegEngine cutlass），並且在持續的迭代中，歡迎大家試用。
總結
本文詳細介紹瞭如何編寫一個高效率的 CUDA SGEMM Kernel，並且介紹了使用 SASS 程式設計這一極限優化效能的手段，並稍稍延伸展開了通過 Implicit Gemm 優化卷積運算的思路，希望可以給予有志於極致挖掘硬體效能的同學們一定的啟發。最後歡迎各位同學加入 MegEngine 團隊，一起優化深度學習的訓練和推理效能。
bot 說：文章都看完啦~ 是否有興趣加入到深度學習框架開發中來？
Megengine 團隊現正火熱招聘中！期待你的加入~
簡歷投遞或詳細資訊瞭解可新增微信：duoduo715495【框架開發工程師（C++）】
職位描述:
- 負責曠視核心深度框架 MegEngine 的設計，演進，實現，維護和優化
- 優化 MegEngine 在各個計算平臺（CUDA / Arm / x86 等）上的效能
- 持續探索改進深度學習框架的先進優化方法（例如圖優化，自動程式碼生成，超低 bit 量化，稀疏化等）
技能要求:
- 1-3 年的效能優化經驗（X86，CUDA，ARM，OpenCL 等）
- 有良好的資料結構與演算法功底，能夠熟練使用 C++ 語言編寫較複雜的演算法
- 對深度學習和深度學習框架（Caffe，Tensorflow，PyTorch 等）有基本瞭解者優先
- 有複雜系統設計經驗者優先