CUDA矩陣乘法優化手段詳解

單精度矩陣乘法（SGEMM）幾乎是每一位學習 CUDA 的同學繞不開的案例，這個經典的計算密集型案例可以很好地展示 GPU 編程中常用的優化技巧。本文將詳細介紹 CUDA SGEMM 的優化手段，適合認真閱讀過《CUDA C++Programming Guide》，具備一定 CUDA 編程基礎的同學閱讀，希望能給追求極致性能的同學們一些啟發。

CUDA 矩陣乘法優化手段詳解

Naive 實現的分析：到底差在哪里？

筆者面試過不少具有 CUDA 編程經驗的校招同學，當提問使用 CUDA 編寫一個 SGEMM Kernel 的時候，通常會獲得這么一個答案：

__global__voidmatrixMul(constfloat*A,constfloat*B,float*C,
intM,intN,intK){
inttx=blockIdx.x*blockDim.x+threadIdx.x;
intty=blockIdx.y*blockDim.y+threadIdx.y;
if(ty

	

	這樣一個 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 會是什么樣呢，偽代碼如下所示：

	
floatc[4][4]={{0}};
floata_reg[4];
floatb_reg[4];
for(inti=0;i

	

	分析可以得出從 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 個線程劃為二維，令：

	
inttx=threadIdx.x%16;
intty=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 而空閑太久，最終的偽代碼如下所示：

	
#defineTILE_K16
__shared__float4smemA[2][TILE_K*128/4];
__shared__float4smemB[2][TILE_K*128/4];
float4c[8][2]={{make_float4(0.f,0.f,0.f,0.f)}};
float4ldg_a_reg[2];
float4ldg_b_reg[2];
float4a_reg[2][2];
float4b_reg[2][2];

//transferfirsttilefromglobalmemtosharedmem
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();

//loadfirsttilefromsharedmemtoregister
load_smem_tile_to_reg(smemA[0],0,a_reg[0]);
load_smem_tile_to_reg(smemB[0],0,b_reg[0]);

intwrite_stage_idx=1;//pingpongswitch
do{
i+=TILE_K;
//loadnexttilefromglobalmem
load_gmem_tile_to_reg(A,i,ldg_a_reg);
load_gmem_tile_to_reg(B,i,ldg_b_reg);

intload_stage_idx=write_stage_idx^1;

#pragmaunroll
for(intj=0;j

	

	注：此處偷懶假設了 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：

	
#MaxwellGPU
FFMAR2,R64.reuse,R73,R2;#R64進入ReuseCache
FFMAR3,R64.reuse,R72,R3;#R64從ReuseCache中獲取，避免與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。

	總結

	本文詳細介紹了如何編寫一個高效率的 CUDA SGEMM Kernel，并且介紹了使用 SASS 編程這一極限優化性能的手段，并稍稍延伸展開了通過 Implicit Gemm 優化卷積運算的思路，希望可以給予有志于極致挖掘硬件性能的同學們一定的啟發。

	審核編輯：湯梓紅