本文將從一個(gè) cuda 初學(xué)者的角度來闡述如何優(yōu)化一個(gè)形狀較大的正方形乘正方形的 FP32 矩陣乘。

矩陣乘作為目前神經(jīng)網(wǎng)絡(luò)計(jì)算中占比最大的一個(gè)部分，其快慢會(huì)顯著影響神經(jīng)網(wǎng)絡(luò)的訓(xùn)練與推斷所消耗的時(shí)間。雖然現(xiàn)在市面上已經(jīng)有非常多的矩陣乘的高效實(shí)現(xiàn)——如基于 cpu 的 mkl、基于 arm 設(shè)備的 ncnn 與 emll、基于 cuda 的 cublas ——掌握了矩陣乘優(yōu)化的思路不僅能幫助你更好的理解編寫高性能代碼的一些基本原則，而且許多神經(jīng)網(wǎng)絡(luò)加速領(lǐng)域進(jìn)階的技巧如算子融合都是與矩陣乘交互從而達(dá)到更高的性能。

由于矩陣乘的性能優(yōu)化與兩個(gè)矩陣的形狀有著非常密切的聯(lián)系，因此，為了降低本文的撰寫難度（以及輔助讀者更好的理解矩陣乘優(yōu)化），本文將從一個(gè) cuda 初學(xué)者的角度來闡述如何優(yōu)化一個(gè)形狀較大的正方形乘正方形的 FP32 矩陣乘。同時(shí)本文按如下順序講解：

• Goals：本文的目標(biāo)是什么？

• Performance：我們達(dá)到了多少性能？

• 樸素 GEMM 與前置知識：簡單介紹一下我們的任務(wù)是什么，我們需要提前了解什么。

• Tiling：如何做矩陣分塊？即如何將一個(gè)巨大的矩陣乘任務(wù)合理的分配到 GPU 的不同線程上。

• Thread 級優(yōu)化：在 Thread 這個(gè)維度，我們能做什么優(yōu)化？

• Warp 級優(yōu)化：在 Warp 這個(gè)維度，我們能做什么優(yōu)化？

• Block 級優(yōu)化：在 Block 這個(gè)維度，我們能做什么優(yōu)化？

• Epilogue：尾聲。

Goals

首先明確一下本文的目標(biāo)是：

? 實(shí)現(xiàn)一個(gè)比 cublas 更快的形狀較大的正方形乘正方形的 FP32 矩陣乘。

? 從理論角度與硬件規(guī)格能夠簡單的推導(dǎo)矩陣分塊與排布的方法。

? 可以大致清楚各個(gè)優(yōu)化技術(shù)效果的階段性的 benchmark。

? 如何使用 Nsight Compute 等性能分析工具分析潛在的性能瓶頸。

本文不含：

? 使用 Tensor Core 加速矩陣乘。（這也是為什么這篇文章叫傳統(tǒng) CUDA GEMM）

? 使用安培架構(gòu)新提出的 async memcpy。

? CUDA 語法知識。

? 匯編。（主要是現(xiàn)在并沒有官方支持匯編的操作，目前的匯編器幾乎都是逆向的產(chǎn)物，不是很穩(wěn)定。同時(shí)匯編帶來的好處如消除寄存器的 bank conflict nvcc 也在不斷的做相應(yīng)的改進(jìn)，因此就不介紹了）

開源地址：https://github.com/AyakaGEMM/Hands-on-GEMM

同時(shí)本文在相當(dāng)程度上參考了李少俠的 GEMM 優(yōu)化指南（寫得非常！非常！非常！不錯(cuò)），本文的優(yōu)勢在于補(bǔ)了階段性代碼和在某些少俠一筆帶過的地方做了一些擴(kuò)展。

Performance

為了讓大家更有動(dòng)力閱讀下去，這里先放出來性能效果！

測試平臺

? 系統(tǒng)：Arch Linux

? 驅(qū)動(dòng)：520.56.06

?CUDA：11.8

?GPU：Nvidia RTX 2080

測試結(jié)果

我們也可以注意到，在較大形狀上手寫的矩陣乘有著與 cublas 相近，甚至更優(yōu)的性能。

從這張圖我們可以看出，手寫的矩陣乘能夠達(dá)到硬件 95% 的峰值性能，效果還是很不錯(cuò)的。

樸素 GEMM 與前置知識

首先寫一個(gè)樸素矩陣乘。

#數(shù)組 A：M 行 K 列的行主序矩陣
#數(shù)組 B：K 行 N 列的行主序矩陣
#數(shù)組 C：M 行 N 列的行主序矩陣
# alpha：一個(gè)標(biāo)量
# beta：一個(gè)標(biāo)量
#計(jì)算方法：
#    c=alpha*A*B+beta*C;


__global__ void matrixMul(const float *A, const float *B, float *C,


intM,intN,intK,floatalpha,floatbeta)
{
inttx=blockIdx.x*blockDim.x+threadIdx.x;
intty=blockIdx.y*blockDim.y+threadIdx.y;
intbaseX=blockIdx.x*blockDim.x;
intbaseY=blockIdx.y*blockDim.y;
floatc=0;
if(tx
{
for(inti=0;i
{
c+=A[tx*K+i]*B[i*N+ty];
}
C[tx*N+ty]=beta*C[tx*N+ty]+alpha*c;//wemultiplyalphaheretoreducethealphacalnum.
}
}

這里 GPU 矩陣乘與 CPU 矩陣乘最大的區(qū)別就在于 GPU 可以為目標(biāo)矩陣 C 中每一個(gè)元素分配一個(gè) thread 進(jìn)行計(jì)算。這也是可以切實(shí)的感知到 GPU 多線程編程的一點(diǎn)。但這個(gè)矩陣乘的樸素實(shí)現(xiàn)會(huì)非常慢，而分析性能瓶頸中最常見的兩個(gè)指標(biāo)即是帶寬和延遲。這里借用 Nvidia 在 GTC 2018 上的分享來做說明。

這里以一個(gè)自動(dòng)扶梯作為例子來講解。

• 帶寬：指這個(gè)自動(dòng)扶梯每秒能夠運(yùn)送多少個(gè)人，以這張圖為例子，這個(gè)扶梯每秒能運(yùn) 0.5 個(gè)人，這就是這個(gè)自動(dòng)扶梯的帶寬。

• 延遲：指一個(gè)人踏上這個(gè)扶梯直到他被運(yùn)到頂所需要的時(shí)間。同樣以這張圖為例，這個(gè)扶梯需要 40s 的延遲。

那么回到指令上來，每一個(gè)指令都有對應(yīng)的延遲和帶寬，而以樸素矩陣乘為例，每一個(gè)乘法運(yùn)算需要讀兩次內(nèi)存和一次 FFMA，假如沒有其他額外的優(yōu)化（如循環(huán)展開與指令重排），相當(dāng)于是兩個(gè)級聯(lián)的自動(dòng)扶梯，一個(gè)負(fù)責(zé)運(yùn)送數(shù)據(jù)，一個(gè)負(fù)責(zé)做數(shù)學(xué)運(yùn)算。假設(shè)數(shù)據(jù)運(yùn)送扶梯的帶寬與延遲與圖中一致而不考慮 FFMA 的帶寬與延遲，那么一次 FFMA 需要等待 40s（扶梯延遲）+ (1/0.5)s（第一個(gè)數(shù)據(jù)到達(dá)后第二個(gè)數(shù)據(jù)到達(dá)的時(shí)間）才能拿到所需的數(shù)據(jù)，這與扶梯的帶寬 0.5s / 人的峰值性能相去甚遠(yuǎn)。那么此時(shí)這個(gè) kernel 就完全被延遲卡住了，而無法發(fā)揮出應(yīng)有的性能。

而對于帶寬部分，這里我們引用李少俠的帶寬分析：

對于 FP32 數(shù)據(jù)，如上圖所示，一個(gè) warp 一次做 32 次 FFMA，對應(yīng) 64 OP，需讀取 A 矩陣 1 個(gè)元素和 B 矩陣 32 個(gè)元素，共 132byte。 通過寄存器累加，且忽略 C 矩陣寫回開銷，那么計(jì)算訪存比為 64OP / 132byte = 0.48。雖然 dram 最小訪問單位為一個(gè) memory transaction，但考慮到 L1 cache 的存在也不會(huì)影響實(shí)際的計(jì)算訪存比。

通過 repo 中提供的 l2cache_bandwidth.cu 可測得 Titan V L2 cache 帶寬約 1.9TB/s，那么最樂觀的結(jié)果即使 L2 cache 100% 命中，此方案的理論上限也只有 1.9T * 0.48 = 0.912 tflops，遠(yuǎn)低于 14.9 tflops 的硬件算力。

由此我們可以看出，樸素的矩陣乘實(shí)現(xiàn)方法無論從延遲和帶寬上都無法滿足需要。

這里一個(gè) warp（即 32 個(gè)線程）是指 GPU 調(diào)度線程的粒度，可以簡單的理解為同一個(gè) warp 內(nèi)的線程總是同時(shí)運(yùn)行、同時(shí)休眠的。當(dāng)然這種說法并不完全準(zhǔn)確，畢竟還有 warp divergent 問題，感興趣的同學(xué)可以自行了解。但總之，思考 GPU 執(zhí)行時(shí)總是從 warp 的角度思考是非常合理的。那么對于一個(gè) warp 而言，我們可以根據(jù)李少俠的分析看出，就算我們假設(shè)延遲能夠被完全覆蓋，這種分配方案也并不能達(dá)到硬件的峰值性能。

這里我用自己的話總結(jié)一遍就是：在每個(gè)線程執(zhí)行指令設(shè)計(jì)時(shí)，需要盡可能的覆蓋掉每個(gè)指令的延遲；在性能分析時(shí)，則從帶寬角度分析矩陣分塊是否合理。

而在于延遲部分還有一頓免費(fèi)的午餐。在實(shí)際應(yīng)用中，編譯器會(huì)自動(dòng)的做一些優(yōu)化，如循環(huán)展開與指令重排等。例如展開循環(huán)后可以將多個(gè)讀取 A 矩陣的元素和讀取 B 矩陣的元素排在一起，使得取數(shù)據(jù)的自動(dòng)扶梯能夠一次多上幾個(gè)人，從而去覆蓋掉扶梯的延遲。而且 GPU 與 CPU 還有一個(gè)非常不同的地方在于，GPU 的線程切換代價(jià)非常低，因此可以在等待延遲的時(shí)候轉(zhuǎn)而去運(yùn)行其他線程從而達(dá)到延遲覆蓋的目的。還是以扶梯為例子，GPU 上有很多個(gè)扶梯，在等待第一個(gè)人到達(dá)扶梯末尾時(shí)，GPU 可以轉(zhuǎn)到第二個(gè)扶梯送幾個(gè)人上扶梯。理想情況下，在 GPU 送完第 N 個(gè)扶梯的人時(shí)，第一個(gè)扶梯的人剛好達(dá)到扶梯頂部，那么這個(gè)運(yùn)送的延遲就被覆蓋掉了。

Tiling

矩陣乘分塊是為了將一個(gè)大問題化解為小問題求解，這里 CPU 與 GPU 分塊的需求也不盡相同。CPU 是希望保持計(jì)算的局部性，可以充分利用 L1、L2 高速緩存來避免緩慢的內(nèi)存訪問。而 GPU 在此基礎(chǔ)之上還需要將一個(gè)大問題合理拆分到不同的 thread 上，使得其能夠充分利用 GPU 上的硬件資源。下面我將從局部性和合理拆分兩個(gè)方面講解如何做矩陣分塊。

局部性原理

局部性原理，在我的理解中便是為了能夠盡可能的使用高速緩存器內(nèi)的數(shù)據(jù)進(jìn)行運(yùn)算所提出的一個(gè)程序設(shè)計(jì)理念。由于高速緩存器往往十分昂貴（或者需要很高的功耗），因此空間都不大。由此我們需要盡可能的將一些重復(fù)訪存聚合起來，放到高速緩存器里面來加速數(shù)據(jù)訪問，或者在進(jìn)行訪存的時(shí)候盡可能連續(xù)訪存來使用 cache 加速訪存。我們先還是讓每一個(gè) thread 負(fù)責(zé)一個(gè)目標(biāo)矩陣元素的計(jì)算。雖然這種分配方式十分樸素、十分直接，同時(shí)各個(gè) thread 之間也沒有數(shù)據(jù)依賴關(guān)系，不需要做額外的同步之類的操作，但這種分配方式卻是十分訪存不友好的，因?yàn)槊恳粋€(gè) thread 都是直接與內(nèi)存做交互，而 GPU 的全局內(nèi)存訪問帶寬完全不足以匹配上它的計(jì)算速度。

同時(shí)我們注意到處于同一行的 thread 總是會(huì)同樣的讀取 A 矩陣的同一行數(shù)據(jù)；同一列的 thread 總是會(huì)讀取 B 矩陣的同一列數(shù)據(jù)。那么一個(gè)非常自然的想法則是對于每一個(gè) Block，我們將數(shù)據(jù)移動(dòng)到這個(gè) Block 共享的一塊高速存儲區(qū) shared memory 上，從而減少與全局內(nèi)存交互的次數(shù)。同時(shí)我們考慮到 shared memory 的容量有限，因此可以一次只取一部分 k，然后通過循環(huán)迭代的方式完成這個(gè) Block 所負(fù)責(zé)的矩陣乘區(qū)域。

值得一提的事，shared memory 雖然叫做 memory，但他卻有著非常高的訪存速度與極低的延遲。實(shí)際上，shared memory 可以被看作是一塊可以顯式控制的 L1 cache。從圖靈架構(gòu)開始，在硬件上 shared memory 與 GPU 上的 L1 cache 共享同一塊區(qū)域，同時(shí) shared memory 與 Load/Store 單元交互也是直連的（沒有中間商賺差價(jià)）。

在將一個(gè)大型矩陣乘劃分為一個(gè)個(gè)由 Block 負(fù)責(zé)的小型矩陣乘之后，我們接下來還需要把一個(gè) Block 負(fù)責(zé)的矩陣乘分配給 Block 內(nèi)部的 warp；分配到 warp 之后我們還需要把一個(gè) warp 負(fù)責(zé)的矩陣乘分配給 warp 內(nèi)部的 thread。經(jīng)過這么一步一步的劃分，我們便可以把一個(gè)巨大的矩陣乘任務(wù)高效的分配到各級速度不一的存儲器上，最終盡可能打滿硬件峰值性能，實(shí)現(xiàn)高效矩陣乘。有了前面劃分 Block 的經(jīng)驗(yàn)，我們也就可以依葫蘆畫瓢，實(shí)現(xiàn)大矩陣的拆分（Tiling），在此就不過多贅述了，最終整體流程圖如下。

當(dāng)然這只是一個(gè)較為粗糙的流程圖，例如每一個(gè) thread 負(fù)責(zé)的分塊也并不是圖中所示的連續(xù)一塊矩陣乘，我們也將在后續(xù)一步一步完善細(xì)節(jié)，但這種分解的框架卻是一種非常經(jīng)典的思路。

如何確定分塊大??？

在擁有分塊的基本理念之后，我們還有一個(gè)問題沒有解決。那便是每一個(gè) Block 該負(fù)責(zé)多大的矩陣乘？每一個(gè) thread 又應(yīng)該負(fù)責(zé)多大的矩陣乘？為了讓文字變得清晰起來，我們定義每一個(gè) Block 負(fù)責(zé)的矩陣大小為，每次迭代的 k 維數(shù)據(jù)，每一個(gè) warp 負(fù)責(zé)的矩陣大小為，每一個(gè) thread 負(fù)責(zé)的大小為。其中這些符號都在上圖出現(xiàn)過，可以自行對照一下。

這里我們同樣引用李少俠的計(jì)算訪存分析：

假設(shè)我們不考慮 shared memory 的訪存代價(jià)（因?yàn)榭梢宰龅?/span>覆蓋掉shared memory 的訪存延遲，而且其帶寬能夠滿足 FFMA 單元的計(jì)算速度），只考慮全局內(nèi)存的訪問，可以看到選擇在K 上縮水（即不把整個(gè) K 維度都放到 shared memory 里）還是比較合理的，因?yàn)?/span>的大小其實(shí)并不影響計(jì)算訪存比。而對計(jì)算訪存比有決定性影響的是每一個(gè) Block 計(jì)算的大小。如果取為 64，帶入 RTX 2080 的數(shù)據(jù)，可以得到 10.1 Tflops / 16 = 631.25 GB/s。即內(nèi)存訪問帶寬達(dá)到 631.25 GB/s 就能避免內(nèi)存訪問瓶頸了。同樣，我們?nèi)?L2 命中率為 20%（還是比較好達(dá)到的），加權(quán)內(nèi)存訪問帶寬為：，即可避免內(nèi)存訪問瓶頸。

那是否我們只要取分塊大小為 64x64 就萬事大吉了呢？也不盡然。我們前面只分析了帶寬，而在延遲無法被覆蓋的情況下，整個(gè) kernel 性能也不會(huì)太好。而更大的分塊意味著每一個(gè) thread 會(huì)計(jì)算更多的數(shù)據(jù)，可以使用一些手段實(shí)現(xiàn)更優(yōu)的延遲覆蓋。這一點(diǎn)會(huì)在后面討論如何具體實(shí)現(xiàn)，大致思想也是局部性的原理，只不過這次是將數(shù)據(jù)從 shared memory 保存到寄存器，從而實(shí)現(xiàn)使用更高速的緩存計(jì)算的目的。

那是否我們?nèi)》謮K越大越好呢？那也不一定。更大的分塊使用了更多的寄存器，從而使得同一個(gè) SM 能夠同時(shí)承載的線程數(shù)變少，這里 Nvidia 將之稱為 Occupancy。如前文所述，當(dāng)一個(gè) warp 被卡住時(shí)，GPU 可以切換到另一個(gè) warp 執(zhí)行指令，Occupancy 越低，可供 GPU 切換的線程就越少。

而 Occupancy 也是和硬件強(qiáng)相關(guān)的。一個(gè) GPU 由多個(gè) SM 構(gòu)成，每一個(gè) SM 擁有有限的寄存器數(shù)量、 shared memory 和最大可調(diào)度線程數(shù)量。而 Occupancy 是指每個(gè) SM 能夠同時(shí)調(diào)度的線程數(shù)量除以一個(gè) SM 的最大可調(diào)度線程數(shù)量。關(guān)于 Occupancy 的計(jì)算我們可以通過在編譯時(shí)添加 --ptxas-options=-v 參數(shù)，使編譯器在編譯時(shí)輸出每個(gè) kernel 所花費(fèi)的寄存器數(shù)量和 shared memory，然后通過隨 cuda 提供的一個(gè) excel 表格進(jìn)行計(jì)算。（盡管這個(gè) Excel 已經(jīng) deprecated 了，但他用起來確實(shí)挺方便的。）

例如我們每個(gè) thread 需要 128 個(gè)寄存器，2048 bytes 的 shared memory，那么由于 RTX 2080 每個(gè) SM 只有 65536 個(gè)寄存器，因此每個(gè) SM 最多只能同時(shí)跑 512 個(gè) threads。又因?yàn)槊總€(gè) SM 最多能夠承載 1024 個(gè) threads，所以此時(shí) Occupancy 為。

值得一提的是，雖然較高的 Occupancy 使得在一個(gè)線程卡住時(shí)，SM 能夠馬上切換到別的線程，通過將其他線程需要執(zhí)行的指令填充到流水線中從而達(dá)到覆蓋延遲的目的，但這并不代表高性能。例如，如果每一個(gè)線程本身就能夠通過更多的寄存器占用從而達(dá)到延遲覆蓋的目的，自然也就不需要 SM 來做這件事了，反倒是如果無腦的去提高 Occupancy 使得一些 thread 內(nèi)的延遲甚至都無法被 SM 通過切換執(zhí)行線程的方式覆蓋，那屬實(shí)是得不償失了。

因此，我們能夠做的就是在有一定理論分析的情況下確定好一些矩陣的分塊大小的方案，然后要不就是經(jīng)驗(yàn)性的去選擇最終用哪個(gè)分塊，要不就是跑一個(gè) profile 來直接得到最快的分塊。這里由于已經(jīng)有非常多的先例證明了 128x128x8 是一個(gè)較優(yōu)的選擇，因此本文則遵從這個(gè)分塊方案。那么，目前我們能夠確定的分塊如下表。

當(dāng)然有些同學(xué)可能會(huì)問，既然最終還是需要用跑 profile 的方式來確定最優(yōu)分塊，那理論分析還有什么意義呢？答案就是如果提前通過理論分析，那么就能夠在一定程度上縮小需要跑 profile 的分塊數(shù)量。用算法上的語言來講就是如果我們將需要搜索的所有分塊作為搜索空間，那么理論分析便是搜索算法中的 A* 算法，你掌握了越多的理論分析知識那么這個(gè)搜索過程就會(huì)越高效。同時(shí)對 CUDA 底層越了解，在同一個(gè)分塊策略下，你更容易寫出能達(dá)到理論性能的 kernel。

Thread 級優(yōu)化

對于一個(gè) thread 能做的優(yōu)化其實(shí)并不多，因?yàn)?GPU 是以一個(gè) warp（即 32 個(gè) thread）進(jìn)行調(diào)度的，所以許多基于單線程的優(yōu)化，如訪存優(yōu)化，其實(shí)并不能直接套到 GPU 上。而為數(shù)不多值得一提的優(yōu)化手段便是單個(gè)線程在計(jì)算時(shí)應(yīng)該采用向量內(nèi)積還是向量外積以及 double buffer。但實(shí)質(zhì)上向量外積嚴(yán)格意義上也不能算作是一個(gè)優(yōu)化，因?yàn)檫@一步編譯器就能在編譯階段幫忙做了。之所以提一句是還是為了給 double buffer 做鋪墊，即我們應(yīng)該怎么預(yù)取數(shù)據(jù)。

首先我們?nèi)×?128x128 的分塊策略，一個(gè) Block 有 256 個(gè)線程，那么每個(gè)線程需要負(fù)責(zé)一個(gè) 8x8 的矩陣乘運(yùn)算。而一個(gè)線程完成一個(gè)小型矩陣乘有兩種實(shí)現(xiàn)方法。

向量內(nèi)積

向量內(nèi)積的實(shí)現(xiàn)方法如圖所示，即將 A 矩陣拆分為多個(gè)行向量、B 矩陣拆分為多個(gè)列向量，這些向量通過向量內(nèi)積的方法求得最終答案。

用代碼描述如下：

M=N=K=8;
floata[M*N];
floatb[N*K];
floatc[M*N];
foriinrange(M):
forjinrange(N):
forkinrange(K):
                  c[i*N+j]+=a[i*K+k]*b[k*N+j];

向量外積

向量外積的實(shí)現(xiàn)方法如圖所示，即將 A 矩陣拆分為多個(gè)列向量、B 矩陣拆分為多個(gè)行向量，這些向量通過向量外積的方法求得最終答案。

用代碼描述如下：

M=N=K=8;
floata[M*N];
floatb[N*K];
floatc[M*N];
forkinrange(K):
foriinrange(M):
forjinrange(N):
                   c[i*N+j]+=a[i*K+k]*b[k*N+j];

可以看到，向量內(nèi)積和向量外積的區(qū)別在代碼上僅僅體現(xiàn)在循環(huán)方式上。

為何我們需要關(guān)心這個(gè)？

有做過 CPU 矩陣乘優(yōu)化的同學(xué)可能知道，僅僅調(diào)整循環(huán)順序就已經(jīng)能夠帶來顯著的性能差異了。有許多分析都是從局部性的角度進(jìn)行分析的。即使用向量外積的方案可以利用到循環(huán)遍歷的局部性，將一些重復(fù)訪存使用寄存器緩存而避免無意義訪存。例如我們補(bǔ)充一下采用向量外積方案關(guān)于寄存器的細(xì)節(jié)。

floata[M*N];
floatb[N*K];
floatc[M*N];
forkinrange(K):
regB[0:N]=b[k*N:(k+1)*N]
foriinrange(M):
regA=a[i*K+k];
forjinrange(N):
c[i*N+j]+=regA*regB[j];

其中 regA 和 regB 均為寄存器。其中我們不難發(fā)現(xiàn)，對于每一次循環(huán) j ，使用的都是完全相同的 A 矩陣?yán)锏脑?，因此可以用一個(gè)寄存器來緩存該值；對于每一次循環(huán) k，使用的都是完全相同的一行 B 矩陣中的值，因此我們可以用 N 個(gè)寄存器緩存該值。于是將原本次訪存（底下兩層循環(huán)需要訪問一次 A 矩陣和一次 B 矩陣），通過使用個(gè)寄存器緩存（B 使用 N 個(gè)，A 使用一個(gè)），優(yōu)化為 N+M 次訪存。同時(shí)我們也注意到， M 和 N 越大的情況下，提升效果越發(fā)顯著，這也是為什么我們希望每一個(gè)線程負(fù)責(zé)的分塊大一點(diǎn)比較好。但同時(shí) M 和 N 越大，每一個(gè)線程多使用的寄存器就越多，而在 GPU 的語境下，更高的寄存器占用意味著更低的 Occupancy。因此當(dāng) M 和 N 大到 shared memory 帶寬不是性能瓶頸即可。更詳細(xì)的分析可以看李少俠的分析。

而我則從循環(huán)展開的角度解釋一下為什么我們需要了解這個(gè)優(yōu)化方案，同時(shí)解釋一下為什么該優(yōu)化方案在 GPU 上并不如 CPU 上那么有效。從循環(huán)展開的角度來看，第二種循環(huán)體構(gòu)造與第一種循環(huán)最大的區(qū)別就在于它能在不展開 k 的情況下通過展開 m 和 n 處的循環(huán)就能自動(dòng)的識別到重復(fù)訪存，并使用相應(yīng)的寄存器來避免重復(fù)訪存。例如我們假定，那么展開 m 和 n 處循環(huán)的結(jié)果如下。

M=N=2;
floata[M*N];
floatb[N*K];
floatc[M*N];
forkinrange(K):
c[0*N+0]+=a[0*K+k]*b[k*N+0]
c[0*N+1]+=a[0*K+k]*b[k*N+1]
c[1*N+0]+=a[1*K+k]*b[k*N+0]
c[1*N+1]+=a[1*K+k]*b[k*N+1]

只要是稍微現(xiàn)代一點(diǎn)的編譯器，都能一眼看出這四條指令的 8 次訪存，有 4 次是可以合并的。同時(shí)現(xiàn)代一點(diǎn)的編譯器也能在一定程度上根據(jù)生成的匯編交叉排列計(jì)算和訪存達(dá)到延遲覆蓋的目的。而向量內(nèi)積的方案需要把整個(gè) k 維度展開才能看到這些潛在的訪存合并機(jī)會(huì)。在 CPU 矩陣乘的語境下，一般計(jì)算 kernel 的都比較大（好幾百），而都很小（一般取 6x16，根據(jù)架構(gòu)來做具體確定），寄存器數(shù)量又非常少，因此基本上無法在 K 維上將循環(huán)完全展開并做優(yōu)化。因?yàn)檎归_一個(gè)超長的循環(huán)不僅會(huì)帶來額外的寄存器占用、優(yōu)化難度，還會(huì)帶來更多的匯編指令，使得最終的二進(jìn)制文件臃腫不堪。但在 GPU 上，情況卻恰恰相反。對于已知循環(huán)次數(shù)的小循環(huán)，即便你沒有指定 #pragma unroll，nvcc 也會(huì)自動(dòng)的展開這些循環(huán)。而對于一個(gè) thread 所負(fù)責(zé)的小型矩陣乘，這三層循環(huán)的值均為 8，符合 nvcc 自動(dòng)展開循環(huán)的條件。而在展開完成后，nvcc 會(huì)對所有的訪存以及計(jì)算指令重排得到一個(gè)不錯(cuò)的匯編指令排列。

那么這就引出了下一個(gè)問題，我們?yōu)楹芜€需要關(guān)心他究竟是向量內(nèi)積還是向量外積？

答案就是 double buffer。如果我們能夠提前知道一個(gè)循環(huán)需要什么數(shù)據(jù)，我們就能提前預(yù)取該循環(huán)第一次所需的數(shù)據(jù)，同時(shí)在該循環(huán)進(jìn)行運(yùn)算的時(shí)候預(yù)取下一次計(jì)算所需的數(shù)據(jù)。而顯然這在向量內(nèi)積的情況下是無法做到的。同時(shí)由于 double buffer 需要額外的寄存器保存從 global memory 轉(zhuǎn)移到 shared memory 的數(shù)據(jù)，所以當(dāng)一開始循環(huán)展開使用的寄存器過多時(shí)，盡管后續(xù)能優(yōu)化到較少的寄存器，但編譯器依然無法正確的在限定寄存器數(shù)量下實(shí)現(xiàn) double buffer。這一點(diǎn)在優(yōu)化 sgemm 的時(shí)候并不是那么重要（因?yàn)槎嗍褂靡稽c(diǎn)寄存器也就從每個(gè) SM 跑兩個(gè) block 變?yōu)橐粋€(gè) block），但是在優(yōu)化 int8 矩陣乘時(shí)需要額外的關(guān)注（因?yàn)楸旧硭椭荒茉谝粋€(gè) SM 上跑一個(gè) block，如果實(shí)現(xiàn)不得當(dāng)將會(huì)完全失去 double buffer）。

那么此時(shí)樸素的利用到向量外積和 shared memory 的代碼：https://github.com/AyakaGEMM/Hands-on-GEMM/blob/main/src/cuda/shared_mem_gemm.cu

Double Buffer

由于 GPU 沒有 prefetch 這種指令，同時(shí)我們又有 shared memory 這種可編程的 L1 cache，因此需要手動(dòng)實(shí)現(xiàn) prefetch 功能，而在 GPU 語境下一般被稱作 double buffer。double buffer 的好處自不必多說，即它可以實(shí)現(xiàn)數(shù)據(jù)讀取與計(jì)算在時(shí)間上重疊，利用 FFMA 單元與 Load/Store 單元可以并行執(zhí)行指令的特點(diǎn)，達(dá)到覆蓋延遲的目的。而盡管 GPU 可以在一個(gè) warp 有延遲的情況下，通過切換去運(yùn)行另一個(gè) warp 達(dá)到延遲覆蓋到目的，但由于可供 warp 調(diào)度器能切換到線程數(shù)量的限制，過于長的延遲并不能通過這種方式覆蓋掉。這里引用一下李少俠更詳細(xì)的分析：

若每 SM 有 4 個(gè)調(diào)度器，若每個(gè)調(diào)度器只有 4 個(gè)可調(diào)度 warp，當(dāng)指令平均間隔超過 4 cycle 后就無法靠 warp 調(diào)度掩蓋延遲了?？紤]到 GEMM 中涉及 smem 讀寫的過程需要同步 thread block，進(jìn)一步限制了 warp 調(diào)度空間，所以很難靠 warp 并行掩蓋延遲。

而本文最終實(shí)現(xiàn)的 kernel occupancy 只有 50%，即每個(gè) SM 只能調(diào)度 512 個(gè) threads（16 warps），加上圖靈架構(gòu)每 SM 有 4 個(gè) warp 調(diào)度器，最終結(jié)果與李少俠分析的一致。因此 double buffer 從指令角度提供的延遲覆蓋方法最終還是會(huì)有效的。

但值得一提的是，在你自己動(dòng)手實(shí)踐時(shí)，盡可能的考慮在其他優(yōu)化已經(jīng)加無可加的情況下再加入這個(gè)優(yōu)化。這是由于這個(gè)優(yōu)化會(huì)大幅修改數(shù)據(jù)讀取部分的代碼，而且還會(huì)產(chǎn)生重復(fù)代碼，不利于代碼維護(hù)。同時(shí)在我自己的實(shí)踐中發(fā)現(xiàn)，如果在一開始 kernel 寫的比較垃圾，加了 double buffer 也沒有什么卵用，還會(huì)讓后續(xù)的優(yōu)化不太好加上去。當(dāng)然，這只是我的個(gè)人建議，如果你想實(shí)際看看 double buffer 的效果也可以一開始就加上去。

首先我們看一下每個(gè) thread 的運(yùn)行流程。

那么能實(shí)現(xiàn) double buffer 的機(jī)會(huì)有兩個(gè)地方：Global Memory to Shared Memory 與 Shared Memory to Register。即在每一次 FFMA 開始之前我們讀取 Global Memory 的數(shù)據(jù)到寄存器中，在 FFMA 之后將該寄存器中的值寫到 shared memory 中。由于在讀取數(shù)據(jù)后 load from shared memory 以及 FFMA 兩個(gè)流程中我們并不依賴于該寄存器中的數(shù)據(jù)，因此可以覆蓋 Global Memory 的讀取延遲。而同時(shí)在計(jì)算每一次 FFMA 之前，我們可以用寄存器提前取下一次 FFMA 需要的數(shù)據(jù)，也就能做到覆蓋 shared memory 的延遲。

大概就是這樣！我們在每一次運(yùn)算之前提前將第一次循環(huán)所需的數(shù)據(jù)移動(dòng)到寄存器中，這樣我們就可以實(shí)現(xiàn)數(shù)據(jù)運(yùn)算和數(shù)據(jù)存取指令級并行的功能了。

Warp 級優(yōu)化

在做了不少鋪墊之后，接下來的優(yōu)化終于是可以帶來一些看得見的性能提升的了。首先回顧我們之前的代碼，可以看到每個(gè) thread 負(fù)責(zé)的部分完全沒有考慮到它們之間可能的協(xié)作關(guān)系，即同一個(gè) warp 內(nèi)的 thread 此時(shí)在同一塊硬件上同時(shí)執(zhí)行——它們共享同一個(gè) register file，這表明它們可以通過寄存器快速共享數(shù)據(jù)（即 shared memory 的 broadcast 機(jī)制）；它們會(huì)同時(shí)訪存，這表明如何安排每一個(gè) warp 內(nèi)的 thread 訪存是至關(guān)重要的。

Warp Tiling

已知我們指定一個(gè) Block 計(jì)算 128x128 的矩陣，一個(gè) Block 有 8 個(gè) warp，一個(gè) warp 有 32 個(gè) thread，每個(gè) thread 需要負(fù)責(zé) 8x8 的小型矩陣乘，那么我們沿用李少俠的定義：

一個(gè) warp 由個(gè)線程組成，可以是?，我們把這些線程對應(yīng)的 thread tile 拼在一起的區(qū)域稱為 warp tile，尺寸為，如下圖所示。

這里的圖給的是的排列方式。由于同一個(gè) warp 在訪問 shared memory 時(shí)有 broadcast 機(jī)制（即同一個(gè) warp 在訪問同一個(gè)內(nèi)存地址內(nèi)的值時(shí)只會(huì)實(shí)質(zhì)發(fā)生一次數(shù)據(jù)讀?。虼诉@一個(gè) warp 計(jì)算時(shí)只會(huì)實(shí)際讀取個(gè) float。與之相對的，這個(gè) warp 會(huì)進(jìn)行次 FFMA。不難看出，在固定為 32 的情況下，與越相近，計(jì)算訪存比就越大，因此取最為合適。

而在確定了 warp tiling 后，如何讀取和存儲數(shù)據(jù)的細(xì)節(jié)還需要細(xì)扣，接下來我將會(huì)按照 GPU 的硬件特性講解讀寫數(shù)據(jù)的細(xì)節(jié)。但這一部分的大致思路基本已經(jīng)介紹完畢了，動(dòng)手能力強(qiáng)的同學(xué)現(xiàn)在就可以自己試試如何寫一個(gè)高效矩陣乘了！

向量化訪存

向量化訪存即是一條指令同時(shí)請求多個(gè) float 數(shù)據(jù)，目前 CUDA 最高支持 128 bit 的向量化訪存，即一條指令請求 4 個(gè) float 數(shù)據(jù)。向量化訪問主要的好處在于可以用更少的指令讀取更多的數(shù)據(jù)。由于在訪問全局內(nèi)存時(shí)是以 32 Byte 為粒度進(jìn)行訪問的，因此如果同一個(gè) warp 內(nèi)的 thread 請求了一段連續(xù)內(nèi)存的數(shù)據(jù)，每一個(gè) thread 都請求兩次 4 Byte 的數(shù)據(jù)（小于 GPU 全局訪存的最小單位），那么 GPU 會(huì)在硬件處將 64 次數(shù)據(jù)請求按照 32 Byte 進(jìn)行合并，最終形成 8 次 32 Byte 內(nèi)存訪問。

而如果每一個(gè) thread 請求 8 Byte 數(shù)據(jù)，那么 GPU 會(huì)在硬件處同樣將 32 次數(shù)據(jù)請求按照 32 Byte 進(jìn)行合并，最終形成也形成 8 次 32 Byte 內(nèi)存訪問。

那么我們可以看出，對于訪問同一數(shù)據(jù)量的數(shù)據(jù)，請求的指令越多，GPU 的聚合訪問的壓力就會(huì)越大。在極端情況下，盡管帶寬足夠，但大量的訪存請求會(huì)塞滿訪問隊(duì)列導(dǎo)致 stall。這在 Nsight Compute 中顯示為 MIO Throttle 和 LG Throttle，即對應(yīng) shared memory 和 global memory。因此采用向量化訪存能在一定程度上緩解 GPU 硬件層面的聚合訪存壓力（因?yàn)槲覀冿@式的用指令告訴 GPU 某些數(shù)據(jù)請求不需要聚合，直接用一個(gè) sector 來處理就好了）。

但使用向量化訪存——即用 float4 讀寫數(shù)據(jù)——也不是完美的。它的一個(gè)嚴(yán)重缺陷在于使用 float4 訪存要求請求的數(shù)據(jù)地址要按照 float4 對齊，因此當(dāng) M、N、K 不為 4 的倍數(shù)時(shí)將會(huì)報(bào) missaligned address 錯(cuò)誤（因?yàn)榈诙虚_始就不能按照 float4 對齊了）。

這么干對輸入矩陣形狀有一定要求，寫出來的矩陣乘沒有特別好的通用性。同時(shí) sgemm 受聚合訪存的影響也并不是那么大，因此在實(shí)操中往往并不會(huì)選擇使用 float4 讀寫全局內(nèi)存，而只會(huì)使用 float4 讀寫 shared memory。但由于我一開始學(xué) CUDA 的時(shí)候?qū)@一塊理解也不深，然后發(fā)現(xiàn)許多人（李少俠除外）都很暴力的直接用 float4 讀寫全局內(nèi)存，于是我也用了 float4 讀寫全局內(nèi)存。

而我們這里對比李少俠的 kernel profile 和我們最終的成品發(fā)現(xiàn)，在 global memory 讀取處是否使用向量化讀取其實(shí)并不會(huì)對性能有多少影響?？梢钥吹阶罱K profile 出來的 Stall LG Throttle 和 Stall MIO Throttle 占比都不高。

上圖為李少俠的 kernel 下圖為我最終寫的 kernel。這兩個(gè) kernel 在數(shù)據(jù)讀取方面的區(qū)別就是李少俠是以 4B 為單位訪存的，而我是以 16B 為單位做訪存的。這進(jìn)一步印證了 sgemm 其實(shí)并不是非常關(guān)心讀取 global memory 時(shí)是以怎樣的粒度讀取的。而向量化訪存對于 shared memory 的影響就留給讀者自行驗(yàn)證了。同時(shí)值得注意的是，在把數(shù)據(jù)讀取方式從向量化訪存修改為一個(gè)一個(gè)訪存時(shí)需要注意 bank conflict 的問題。因?yàn)橐粋€(gè) warp 在執(zhí)行 128-bit load 和 32-bit load 時(shí)的調(diào)度并不相同（這點(diǎn)會(huì)在后面提到）。

還有一個(gè)值得注意的是在 Global Memory 訪存時(shí)，并不能直接將原先的向量化存取代碼直接改成一個(gè)一個(gè)的讀取。因?yàn)檫@樣訪存從原來一個(gè) warp 并行訪問一段連續(xù)的內(nèi)存變成一個(gè) warp 分成四次訪問不連續(xù)的內(nèi)存。雖然有 L2 cache 來平滑這種不規(guī)則的訪存，但最終會(huì)帶來 10% 左右的性能下降。代碼如下：

//OriginalCode
preA=*reinterpret_cast<constfloat4*>(baseA+i+rowA*K+colA);


//ModifiedCode
preA.x=baseA[rowA*K+i+colA];
preA.y=baseA[rowA*K+i+colA+1];
preA.z=baseA[rowA*K+i+colA+2];
preA.w = baseA[rowA * K + i + colA + 3];

可以看到這種簡單的更改其實(shí)并不可取，更優(yōu)的寫法是每一條指令都是在 warp 視圖下的連續(xù)訪存。

Global Memory

前面提到 GPU 訪存時(shí)以 32 Byte 為粒度進(jìn)行訪問的，那么一個(gè) 32 Byte 訪問被稱為一個(gè) sector。那么值得注意的就是在搬運(yùn)數(shù)據(jù)時(shí)，盡可能的讓同一個(gè) warp 搬運(yùn)同一行的數(shù)據(jù)來避免使用額外的 sector（本文采用現(xiàn)代的行主序來存儲矩陣）。

這里借用一下 Nvidia 的圖。如果同一個(gè) warp 內(nèi)的 thread 都訪問每一行的開頭，那么如果一行超過 8 個(gè) float，那么每一個(gè) thread 都需要一個(gè) sector 去請求它們需要的數(shù)據(jù)，這就造成了 sector 浪費(fèi)。而實(shí)際中每一行的元素往往都會(huì)大于 8 個(gè) float，因此會(huì)有非常大的性能損失。下圖為一個(gè) warp 在拷貝時(shí)，每個(gè) thread 之間間隔的大小，單位為 float?？梢钥吹皆陂g隔為 2 時(shí)就已經(jīng)有一半的性能損失了，這很不好。

因此我們采用下圖所示的訪問方式。即盡可能的讓一個(gè) warp 中的 thread 連續(xù)的讀取 Global Memory 中的元素。

Shared Memory

前文已經(jīng)講過，shared memory 在圖靈架構(gòu)之后可以完全被看作是 L1 cache。而在此基礎(chǔ)之上，shared memory 的訪問粒度是 32 bit 也就是 4 Byte，剛好是一個(gè) float 數(shù)據(jù)的大小。而后 shared memory 按照 4 Byte 連續(xù)的劃分為一個(gè)個(gè) bank。對于 bank 可以簡單的理解為雙通道內(nèi)存中通道的概念，即在不同的 bank 中的數(shù)據(jù)可以并行訪問，同一個(gè) bank 內(nèi)不同地址的數(shù)據(jù)只能串行訪問。在 Compute Capability 5.x 及之后的卡上，shared memory 具有 32 個(gè) bank，剛好是一個(gè) warp 中線程的數(shù)量。而如果同一個(gè) warp 中不同 thread 均只訪問 4 Byte 數(shù)據(jù)且希望同時(shí)訪問同一個(gè) bank 的數(shù)據(jù)將會(huì)有兩種結(jié)果。（對于每一個(gè) thread 訪問更多數(shù)據(jù)的行為將在后面提到）

1. 兩個(gè)或多個(gè) thread 訪問的剛好是同一個(gè)地址內(nèi)的數(shù)據(jù)，那么此時(shí)將會(huì)觸發(fā) broadcast 機(jī)制，即實(shí)際只讀取一次數(shù)據(jù)，而后廣播到這些 thread 中。

2. 兩個(gè)或多個(gè) thread 訪問的是同一個(gè) bank 內(nèi)的數(shù)據(jù)，那么此時(shí)這些 thread 的訪問將會(huì)被強(qiáng)制安排為串行執(zhí)行。這種訪問情況被稱為 bank conflict。

這里給出 cuda programming guide 的兩張圖來直觀的體現(xiàn) broadcast 和 bank conflict。

這張圖表示同一個(gè) warp 中的 thread 按不間隔、隔一個(gè)、隔兩個(gè) bank 對 shared memory 訪問。中間的訪問每兩個(gè) thread 都會(huì)發(fā)生一次 bank conflict，而其他兩種訪問都不會(huì)發(fā)生 bank conflict。值得注意的一點(diǎn)是這張圖最右側(cè)的圖的訪問方式剛好可以達(dá)到每一個(gè) thread 都訪問了不同的 bank 的效果。

同時(shí)考慮到 shared memory 是按照 bank 來訪問的，且與 Load/Store 單元直連，并沒有中間商賺差價(jià)，所以對于 shared memory 的訪存并不講究連續(xù)訪存，而只需要考慮是否有 bank conflict 就足夠了。因此理論上最左和最右兩列圖的訪問性能是一樣的，這與訪問全局內(nèi)存有一點(diǎn)區(qū)別。同理，每一個(gè) warp 連續(xù)的多次訪存也并不要求連續(xù)訪存，而在拷貝數(shù)據(jù)到 shared memory 時(shí)對 A 矩陣做矩陣轉(zhuǎn)置的目的是為了向量化訪存，而不是為了連續(xù)訪存。

這張圖則展示了 broadcast 機(jī)制，沒啥好說的。

128-bit conflict-free store

而前文中提到，我們使用 float4 來做數(shù)據(jù)傳輸來緩解 GPU 聚合訪問的壓力，使得每一個(gè)指令都更加高效。而又因?yàn)榍拔乃?，每個(gè)線程需要使用向量外積的方法計(jì)算矩陣乘，因此我們需要在 A 矩陣轉(zhuǎn)存到 shared memory 時(shí)做一次轉(zhuǎn)置。

但細(xì)心的同學(xué)可能注意到，如果就這么平鋪直敘的做轉(zhuǎn)置那么將會(huì)發(fā)生非常嚴(yán)重的 bank conflict，因?yàn)橐粋€(gè) warp 內(nèi)的奇數(shù) thread 和偶數(shù) thread 使用同一個(gè) bank。那么此時(shí)解決 bank conflict 的方法有兩種，第一種便是將 shared memory 的 k 維度縮小，然后直接把奇數(shù) thread 所取的數(shù)據(jù)直接并到 M 維上，就不會(huì)有 bank conflict 的問題了。這種方法通過 index 變換，直接就能避免 bank conflict，非常巧妙，而我當(dāng)時(shí)沒有想到，就沒有用這種方法。值得注意的是，盡管圖是按行隔開的，但那只是為了表示數(shù)據(jù)是如何在一個(gè) thread 里保存的，實(shí)際寫到 shared memory 中是以一個(gè) float 為單位，按列主序存儲到 shared memory 中。

而第二種方法就非常簡單粗暴了，直接往 lda 上加 4，然后就不會(huì)有 bank conflict 了。當(dāng)然這種方法的弊端也是有的，那就是會(huì)造成一部分 shared memory 的浪費(fèi)。但對 sgemm 來說倒也還好， shared memory 的占用也不是導(dǎo)致 Occupancy 降低的原因，所以我就用了這個(gè)方法。

128-bit conflict-free load

而我們把數(shù)據(jù)存儲到 shared memory 之后，下一步便是考慮如何在沒有 bank conflict 的情況下將數(shù)據(jù)讀取出來。在本文中，我們?nèi)?/span>為 8，在采用向量化存取時(shí)，直接按照 Warp Tiling 采用樸素的存取方法就能在沒有 bank conflict 的情況下把數(shù)據(jù)讀出來了。

當(dāng)然有的同學(xué)可能會(huì)問：既然訪存是按照一個(gè) warp 為單位進(jìn)行的，而圖中明顯讀取 B 矩陣時(shí)，t16 會(huì)和 t0 發(fā)生 bank conflict，那為什么又說不會(huì)有 bank conflict 呢？那么答案就是在做 128-bit 訪存時(shí)，warp 并不是同時(shí)讀取數(shù)據(jù)的。這里還是借用 Nvidia 在 GTC 2018 上的分享來做說明。

當(dāng) warp 中每個(gè) thread 只讀取 4B 或更小數(shù)據(jù)時(shí)，warp 才是同時(shí)讀取的。而本文中采用 128-bit 也就是 16B 讀取，那么一個(gè) warp 會(huì)分成 4 次操作讀取，每次操作只有 1/4 warp 工作。那么只要同一次操作內(nèi)的 thread 沒有發(fā)生 bank conflict，那么就沒有 bank conflict。而上圖中 t0-t7 同時(shí)操作，它們之間并沒有 bank conflict，后面的 thread 依此類推，那么也就不會(huì)有 bank conflict。那么樸素的 warp tiling 實(shí)現(xiàn)代碼在這：https://github.com/AyakaGEMM/Hands-on-GEMM/blob/main/src/cuda/warp_tile_gemm.cu

而李少俠在代碼中采用了一種更高級的排布方式，即 z 字排布。與之相對應(yīng)的，他將一個(gè) thread 負(fù)責(zé)的小型矩陣乘拆分成四個(gè)更小的矩陣乘。同時(shí)這個(gè)拆分雖然是在地址上做的拆分，但在運(yùn)算中依然可以看作是一個(gè)整體，即運(yùn)算部分不用更改任何代碼而只需要在 index 上做一些變換即可。而他這么做的理由是為了更快的 broadcast。但說實(shí)話，我不是很理解，也沒搜到為什么這樣能有更快的 broadcast 性能。（而且我這么試了一下，發(fā)現(xiàn)確實(shí)是快了，這實(shí)在是太神奇了，歡迎大家提供一些看法。）

這里我們跑一個(gè) profile 發(fā)現(xiàn)，確實(shí)是沒有 bank conflict 的，挺好。代碼在這：

https://github.com/AyakaGEMM/Hands-on-GEMM/blob/main/src/cuda/z_thread_map_gemm.cu

Block 級優(yōu)化

Block 在 GPU 上基本等同于不同的 kernel 在 GPU 上運(yùn)行了，所以它們之間的聯(lián)系并不是特別強(qiáng)烈。而它們之間的相互關(guān)系在 GEMM 語境下基本就只有 wave 和 L2 cache（一個(gè) wave 里的 Block 共享這一塊 cache）了，良好的 Block Tiling 能提升相當(dāng)可觀的 L2 cache 命中率。

但這一部分屬于 sgemm 并不是特別關(guān)心的部分，因?yàn)楸旧?FFMA 單元算的就不是很快，所以 Block Tiling 隨便搞搞就能夠滿足 FFMA 單元的帶寬和延遲需求了。因此，這一節(jié)的內(nèi)容主要是為了有些有用到 tensor core 的同學(xué)提一個(gè)需要注意的性能提升點(diǎn)，其次就是有些同學(xué)可能會(huì)發(fā)現(xiàn)自己寫的 kernel 可能會(huì)比本文中的示例慢一點(diǎn)（大約 10% 左右），因此在此提一下在 sgemm 中應(yīng)該怎么隨便搞搞 Block Tiling。

Wave & L2 cache Hit Rate

首先明確一下 wave 的概念，即一個(gè) GPU 上能夠同時(shí)運(yùn)行的 Block 數(shù)量。關(guān)于 GPU 是如何決定一個(gè) wave 由哪些 Block 組成的我并沒有找到非常明確的文檔說明，但我一拍腦袋想，說不定就是樸素的按順序決定的，即 index 處于范圍內(nèi)的 Block 處在第一個(gè) wave 中，后面的 Block 依此類推。后面試了試好像的確是這樣劃分的。

在明確了 wave 的概念后，我們便可以對 L2 cache 命中率做一個(gè)簡單的分析了。我們指定代表一個(gè) wave 同時(shí)運(yùn)行的 Block 數(shù)量，假設(shè)一個(gè) wave 剛好能計(jì)算 C 矩陣的整數(shù)行，那么我們不難發(fā)現(xiàn)對于一個(gè) wave 而言，它需要從 Global Memory 中讀取個(gè) float。但由于有 L2 cache 的存在，假設(shè)一個(gè) wave 讀取的數(shù)據(jù)全能被 L2 cache 裝下，那么實(shí)際只讀取了的數(shù)據(jù)。最終 L2 cache 的命中率為：

即差距越大，L2 cache 的命中就越低。那么如果想要去優(yōu)化 L2 cache 命中，一個(gè)比較直接的想法就是盡可能把一個(gè) wave 的 Block 變成方的。但就算不搞，sgemm 也不在乎，因?yàn)槠鋵?shí)對性能來講并沒有什么區(qū)別，所以就沒搞。

SGEMM Block Tiling

而在 sgemm 的語境下，假設(shè)最壞的情況即一個(gè) wave 都不能覆蓋目標(biāo)矩陣 C 的一行，且 RTX 2080 有 46 個(gè) SM，一個(gè) SM 能跑兩個(gè) Block，此時(shí)

，

帶入上式可得，此時(shí) L2 cache 命中率大概是 49.4%。這里我們并沒有考慮訪問 C 矩陣的影響，在實(shí)踐中會(huì)把 L2 cache 的命中率拉低一點(diǎn)。但即便是如此，前文我們分析過只要 L2 cache 命中達(dá)到 20%，在帶寬上就不會(huì)造成性能瓶頸了。因此發(fā)現(xiàn)，就算我們采用樸素的 Block Tiling，Global Memory 訪問也不會(huì)成為訪存瓶頸。

但事實(shí)真的是這樣嗎？

細(xì)心的同學(xué)可能會(huì)發(fā)現(xiàn)，上圖所采用的 tiling 方式并不是直覺上的用 blockIdx.x 表示 Block 在 M 維上的位置，而是用 blockIdx.y 表示 Block 在 M 上的位置。而我們簡單調(diào)換一下代碼中 blockIdx.x 與 blockIdx.y 的順序，瞬間就有了 10% 左右的性能差距。目前網(wǎng)上并沒有針對這個(gè)現(xiàn)象的分析（因?yàn)閹缀跛腥硕际怯玫?col major 的 data layout，而且李少俠直接就在代碼里使用了更優(yōu)的 tiling 方式，所以沒有人遇到這個(gè)問題），因此我這里提出一點(diǎn)個(gè)人的猜想，如果猜的不對還請指正。

L2 cache

首先我們看一下這兩種 tiling 方式的區(qū)別在哪。最為直觀的區(qū)別就是當(dāng) N 或 M 足夠大時(shí)，采用上圖中的 tiling 方式的 wave 形狀是橫著的，而另一種 tiling 方式的 wave 形狀是豎著的，而這種豎著的形狀看起來就不是 cache 友好的訪存方式。

為什么這么說呢？因?yàn)槲也捎玫氖切兄餍虻姆绞酱鎯Φ木仃?，因此如果一個(gè) wave 的形狀是扁平的，那么每個(gè) Block 在每一次循環(huán)遍歷 B 矩陣時(shí)只會(huì)有次 cache miss。這是由于 L2 cache 的 cache line 大小為 128 bytes，因此當(dāng)數(shù)據(jù)從 Global Memory 中移動(dòng)到 L2 cache 后，許多 Block 就能直接從 L2 cache 中讀取數(shù)據(jù)了。然而如果一個(gè) wave 的形狀是狹長的。那么每個(gè) Block 在第一次訪問 A 矩陣的每一行時(shí)都會(huì)有 cache miss 的情況出現(xiàn)，即產(chǎn)生次 cache miss，而后 31 次的遍歷都不會(huì)有 cache miss。雖然兩種 tiling 方式最終 cache miss 的次數(shù)是一樣的，但這種短時(shí)間爆發(fā)的 cache miss 所帶來的延遲是非常難被各種優(yōu)化手段覆蓋的，因為這種延遲不僅短時(shí)間內(nèi)有很多次，同時(shí)每一次的延遲都很長，所以會(huì)造成性能損失。因此以后高性能代碼的開發(fā)中，也要注意合理的把 cache miss 分配到 kernel 運(yùn)行的各個(gè)階段。

Bank Conflict

在查看兩種 Tiling 方式的 profile 我發(fā)現(xiàn)，采用橫著 Tiling 方式的 kernel bank conflict 更低一些。

等等，既然我們之前已經(jīng)處理過 bank conflict 了，那么為什么這里還會(huì)有 bank conflict 呢？這個(gè)現(xiàn)象其實(shí)我也不是很清楚。但目前已知的是，在沒有加 double buffer 情況下是沒有 bank conflict 的，但加了 double buffer 之后或多或少會(huì)出現(xiàn)一些 bank conflict。那么至于為什么橫著 Tiling 方式的 bank conflict 更低，我就更不知道了，因此這里還請各位 dalao 賜教。

最終版本的代碼在這：https://github.com/AyakaGEMM/Hands-on-GEMM/blob/main/src/cuda/double_buffer_yhs_refine_gemm.cu

Epilogue

回顧本文，也基本達(dá)成了文章開頭所立的各種 flag。當(dāng)然現(xiàn)在還是有很多問題沒有解決的，如 split K、長尾問題、分塊細(xì)調(diào)等等，這些權(quán)當(dāng)是一些未來展望了。近期也在嘗試寫一下 int8 tensor core 的矩陣乘，在較小形狀上（M、N、K<=2048）能有比 cublas 更高的性能，但在更大形狀上就只有 80% 左右了（這還是 L2 cache 命中率為 90% 的結(jié)果，可能還有啥別的沒做好），所以就沒有寫 int8 tensor core 的部分。不過好歹是寫完了！

審核編輯 ：李倩