CUDA矩陣乘法優(yōu)化手段詳解

單精度矩陣乘法（SGEMM）幾乎是每一位學(xué)習(xí) CUDA 的同學(xué)繞不開(kāi)的案例，這個(gè)經(jīng)典的計(jì)算密集型案例可以很好地展示 GPU 編程中常用的優(yōu)化技巧。本文將詳細(xì)介紹 CUDA SGEMM 的優(yōu)化手段，適合認(rèn)真閱讀過(guò)《CUDA C++Programming Guide》，具備一定 CUDA 編程基礎(chǔ)的同學(xué)閱讀，希望能給追求極致性能的同學(xué)們一些啟發(fā)。

CUDA 矩陣乘法優(yōu)化手段詳解

Naive 實(shí)現(xiàn)的分析：到底差在哪里？

筆者面試過(guò)不少具有 CUDA 編程經(jīng)驗(yàn)的校招同學(xué)，當(dāng)提問(wèn)使用 CUDA 編寫(xiě)一個(gè) SGEMM Kernel 的時(shí)候，通常會(huì)獲得這么一個(gè)答案：

__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

	

	這樣一個(gè) Naive 的 Kernel 當(dāng)然不是筆者所期待的，因?yàn)檫@個(gè) Kernel 的性能基本可以斷定連 cublas 的 1/10 都不到，顯然不符合我們追求高性能的需求。那么這個(gè) Naive 的實(shí)現(xiàn)究竟差在哪呢？

	分析代碼我們可以看到，計(jì)算一次 FMA（乘累加）之前需要讀一次 A 和讀一次 B，眾所周知，讀取 Global Memory 的代價(jià)很大，通常都需要幾百個(gè) cycle（時(shí)鐘周期），而計(jì)算一次 FMA 通常只需要幾個(gè) cycle，大量的時(shí)間被花費(fèi)在了訪存上。也許有思維活絡(luò)的同學(xué)立馬想到，可以將 A 和 B 矩陣先搬運(yùn)到 Shared Memory（SM 中低延遲的 on-chip memory，block 內(nèi)線程共享，附 NVIDIA GPU 內(nèi)存結(jié)構(gòu)圖）中降低訪存的開(kāi)銷，這的確是一個(gè)很好的思路，但是這只能將訪存代價(jià)從幾百 cycle 降低到幾十 cycle，并不改變問(wèn)題的本質(zhì)。問(wèn)題的關(guān)鍵在于主體循環(huán)由兩條 Load 指令與一條 FMA 指令構(gòu)成，計(jì)算指令只占總體的 1/3，計(jì)算訪存比過(guò)低，最終導(dǎo)致了訪存延遲不能被隱藏，從而性能不理想。

	

	讓我們打開(kāi)思路，若一個(gè) thread 并不只計(jì)算一個(gè)結(jié)果，而是計(jì)算 4x4 個(gè)結(jié)果，并且使用 Shared Memory 優(yōu)化，Hot Loop 會(huì)是什么樣呢，偽代碼如下所示：

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

	

	分析可以得出從 smemA 讀取到寄存器 a_reg 中，需要進(jìn)行 4 次訪存操作，B 同理，那么主體的計(jì)算訪存指令比例變成了 16/8，相對(duì)于之前的情況，計(jì)算指令的占比大大提高了。足夠大的計(jì)算訪存比能提升計(jì)算單元的利用率，并能起到隱藏訪存延遲的作用。我們可以進(jìn)一步提升計(jì)算訪存比，從而使得 kernel 的性能接近理論峰值。

	矩陣分塊與資源分配

	顯然我們不能只使用一個(gè) block 計(jì)算一個(gè)超大矩陣，這樣會(huì)造成大量 SM（Streaming Multiprocessor）的閑置浪費(fèi)，這就需要對(duì)矩陣進(jìn)行分塊計(jì)算，如下圖所示：

	

	不同的分塊大小在不同 shape 的矩陣乘法應(yīng)用上性能各有優(yōu)劣，本文選取 128x128 的分塊舉例。

	從上一小節(jié)我們可以看到，提升計(jì)算訪存比有很大的好處，那么計(jì)算訪存比可以無(wú)限提升嗎，答案是否定的。因?yàn)橐嵘?jì)算訪存比，單個(gè) thread 就需要計(jì)算一個(gè)更大的塊，這就需要更多的寄存器，但寄存器的個(gè)數(shù)是有限的。以 Turing 架構(gòu)的 GPU 為例，單個(gè) SM 的寄存器總量為 65536，因?yàn)橹噶罹幋a的限制，單個(gè) thread 能使用的最大寄存器個(gè)數(shù)為 255，并且寄存器個(gè)數(shù)并不是用得越多越好。這里需要引入一個(gè) Occupancy（占用率）的概念，Occupancy 是指每個(gè) SM 中活動(dòng)線程束（Warp）數(shù)量與最大并發(fā)線程束數(shù)量的比值，高的 Occupancy 不一定意味著高性能，但可以通過(guò)切換執(zhí)行 Warp 來(lái)起到一定隱藏延遲的作用。而每個(gè) SM 中的 Active Warp 數(shù)量，取決于 block 使用的資源數(shù)量，具體為每個(gè)線程使用的寄存器個(gè)數(shù)與 Shared Memory 用量。Occupany可通過(guò) CUDA Toolkit 中提供的 CUDA_Occupancy_Calculator.xls 工具獲得。

	考慮一個(gè) block 計(jì)算 128x128 的分塊，若每個(gè)線程計(jì)算 128 個(gè)結(jié)果，需要的 block size 為 128，單個(gè)線程需要 128 個(gè)寄存器儲(chǔ)存計(jì)算結(jié)果，加上所需的 Gmem to Smem，Smem to Reg 等一些所需的寄存器，大概共需要至少 180 多個(gè)，計(jì)算 Occupany 可知此時(shí)的 Active Warp 數(shù)只有 8，Occupany 為 25%；若設(shè)置 block size 為 256，則每個(gè)線程僅需計(jì)算 64 個(gè)結(jié)果，調(diào)整寄存器和 Shared Memory 的使用量并觀察 Occupany，可知若每個(gè)線程只使用 128 個(gè)寄存器，block 內(nèi)的 Shared Memory 使用量限制在 32K，Active Warp 數(shù)可以達(dá)到 16，是一個(gè)更優(yōu)的選擇：

	

	并且此時(shí)的配置計(jì)算訪存比可以達(dá)到 64/4（使用向量讀取），已經(jīng)足夠隱藏訪存延遲。

	極致的訪存優(yōu)化

	通常情況下，在選取了合適的 block 資源配置，利用 Shared Memory 降低訪存延遲，做好循環(huán)展開(kāi)之后，SGEMM Kernel 的性能已經(jīng)能達(dá)到一個(gè)不錯(cuò)的水平（80% cublas），但這并不是我們旅程的終點(diǎn)。首先，我們可以使用向量讀取指令LDS.128優(yōu)化 Shared Memory 訪問(wèn)（對(duì)應(yīng) float4 數(shù)據(jù)類型），這能大幅減少訪存指令的數(shù)量，進(jìn)一步提升計(jì)算訪存比，由此我們需要將 A 矩陣存入 smemA 之前做一次轉(zhuǎn)置：

	

	同時(shí)，我們的 kernel 為 256 個(gè)線程計(jì)算 128x128 的分塊，為了能夠合并訪問(wèn) Shared Memory，我們將 256 個(gè)線程劃為二維，令：

	
inttx=threadIdx.x%16;
intty=threadIdx.x/16;


	

	并按照如下方式向量讀取 Shared Memory 中的數(shù)據(jù)：

	

	最終單個(gè)線程計(jì)算 2x2 個(gè) 4x4 的結(jié)果，結(jié)果布局如圖所示：

	

	并且通過(guò) micro benchmark 可以探測(cè)出，Turing(Tesla T4) 的 Global Memory 的訪存延遲約 300 cycle，Shared Memory 的訪存延遲在約 30 cycle，需要充分利用 Prefetch 的思想，隱藏 Global Memory 讀入中間寄存器、將來(lái)自 Global Memory 的數(shù)據(jù)塊寫(xiě)入 Shared Memory、從 Shared Memory 中讀出數(shù)據(jù)塊的訪存延遲，以免計(jì)算單元因?yàn)?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

	

	注：此處偷懶假設(shè)了 M、N、K 都是 4 的倍數(shù)，若非 4 的倍數(shù)則 Global Memory 不能使用 float4 進(jìn)行讀取，結(jié)果也不能用 float4 進(jìn)行寫(xiě)回，而且為了合并寫(xiě)回，需要通過(guò) Shared Memory 交換 warp 內(nèi)的結(jié)果，保證每個(gè) warp 執(zhí)行一條 Store 指令能夠?qū)懟匾黄B續(xù)的內(nèi)存空間。

	至此我們獲得了一個(gè)充分優(yōu)化的 SGEMM Kernel。另外 Ampere GPU 新增了LDGSTS指令，數(shù)據(jù)塊從 Global Memory 到 Shared Memory 的過(guò)程不需要經(jīng)過(guò)中間寄存器，可以進(jìn)一步的優(yōu)化 SGEMM 的性能。

	性能對(duì)比

	為了避免 cublas 選取到 split K 的 Kernel，我們將 K 固定為 1024，取 M, N = 2048, 4096, 8192 和 16384 作為測(cè)試用例，對(duì)比了上述 SGEMM Kernel 與 cublas 的性能（測(cè)試 GPU 為 Tesla T4，鎖定核心頻率為 1100)：

	

	可以看到所實(shí)現(xiàn)的 SGEMM Kernel 達(dá)到了 cublas 平均 97.5% 的性能。

	超越 cublas：使用 SASS 調(diào)優(yōu) Kernel

	到這里，可能有同學(xué)依然有一個(gè)疑問(wèn)，我們似乎把所有能想到的優(yōu)化手段都用上了，為什么寫(xiě)出來(lái)的 CUDA C Kernel 依然離 cublas 有一定的差距，答案是 cublas 所使用的 kernel 中有一大部分并不是通過(guò) nvcc 編譯的 CUDA Kernel，而是使用 NVIDIA GPU 的匯編語(yǔ)言（Shader Assembly，簡(jiǎn)稱 SASS）編寫(xiě)的深度調(diào)優(yōu)版本。

	盡管 nvcc 編譯器在不斷的進(jìn)步，特別是 CUDA 11 中的 nvcc，所編譯的 Kernel 與手工匯編優(yōu)化版本之間的差距已大幅縮小，但仍然無(wú)法完全避免寄存器 Bank conflict 的影響以及充分利用寄存器的 Reuse Cache（這兩個(gè)概念下面會(huì)進(jìn)行詳細(xì)的介紹），使得差距仍然存在。即使 PTX 這樣的偽匯編語(yǔ)言，也無(wú)法精確控制寄存器的分配，和 CUDA C 面臨著一樣的困境。

	所以為了充分挖掘 GPU 的性能極限，需要對(duì) GPU 指令和寄存器進(jìn)行精確控制，就必須交由 GPU 原生匯編語(yǔ)言 SASS 完成。這方面已經(jīng)有了很多研究，如出自 Citadel 的深入研究 NV GPU 架構(gòu)的 Dissecting the NVidia XXX GPU architecture via microbenchmarking 系列論文，這一系列文章對(duì)底層架構(gòu)做了系統(tǒng)的測(cè)試、分析和總結(jié)，雖然其中某些結(jié)論可能并不準(zhǔn)確，但總體來(lái)講有很高的參考價(jià)值。同時(shí)催生了不少開(kāi)源匯編器如 KeplerAs、maxas（最成熟，影響深遠(yuǎn)）、turingas 和 CuAssembler 等一系列開(kāi)源 SASS 匯編器，使得使用 SASS 編寫(xiě)高性能 Kernel 變成了可能。

	寄存器 Bank conflict

	我們知道 Shared Memory 有 Bank conflict，而寄存器的 Bank conflict 也是類似的概念。NVIDIA GPU 每個(gè) SM 有獨(dú)立的 Register File，而 Register File 被分為若干個(gè) Bank，以 Maxwell 為例，若一條指令所需的源寄存器有 2 個(gè)以上來(lái)自于同一 Bank，則會(huì)產(chǎn)生 conflict，指令會(huì)相當(dāng)于重發(fā)射，浪費(fèi)一個(gè) cycle。Maxwell/Pascal 的 Register File 的 Bank 數(shù)為 4，寄存器的id%4即為該寄存器的所屬 bank（如 R0 屬于 Bank 0，R5 屬于 Bank 1），FFMA R1, R0, R4, R1這樣的指令就會(huì)產(chǎn)生寄存器 Bank conflict。而 Turing 架構(gòu)做了改進(jìn)，Register File 被分為 2 個(gè) Bank，每個(gè) Bank 有 2 個(gè) Port，若非三個(gè)源寄存器 id 同奇偶則不會(huì)產(chǎn)生沖突，大大緩解了寄存器 Bank conflict。

	maxas 中的 Maxwell SGEMM SASS Kernel 為了緩解寄存器 Bank conflict，就對(duì)參與 FFMA 計(jì)算的寄存器做了精巧的分配（參考 maxas 的 SGEMM 文檔），如下圖所示：

	

	經(jīng)過(guò)對(duì) C 的巧妙排布，寄存器 Bank conflict 大大減少，但依然無(wú)法完全避免（如上圖中黑框標(biāo)識(shí)的部分，A/B 所使用的寄存器會(huì)產(chǎn)生 Bank conflict），這部分沖突就需要用到寄存器 Reuse 來(lái)消除。

	Register Reuse

	寄存器 Reuse 是 NVIDIA 為了緩解寄存器 Bank conflict 的問(wèn)題，在 Maxwell 開(kāi)始引入的一種機(jī)制，NVIDIA 在讀取指令操作數(shù)的 Collector 單元加入了寄存器的 Reuse Cache。Reuse Cache 是只讀的，指令獲取 Operand 是否通過(guò)此 Cache 由該指令的 control code（maxas 的 control code wiki中有詳細(xì)的介紹）所指定，使用 cuobjdump 反匯編一些 Kernel 可以發(fā)現(xiàn)一些寄存器后有 .reuse的 flag，即表示該寄存器從 Reuse Cache 而非 Register File 中取值，從而消除寄存器 Bank conflict：

	
#MaxwellGPU
FFMAR2,R64.reuse,R73,R2;#R64進(jìn)入ReuseCache
FFMAR3,R64.reuse,R72,R3;#R64從ReuseCache中獲取，避免與R72沖突


	

	但是使用 .reuse需要滿足一定條件（寄存器將被改寫(xiě)前不能設(shè)置 .reuse），胡亂設(shè)置 reuse flag 會(huì)有可能獲取的是歷史值，造成計(jì)算錯(cuò)誤，根據(jù)筆者的理解，.reuse更像是使該寄存器的值在 Reuse Cache 中 hold 住的標(biāo)識(shí)。nvcc 編譯 CUDA Kernel 也會(huì)使用 Reuse Cache 去規(guī)避一些寄存器 Bank conflict，但是因?yàn)榧拇嫫鞣峙浼爸噶钆挪嫉脑?，Reuse 的利用率并不高，反匯編我們剛才寫(xiě)的 SGEMM Kernel，對(duì)主循環(huán)的所有 FFMA 指令做個(gè)統(tǒng)計(jì)，可以發(fā)現(xiàn) Reuse Cache 僅達(dá)到 20% 左右，而 maxas 的 SASS Kernel 通過(guò)設(shè)計(jì)使得 Reuse 的利用率可以達(dá)到 49%。

	最終通過(guò) SASS 精細(xì)調(diào)優(yōu)的 SGEMM Kernel 的性能可以全面超越 cublas，感興趣的同學(xué)們可以自行編譯 maxas 中的 SGEMM Kernel 在 Maxwell 或者 Pascal GPU 上進(jìn)行測(cè)試。最后，雖然使用 SASS 能充分挖掘 GPU 的性能，但面臨有三大問(wèn)題：1. 第三方 NV GPU 匯編器依賴于對(duì) GPU 架構(gòu)的逆向研究，可能因?yàn)闆](méi)有探究到全部的硬件底層細(xì)節(jié)而存在未知的 BUG；2. 匯編 Kernel 難于開(kāi)發(fā)，更難于調(diào)試；3. NV 每一代 GPU 的 ISA（指令集）都不盡相同，需要不斷開(kāi)發(fā)對(duì)應(yīng)的匯編器和匯編 Kernel。正因?yàn)檫@幾大問(wèn)題的存在，使得使用 SASS 編寫(xiě) Kernel 是個(gè)費(fèi)時(shí)費(fèi)力的工作，除非有追求極致性能的需求，否則不建議輕易嘗試。

	GEMM 的延伸：優(yōu)化卷積運(yùn)算

	我們都知道優(yōu)化卷積運(yùn)算可以通過(guò) im2col 將卷積映射為矩陣乘法來(lái)實(shí)現(xiàn)，對(duì)于上述 SGEMM Kernel，只需要將 Global Memory 的數(shù)據(jù)搬運(yùn)到 Shared Memory 這一過(guò)程稍作修改，由對(duì)應(yīng)位置的映射變?yōu)?im2col 映射，SGEMM Kernel 就搖身一變成為了計(jì)算 Conv 的 Kernel，這即是 cudnn 卷積運(yùn)算的 Implicit Gemm 算法。而在 im2col 過(guò)程中，若直接計(jì)算指針的偏移量的話，會(huì)引入大量的整數(shù)除法和取余運(yùn)算，這是一筆不小的開(kāi)銷，所以可以將地址的偏移量在 host 端預(yù)先計(jì)算好，作為 param 傳入 kernel 中，則可以在需要時(shí)從常量?jī)?nèi)存中讀取，避免整數(shù)除法和取余，實(shí)現(xiàn) Implicit Precomp Gemm。

	總結(jié)

	本文詳細(xì)介紹了如何編寫(xiě)一個(gè)高效率的 CUDA SGEMM Kernel，并且介紹了使用 SASS 編程這一極限優(yōu)化性能的手段，并稍稍延伸展開(kāi)了通過(guò) Implicit Gemm 優(yōu)化卷積運(yùn)算的思路，希望可以給予有志于極致挖掘硬件性能的同學(xué)們一定的啟發(fā)。

	審核編輯：湯梓紅