如何高效實現矩陣乘?萬文長字帶你從CUDA初學者的角度入門

本文將從一個 cuda 初學者的角度來闡述如何最佳化一個形狀較大的正方形乘正方形的 FP32 矩陣乘。

矩陣乘作為目前神經網路計算中佔比最大的一個部分,其快慢會顯著影響神經網路的訓練與推斷所消耗的時間。雖然現在市面上已經有非常多的矩陣乘的高效實現——如基於 cpu 的 mkl、基於 arm 設備的 ncnn 與 emll、基於 cuda 的 cublas ——掌握了矩陣乘最佳化的思路不僅能幫助你更好的理解編寫高性能程式碼的一些基本原則,而且許多神經網路加速領域進階的技巧如運算元融合都是與矩陣乘互動從而達到更高的性能。

由於矩陣乘的性能最佳化與兩個矩陣的形狀有著非常密切的聯繫,因此,為了降低本文的撰寫難度(以及輔助讀者更好的理解矩陣乘最佳化),本文將從一個 cuda 初學者的角度來闡述如何最佳化一個形狀較大的正方形乘正方形的 FP32 矩陣乘。同時本文按如下順序講解:

  • Goals:本文的目標是什麼?
  • Performance:我們達到了多少性能?
  • 樸素 GEMM 與前置知識:簡單介紹一下我們的任務是什麼,我們需要提前了解什麼。
  • Tiling:如何做矩陣分塊?即如何將一個巨大的矩陣乘任務合理的分配到 GPU 的不同執行緒上。
  • Thread 級最佳化:在 Thread 這個維度,我們能做什麼最佳化?
  • Warp 級最佳化:在 Warp 這個維度,我們能做什麼最佳化?
  • Block 級最佳化:在 Block 這個維度,我們能做什麼最佳化?
  • Epilogue:尾聲。

Goals

首先明確一下本文的目標是:

• 實現一個比 cublas 更快的形狀較大的正方形乘正方形的 FP32 矩陣乘。

• 從理論角度與硬體規格能夠簡單的推導矩陣分塊與排布的方法。

• 可以大致清楚各個最佳化技術效果的階段性的 benchmark。

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

本文不含:

• 使用 Tensor Core 加速矩陣乘。(這也是為什麼這篇文章叫傳統 CUDA GEMM)

• 使用安培架構新提出的 async memcpy。

• CUDA 語法知識。

• 彙編。(主要是現在並沒有官方支持彙編的操作,目前的彙編器幾乎都是逆向的產物,不是很穩定。同時彙編帶來的好處如消除暫存器的 bank conflict nvcc 也在不斷的做相應的改進,因此就不介紹了)

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

同時本文在相當程度上參考了李少俠的 GEMM 最佳化指南(寫得非常!非常!非常!不錯),本文的優勢在於補了階段性程式碼和在某些少俠一筆帶過的地方做了一些擴展。

Performance

為了讓大家更有動力閱讀下去,這裡先放出來性能效果!

測試平臺

• 系統:Arch Linux

• 驅動:520.56.06

•CUDA:11.8

•GPU:Nvidia RTX 2080

測試結果

測試結果

測試結果

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

從這張圖我們可以看出,手寫的矩陣乘能夠達到硬體 95% 的峰值性能,效果還是很不錯的。

樸素 GEMM 與前置知識

首先寫一個樸素矩陣乘。

首先寫一個樸素矩陣乘
# 陣列 A:M 行 K 列的行主序矩陣# 陣列 B:K 行 N 列的行主序矩陣# 陣列 C:M 行 N 列的行主序矩陣# alpha:一個標量# beta:一個標量# 計算方法:#    c=alpha*A*B+beta*C;__global__ void matrixMul(const float *A, const float *B, float *C,int M, int N, int K, float alpha, float beta){int tx = blockIdx.x * blockDim.x + threadIdx.x;int ty = blockIdx.y * blockDim.y + threadIdx.y;int baseX = blockIdx.x * blockDim.x;int baseY = blockIdx.y * blockDim.y;float c = 0;if (tx < M && ty < N){for (int i = 0; i < K; i++){c += A[tx * K + i] * B[i * N + ty];}C[tx * N + ty] = beta * C[tx * N + ty] + alpha * c; // we multiply alpha here to reduce the alpha cal num.}}

這裡 GPU 矩陣乘與 CPU 矩陣乘最大的區別就在於 GPU 可以為目標矩陣 C 中每一個元素分配一個 thread 進行計算。這也是可以切實的感知到 GPU 多執行緒程式設計的一點。但這個矩陣乘的樸素實現會非常慢,而分析性能瓶頸中最常見的兩個指標即是頻寬和延遲。這裡借用 Nvidia 在 GTC 2018 上的分享來做說明。

這裡以一個自動扶梯作為例子來講解

這裡以一個自動扶梯作為例子來講解。

  • 頻寬:指這個自動扶梯每秒能夠運送多少個人,以這張圖為例子,這個扶梯每秒能運 0.5 個人,這就是這個自動扶梯的頻寬。
  • 延遲:指一個人踏上這個扶梯直到他被運到頂所需要的時間。同樣以這張圖為例,這個扶梯需要 40s 的延遲。

那麼回到指令上來,每一個指令都有對應的延遲和頻寬,而以樸素矩陣乘為例,每一個乘法運算需要讀兩次記憶體和一次 FFMA,假如沒有其他額外的最佳化(如循環展開與指令重排),相當於是兩個級聯的自動扶梯,一個負責運送資料,一個負責做數學運算。假設資料運送扶梯的頻寬與延遲與圖中一致而不考慮 FFMA 的頻寬與延遲,那麼一次 FFMA 需要等待 40s(扶梯延遲)+ (1/0.5)s(第一個資料到達後第二個資料到達的時間)才能拿到所需的資料,這與扶梯的頻寬 0.5s / 人的峰值性能相去甚遠。那麼此時這個 kernel 就完全被延遲卡住了,而無法發揮出應有的性能。

而對於頻寬部分,這裡我們引用李少俠的頻寬分析:

對於 FP32 資料,如上圖所示,一個 warp 一次做 32 次 FFMA,對應 64 OP,需讀取 A 矩陣 1 個元素和 B 矩陣 32 個元素,共 132byte。

通過暫存器累加,且忽略 C 矩陣寫回開銷,那麼計算訪存比為 64OP / 132byte = 0.48。雖然 dram 最小訪問單位為一個 memory transaction,但考慮到 L1 cache 的存在也不會影響實際的計算訪存比。

通過 repo 中提供的 l2cache_bandwidth.cu 可測得 Titan V L2 cache 頻寬約 1.9TB/s,那麼最樂觀的結果即使 L2 cache 100% 命中,此方案的理論上限也只有 1.9T * 0.48 = 0.912 tflops,遠低於 14.9 tflops 的硬體算力。

由此我們可以看出,樸素的矩陣乘實現方法無論從延遲和頻寬上都無法滿足需要。

這裡一個 warp(即 32 個執行緒)是指 GPU 排程執行緒的粒度,可以簡單的理解為同一個 warp 內的執行緒總是同時運行、同時休眠的。當然這種說法並不完全準確,畢竟還有 warp divergent 問題,感興趣的同學可以自行了解。但總之,思考 GPU 執行時總是從 warp 的角度思考是非常合理的。那麼對於一個 warp 而言,我們可以根據李少俠的分析看出,就算我們假設延遲能夠被完全覆蓋,這種分配方案也並不能達到硬體的峰值性能。

這裡我用自己的話總結一遍就是:在每個執行緒執行指令設計時,需要儘可能的覆蓋掉每個指令的延遲;在性能分析時,則從頻寬角度分析矩陣分塊是否合理

而在於延遲部分還有一頓免費的午餐。在實際應用中,編譯器會自動的做一些最佳化,如循環展開與指令重排等。例如展開循環後可以將多個讀取 A 矩陣的元素和讀取 B 矩陣的元素排在一起,使得取資料的自動扶梯能夠一次多上幾個人,從而去覆蓋掉扶梯的延遲。而且 GPU 與 CPU 還有一個非常不同的地方在於,GPU 的執行緒切換代價非常低,因此可以在等待延遲的時候轉而去運行其他執行緒從而達到延遲覆蓋的目的。還是以扶梯為例子,GPU 上有很多個扶梯,在等待第一個人到達扶梯末尾時,GPU 可以轉到第二個扶梯送幾個人上扶梯。理想情況下,在 GPU 送完第 N 個扶梯的人時,第一個扶梯的人剛好達到扶梯頂部,那麼這個運送的延遲就被覆蓋掉了。

Tiling

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

局部性原理

局部性原理,在我的理解中便是為了能夠儘可能的使用高速快取器內的資料進行運算所提出的一個程式設計理念。由於高速快取器往往十分昂貴(或者需要很高的功耗),因此空間都不大。由此我們需要儘可能的將一些重複訪存聚合起來,放到高速快取器裡面來加速資料訪問,或者在進行訪存的時候儘可能連續訪存來使用 cache 加速訪存。我們先還是讓每一個 thread 負責一個目標矩陣元素的計算。雖然這種分配方式十分樸素、十分直接,同時各個 thread 之間也沒有資料依賴關係,不需要做額外的同步之類的操作,但這種分配方式卻是十分訪存不友好的,因為每一個 thread 都是直接與記憶體做互動,而 GPU 的全局記憶體訪問頻寬完全不足以匹配上它的計算速度。

同時我們注意到處於同一行的 thread 總是會同樣的讀取 A 矩陣的同一行資料;同一列的 thread 總是會讀取 B 矩陣的同一列資料。那麼一個非常自然的想法則是對於每一個 Block,我們將資料移動到這個 Block 共享的一塊高速儲存區 shared memory 上,從而減少與全局記憶體互動的次數。同時我們考慮到 shared memory 的容量有限,因此可以一次只取一部分 k,然後通過循環迭代的方式完成這個 Block 所負責的矩陣乘區域。

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

在將一個大型矩陣乘劃分為一個個由 Block 負責的小型矩陣乘之後,我們接下來還需要把一個 Block 負責的矩陣乘分配給 Block 內部的 warp;分配到 warp 之後我們還需要把一個 warp 負責的矩陣乘分配給 warp 內部的 thread。經過這麼一步一步的劃分,我們便可以把一個巨大的矩陣乘任務高效的分配到各級速度不一的儲存器上,最終儘可能打滿硬體峰值性能,實現高效矩陣乘。有了前面劃分 Block 的經驗,我們也就可以依葫蘆畫瓢,實現大矩陣的拆分(Tiling),在此就不過多贅述了,最終整體流程圖如下。

當然這只是一個較為粗糙的流程圖,例如每一個 thread 負責的分塊也並不是圖中所示的連續一塊矩陣乘,我們也將在後續一步一步完善細節,但這種分解的框架卻是一種非常經典的思路。

如何確定分塊大小?

在擁有分塊的基本理念之後,我們還有一個問題沒有解決。那便是每一個 Block 該負責多大的矩陣乘?每一個 thread 又應該負責多大的矩陣乘?為了讓文字變得清晰起來,我們定義每一個 Block 負責的矩陣大小為

,每次迭代

,每次迭代

,每次迭代

的 k 維資料,每一個 warp 負責的矩陣大小為

,每一個 thread 負責的大小為

,每一個 thread 負責的大小為

,每一個 thread 負責的大小為

。其中這些符號都在上圖出現過,可以自行對照一下。

這裡我們同樣引用李少俠的計算訪存分析:

這裡我們同樣引用李少俠的計算訪存分析

假設我們不考慮 shared memory 的訪存代價(因為可以做到覆蓋掉shared memory 的訪存延遲,而且其頻寬能夠滿足 FFMA 單元的計算速度),只考慮全局記憶體的訪問,可以看到選擇在 K 上縮水(即不把整個 K 維度都放到 shared memory 裡)還是比較合理的,因為

的大小其實並不影響計算訪存比。而對計算訪存比有決定性影響的是每一個 Block 計算的大小。如果取

為 64,帶入 RTX 2080 的資料,可以得到 10.1 Tflops / 16 = 631.25 GB/s。即記憶體訪問頻寬達到 631.25 GB/s 就能避免記憶體訪問瓶頸了。同樣,我們取 L2 命中率為 20%(還是比較好達到的),加權記憶體訪問頻寬為:

,即可避免記憶體訪問瓶頸

,即可避免記憶體訪問瓶頸。

那是否我們只要取分塊大小為 64×64 就萬事大吉了呢?也不盡然。我們前面只分析了頻寬,而在延遲無法被覆蓋的情況下,整個 kernel 性能也不會太好。而更大的分塊意味著每一個 thread 會計算更多的資料,可以使用一些手段實現更優的延遲覆蓋。這一點會在後面討論如何具體實現,大致思想也是局部性的原理,只不過這次是將資料從 shared memory 保存到暫存器,從而實現使用更高速的快取計算的目的。

那是否我們取分塊越大越好呢?那也不一定。更大的分塊使用了更多的暫存器,從而使得同一個 SM 能夠同時承載的執行緒數變少,這裡 Nvidia 將之稱為 Occupancy。如前文所述,當一個 warp 被卡住時,GPU 可以切換到另一個 warp 執行指令,Occupancy 越低,可供 GPU 切換的執行緒就越少。

而 Occupancy 也是和硬體強相關的。一個 GPU 由多個 SM 構成,每一個 SM 擁有有限的暫存器數量、 shared memory 和最大可排程執行緒數量。而 Occupancy 是指每個 SM 能夠同時排程的執行緒數量除以一個 SM 的最大可排程執行緒數量。關於 Occupancy 的計算我們可以通過在編譯時添加 –ptxas-options=-v 參數,使編譯器在編譯時輸出每個 kernel 所花費的暫存器數量和 shared memory,然後通過隨 cuda 提供的一個 excel 表格進行計算。(儘管這個 Excel 已經 deprecated 了,但他用起來確實挺方便的。)

例如我們每個 thread 需要 128 個暫存器,2048 bytes 的 shared memory,那麼由於 RTX 2080 每個 SM 只有 65536 個暫存器,因此每個 SM 最多只能同時跑 512 個 threads。又因為每個 SM 最多能夠承載 1024 個 threads,所以此時 Occupancy 為

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

因此,我們能夠做的就是在有一定理論分析的情況下確定好一些矩陣的分塊大小的方案,然後要不就是經驗性的去選擇最終用哪個分塊,要不就是跑一個 profile 來直接得到最快的分塊。這裡由於已經有非常多的先例證明了 128x128x8 是一個較優的選擇,因此本文則遵從這個分塊方案。那麼,目前我們能夠確定的分塊如下表。

當然有些同學可能會問,既然最終還是需要用跑 profile 的方式來確定最優分塊,那理論分析還有什麼意義呢?答案就是如果提前通過理論分析,那麼就能夠在一定程度上縮小需要跑 profile 的分塊數量。用演算法上的語言來講就是如果我們將需要搜尋的所有分塊作為搜尋空間,那麼理論分析便是搜尋演算法中的 A* 演算法,你掌握了越多的理論分析知識那麼這個搜尋過程就會越高效。同時對 CUDA 底層越了解,在同一個分塊策略下,你更容易寫出能達到理論性能的 kernel。

Thread 級最佳化

對於一個 thread 能做的最佳化其實並不多,因為 GPU 是以一個 warp(即 32 個 thread)進行排程的,所以許多基於單執行緒的最佳化,如訪存最佳化,其實並不能直接套到 GPU 上。而為數不多值得一提的最佳化手段便是單個執行緒在計算時應該採用向量內積還是向量外積以及 double buffer。但實質上向量外積嚴格意義上也不能算作是一個最佳化,因為這一步編譯器就能在編譯階段幫忙做了。之所以提一句是還是為了給 double buffer 做鋪墊,即我們應該怎麼預取資料。

首先我們取了 128×128 的分塊策略,一個 Block 有 256 個執行緒,那麼每個執行緒需要負責一個 8×8 的矩陣乘運算。而一個執行緒完成一個小型矩陣乘有兩種實現方法。

向量內積

向量內積的實現方法如圖所示,即將 A 矩陣拆分為多個向量、B 矩陣拆分為多個向量,這些向量通過向量內積的方法求得最終答案。

用程式碼描述如下

用程式碼描述如下:

M=N=K=8;float a[M*N];float b[N*K];float c[M*N];for i in range(M):for j in range(N):for k in range(K):c[i*N+j]+=a[i*K+k]*b[k*N+j];

向量外積

向量外積的實現方法如圖所示,即將 A 矩陣拆分為多個向量、B 矩陣拆分為多個向量,這些向量通過向量外積的方法求得最終答案。

用程式碼描述如下

用程式碼描述如下:

M=N=K=8;float a[M*N];float b[N*K];float c[M*N];for k in range(K):for i in range(M):for j in range(N):c[i*N+j]+=a[i*K+k]*b[k*N+j];

可以看到,向量內積和向量外積的區別在程式碼上僅僅體現在循環方式上。

為何我們需要關心這個?

有做過 CPU 矩陣乘最佳化的同學可能知道,僅僅調整循環順序就已經能夠帶來顯著的性能差異了。有許多分析都是從局部性的角度進行分析的。即使用向量外積的方案可以利用到循環遍歷的局部性,將一些重複訪存使用暫存器快取而避免無意義訪存。例如我們補充一下采用向量外積方案關於暫存器的細節。

float a[M*N];float b[N*K];float c[M*N];for k in range(K):regB[0:N] = b[k*N:(k+1)*N]for i in range(M):regA = a[i*K+k];for j in range(N):c[i*N+j]+=regA*regB[j];

其中 regA 和 regB 均為暫存器。其中我們不難發現,對於每一次循環 j ,使用的都是完全相同的 A 矩陣裡的元素,因此可以用一個暫存器來快取該值;對於每一次循環 k,使用的都是完全相同的一行 B 矩陣中的值,因此我們可以用 N 個暫存器快取該值。於是將原本

次訪存(底下兩層循環需要訪問一次 A 矩陣和一次 B 矩陣),通過使用

個暫存器快取(B 使用 N 個,A 使用一個),最佳化為 N+M 次訪存。同時我們也注意到, M 和 N 越大的情況下,提升效果越發顯著,這也是為什麼我們希望每一個執行緒負責的分塊大一點比較好。但同時 M 和 N 越大,每一個執行緒多使用的暫存器就越多,而在 GPU 的語境下,更高的暫存器佔用意味著更低的 Occupancy。因此當 M 和 N 大到 shared memory 頻寬不是性能瓶頸即可。更詳細的分析可以看李少俠的分析。

而我則從循環展開的角度解釋一下為什麼我們需要了解這個最佳化方案,同時解釋一下為什麼該最佳化方案在 GPU 上並不如 CPU 上那麼有效。從循環展開的角度來看,第二種循環體構造與第一種循環最大的區別就在於它能在不展開 k 的情況下通過展開 m 和 n 處的循環就能自動的識別到重複訪存,並使用相應的暫存器來避免重複訪存。例如我們假定

,那麼展開 m 和 n 處循環的結果如下。

M=N=2;float a[M*N];float b[N*K];float c[M*N];for k in range(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]

只要是稍微現代一點的編譯器,都能一眼看出這四條指令的 8 次訪存,有 4 次是可以合併的。同時現代一點的編譯器也能在一定程度上根據生成的彙編交叉排列計算和訪存達到延遲覆蓋的目的。而向量內積的方案需要把整個 k 維度展開才能看到這些潛在的訪存合併機會。在 CPU 矩陣乘的語境下,一般計算 kernel 的

都比較大(好幾百),而

都比較大(好幾百),而

都比較大(好幾百),而

都很小(一般取 6×16,根據架構來做具體確定),暫存器數量又非常少,因此基本上無法在 K 維上將循環完全展開並做最佳化。因為展開一個超長的循環不僅會帶來額外的暫存器佔用、最佳化難度,還會帶來更多的彙編指令,使得最終的二進位制檔案臃腫不堪。但在 GPU 上,情況卻恰恰相反。對於已知循環次數的小循環,即便你沒有指定 #pragma unroll,nvcc 也會自動的展開這些循環。而對於一個 thread 所負責的小型矩陣乘,這三層循環的值均為 8,符合 nvcc 自動展開循環的條件。而在展開完成後,nvcc 會對所有的訪存以及計算指令重排得到一個不錯的彙編指令排列。

那麼這就引出了下一個問題,我們為何還需要關心他究竟是向量內積還是向量外積?

答案就是 double buffer。如果我們能夠提前知道一個循環需要什麼資料,我們就能提前預取該循環第一次所需的資料,同時在該循環進行運算的時候預取下一次計算所需的資料。而顯然這在向量內積的情況下是無法做到的。同時由於 double buffer 需要額外的暫存器保存從 global memory 轉移到 shared memory 的資料,所以當一開始循環展開使用的暫存器過多時,儘管後續能最佳化到較少的暫存器,但編譯器依然無法正確的在限定暫存器數量下實現 double buffer。這一點在最佳化 sgemm 的時候並不是那麼重要(因為多使用一點暫存器也就從每個 SM 跑兩個 block 變為一個 block),但是在最佳化 int8 矩陣乘時需要額外的關注(因為本身它就只能在一個 SM 上跑一個 block,如果實現不得當將會完全失去 double buffer)。

那麼此時樸素的利用到向量外積和 shared memory 的程式碼:https://github.com/AyakaGEMM/Hands-on-GEMM/blob/main/src/cuda/shared_mem_gemm.cu

Double Buffer

由於 GPU 沒有 prefetch 這種指令,同時我們又有 shared memory 這種可程式設計的 L1 cache,因此需要手動實現 prefetch 功能,而在 GPU 語境下一般被稱作 double buffer。double buffer 的好處自不必多說,即它可以實現資料讀取與計算在時間上重疊,利用 FFMA 單元與 Load/Store 單元可以並行執行指令的特點,達到覆蓋延遲的目的。而儘管 GPU 可以在一個 warp 有延遲的情況下,通過切換去運行另一個 warp 達到延遲覆蓋到目的,但由於可供 warp 排程器能切換到執行緒數量的限制,過於長的延遲並不能通過這種方式覆蓋掉。這裡引用一下李少俠更詳細的分析:

若每 SM 有 4 個排程器,若每個排程器只有 4 個可排程 warp,當指令平均間隔超過 4 cycle 後就無法靠 warp 排程掩蓋延遲了。考慮到 GEMM 中涉及 smem 讀寫的過程需要同步 thread block,進一步限制了 warp 排程空間,所以很難靠 warp 並行掩蓋延遲。

而本文最終實現的 kernel occupancy 只有 50%,即每個 SM 只能排程 512 個 threads(16 warps),加上圖靈架構每 SM 有 4 個 warp 排程器,最終結果與李少俠分析的一致。因此 double buffer 從指令角度提供的延遲覆蓋方法最終還是會有效的。

但值得一提的是,在你自己動手實踐時,儘可能的考慮在其他最佳化已經加無可加的情況下再加入這個最佳化。這是由於這個最佳化會大幅修改資料讀取部分的程式碼,而且還會產生重複程式碼,不利於程式碼維護。同時在我自己的實踐中發現,如果在一開始 kernel 寫的比較垃圾,加了 double buffer 也沒有什麼卵用,還會讓後續的最佳化不太好加上去。當然,這只是我的個人建議,如果你想實際看看 double buffer 的效果也可以一開始就加上去。

首先我們看一下每個 thread 的運行流程。

那麼能實現 double buffer 的機會有兩個地方:Global Memory to Shared Memory 與 Shared Memory to Register。即在每一次 FFMA 開始之前我們讀取 Global Memory 的資料到暫存器中,在 FFMA 之後將該暫存器中的值寫到 shared memory 中。由於在讀取資料後 load from shared memory 以及 FFMA 兩個流程中我們並不依賴於該暫存器中的資料,因此可以覆蓋 Global Memory 的讀取延遲。而同時在計算每一次 FFMA 之前,我們可以用暫存器提前取下一次 FFMA 需要的資料,也就能做到覆蓋 shared memory 的延遲。

大概就是這樣!我們在每一次運算之前提前將第一次循環所需的資料移動到暫存器中,這樣我們就可以實現資料運算和資料存取指令級並行的功能了。

Warp 級最佳化

在做了不少鋪墊之後,接下來的最佳化終於是可以帶來一些看得見的性能提升的了。首先回顧我們之前的程式碼,可以看到每個 thread 負責的部分完全沒有考慮到它們之間可能的協作關係,即同一個 warp 內的 thread 此時在同一塊硬體上同時執行——它們共享同一個 register file,這表明它們可以通過暫存器快速共享資料(即 shared memory 的 broadcast 機制);它們會同時訪存,這表明如何安排每一個 warp 內的 thread 訪存是至關重要的。

Warp Tiling

已知我們指定一個 Block 計算 128×128 的矩陣,一個 Block 有 8 個 warp,一個 warp 有 32 個 thread,每個 thread 需要負責 8×8 的小型矩陣乘,那麼我們沿用李少俠的定義:

一個 warp 由

個執行緒組成,可以是

個執行緒組成,可以是

個執行緒組成,可以是

,我們把這些執行緒對應的 thread tile 拼在一起的區域稱為 warp tile,尺寸為

,如下圖所示

,如下圖所示。

這裡的圖給的是

這裡的圖給的是

這裡的圖給的是

的排列方式。由於同一個 warp 在訪問 shared memory 時有 broadcast 機制(即同一個 warp 在訪問同一個記憶體地址內的值時只會實質發生一次資料讀取),因此這一個 warp 計算時只會實際讀取

個 float。與之相對的,這個 warp 會進行

次 FFMA不難看出,在

次 FFMA。不難看出,在

固定為 32 的情況下,

固定為 32 的情況下,

固定為 32 的情況下,

越相近,計算訪存比就越大,因此取

越相近,計算訪存比就越大,因此取

最為合適

最為合適。

而在確定了 warp tiling 後,如何讀取和儲存資料的細節還需要細扣,接下來我將會按照 GPU 的硬體特性講解讀寫資料的細節。但這一部分的大致思路基本已經介紹完畢了,動手能力強的同學現在就可以自己試試如何寫一個高效矩陣乘了!

向量化訪存

向量化訪存即是一條指令同時請求多個 float 資料,目前 CUDA 最高支持 128 bit 的向量化訪存,即一條指令請求 4 個 float 資料。向量化訪問主要的好處在於可以用更少的指令讀取更多的資料。由於在訪問全局記憶體時是以 32 Byte 為粒度進行訪問的,因此如果同一個 warp 內的 thread 請求了一段連續記憶體的資料,每一個 thread 都請求兩次 4 Byte 的資料(小於 GPU 全局訪存的最小單位),那麼 GPU 會在硬體處將 64 次資料請求按照 32 Byte 進行合併,最終形成 8 次 32 Byte 記憶體訪問。

而如果每一個 thread 請求 8 Byte 資料,那麼 GPU 會在硬體處同樣將 32 次資料請求按照 32 Byte 進行合併,最終形成也形成 8 次 32 Byte 記憶體訪問。

那麼我們可以看出,對於訪問同一資料量的資料,請求的指令越多,GPU 的聚合訪問的壓力就會越大。在極端情況下,儘管頻寬足夠,但大量的訪存請求會塞滿訪問隊列導致 stall。這在 Nsight Compute 中顯示為 MIO Throttle 和 LG Throttle,即對應 shared memory 和 global memory。因此採用向量化訪存能在一定程度上緩解 GPU 硬體層面的聚合訪存壓力(因為我們顯式的用指令告訴 GPU 某些資料請求不需要聚合,直接用一個 sector 來處理就好了)。

但使用向量化訪存——即用 float4 讀寫資料——也不是完美的。它的一個嚴重缺陷在於使用 float4 訪存要求請求的資料地址要按照 float4 對齊,因此當 M、N、K 不為 4 的倍數時將會報 missaligned address 錯誤(因為第二行開始就不能按照 float4 對齊了)。

這麼幹對輸入矩陣形狀有一定要求,寫出來的矩陣乘沒有特別好的通用性。同時 sgemm 受聚合訪存的影響也並不是那麼大,因此在實操中往往並不會選擇使用 float4 讀寫全局記憶體,而只會使用 float4 讀寫 shared memory。但由於我一開始學 CUDA 的時候對這一塊理解也不深,然後發現許多人(李少俠除外)都很暴力的直接用 float4 讀寫全局記憶體,於是我也用了 float4 讀寫全局記憶體。

而我們這裡對比李少俠的 kernel profile 和我們最終的成品發現,在 global memory 讀取處是否使用向量化讀取其實並不會對性能有多少影響。可以看到最終 profile 出來的 Stall LG Throttle 和 Stall MIO Throttle 佔比都不高。

上圖為李少俠的 kernel 下圖為我最終寫的 kernel。這兩個 kernel 在資料讀取方面的區別就是李少俠是以 4B 為單位訪存的,而我是以 16B 為單位做訪存的。這進一步印證了 sgemm 其實並不是非常關心讀取 global memory 時是以怎樣的粒度讀取的。而向量化訪存對於 shared memory 的影響就留給讀者自行驗證了。同時值得注意的是,在把資料讀取方式從向量化訪存修改為一個一個訪存時需要注意 bank conflict 的問題。因為一個 warp 在執行 128-bit load 和 32-bit load 時的排程並不相同(這點會在後面提到)。

還有一個值得注意的是在 Global Memory 訪存時,並不能直接將原先的向量化存取程式碼直接改成一個一個的讀取。因為這樣訪存從原來一個 warp 並行訪問一段連續的記憶體變成一個 warp 分成四次訪問不連續的記憶體。雖然有 L2 cache 來平滑這種不規則的訪存,但最終會帶來 10% 左右的性能下降。程式碼如下:

// Original CodepreA = *reinterpret_cast(baseA + i + rowA * K + colA);// Modified CodepreA.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];

可以看到這種簡單的更改其實並不可取,更優的寫法是每一條指令都是在 warp 視圖下的連續訪存。

Global Memory

前面提到 GPU 訪存時以 32 Byte 為粒度進行訪問的,那麼一個 32 Byte 訪問被稱為一個 sector。那麼值得注意的就是在搬運資料時,儘可能的讓同一個 warp 搬運同一行的資料來避免使用額外的 sector(本文采用現代的行主序來儲存矩陣)。

這裡借用一下 Nvidia 的圖。如果同一個 warp 內的 thread 都訪問每一行的開頭,那麼如果一行超過 8 個 float,那麼每一個 thread 都需要一個 sector 去請求它們需要的資料,這就造成了 sector 浪費。而實際中每一行的元素往往都會大於 8 個 float,因此會有非常大的性能損失。下圖為一個 warp 在複製時,每個 thread 之間間隔的大小,單位為 float。可以看到在間隔為 2 時就已經有一半的性能損失了,這很不好。

因此我們採用下圖所示的訪問方式。即儘可能的讓一個 warp 中的 thread 連續的讀取 Global Memory 中的元素。

Shared Memory

Shared Memory

前文已經講過,shared memory 在圖靈架構之後可以完全被看作是 L1 cache。而在此基礎之上,shared memory 的訪問粒度是 32 bit 也就是 4 Byte,剛好是一個 float 資料的大小。而後 shared memory 按照 4 Byte 連續的劃分為一個個 bank。對於 bank 可以簡單的理解為雙通道記憶體中通道的概念,即在不同的 bank 中的資料可以並行訪問,同一個 bank 內不同地址的資料只能序列訪問。在 Compute Capability 5.x 及之後的卡上,shared memory 具有 32 個 bank,剛好是一個 warp 中執行緒的數量。而如果同一個 warp 中不同 thread 均只訪問 4 Byte 資料且希望同時訪問同一個 bank 的資料將會有兩種結果。(對於每一個 thread 訪問更多資料的行為將在後面提到)

1. 兩個或多個 thread 訪問的剛好是同一個地址內的資料,那麼此時將會觸發 broadcast 機制,即實際只讀取一次資料,而後廣播到這些 thread 中。

2. 兩個或多個 thread 訪問的是同一個 bank 內的資料,那麼此時這些 thread 的訪問將會被強制安排為序列執行。這種訪問情況被稱為 bank conflict。

這裡給出 cuda programming guide 的兩張圖來直觀的體現 broadcast 和 bank conflict。

這張圖表示同一個 warp 中的 thread 按不間隔、隔一個、隔兩個 bank 對 shared memory 訪問。中間的訪問每兩個 thread 都會發生一次 bank conflict,而其他兩種訪問都不會發生 bank conflict。值得注意的一點是這張圖最右側的圖的訪問方式剛好可以達到每一個 thread 都訪問了不同的 bank 的效果。

同時考慮到 shared memory 是按照 bank 來訪問的,且與 Load/Store 單元直連,並沒有中間商賺差價,所以對於 shared memory 的訪存並不講究連續訪存,而只需要考慮是否有 bank conflict 就足夠了。因此理論上最左和最右兩列圖的訪問性能是一樣的,這與訪問全局記憶體有一點區別。同理,每一個 warp 連續的多次訪存也並不要求連續訪存,而在複製資料到 shared memory 時對 A 矩陣做矩陣轉置的目的是為了向量化訪存,而不是為了連續訪存。

這張圖則展示了 broadcast 機制,沒什麼好說的。

128-bit conflict-free store

而前文中提到,我們使用 float4 來做資料傳輸來緩解 GPU 聚合訪問的壓力,使得每一個指令都更加高效。而又因為前文所述,每個執行緒需要使用向量外積的方法計算矩陣乘,因此我們需要在 A 矩陣轉存到 shared memory 時做一次轉置。

但細心的同學可能注意到,如果就這麼平鋪直敘的做轉置那麼將會發生非常嚴重的 bank conflict,因為一個 warp 內的奇數 thread 和偶數 thread 使用同一個 bank。那麼此時解決 bank conflict 的方法有兩種,第一種便是將 shared memory 的 k 維度縮小,然後直接把奇數 thread 所取的資料直接併到 M 維上,就不會有 bank conflict 的問題了。這種方法通過 index 變換,直接就能避免 bank conflict,非常巧妙,而我當時沒有想到,就沒有用這種方法。值得注意的是,儘管圖是按行隔開的,但那只是為了表示資料是如何在一個 thread 裡保存的,實際寫到 shared memory 中是以一個 float 為單位,按列主序儲存到 shared memory 中。

而第二種方法就非常簡單粗暴了,直接往 lda 上加 4,然後就不會有 bank conflict 了。當然這種方法的弊端也是有的,那就是會造成一部分 shared memory 的浪費。但對 sgemm 來說倒也還好, shared memory 的佔用也不是導致 Occupancy 降低的原因,所以我就用了這個方法。

128-bit conflict-free load

而我們把資料儲存到 shared memory 之後,下一步便是考慮如何在沒有 bank conflict 的情況下將資料讀取出來。在本文中,我們取

為 8,在採用向量化存取時,直接按照 Warp Tiling 採用樸素的存取方法就能在沒有 bank conflict 的情況下把資料讀出來了。

當然有的同學可能會問:既然訪存是按照一個 warp 為單位進行的,而圖中明顯讀取 B 矩陣時,t16 會和 t0 發生 bank conflict,那為什麼又說不會有 bank conflict 呢?那麼答案就是在做 128-bit 訪存時,warp 並不是同時讀取資料的。這裡還是借用 Nvidia 在 GTC 2018 上的分享來做說明。

當 warp 中每個 thread 只讀取 4B 或更小資料時,warp 才是同時讀取的。而本文中採用 128-bit 也就是 16B 讀取,那麼一個 warp 會分成 4 次操作讀取,每次操作只有 1/4 warp 工作。那麼只要同一次操作內的 thread 沒有發生 bank conflict,那麼就沒有 bank conflict。而上圖中 t0-t7 同時操作,它們之間並沒有 bank conflict,後面的 thread 依此類推,那麼也就不會有 bank conflict。那麼樸素的 warp tiling 實現程式碼在這:https://github.com/AyakaGEMM/Hands-on-GEMM/blob/main/src/cuda/warp_tile_gemm.cu

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

這裡我們跑一個 profile 發現,確實是沒有 bank conflict 的,挺好。程式碼在這:

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

Block 級最佳化

Block 級最佳化

Block 在 GPU 上基本等同於不同的 kernel 在 GPU 上運行了,所以它們之間的聯繫並不是特別強烈。而它們之間的相互關係在 GEMM 語境下基本就只有 wave 和 L2 cache(一個 wave 裡的 Block 共享這一塊 cache)了,良好的 Block Tiling 能提升相當可觀的 L2 cache 命中率。

但這一部分屬於 sgemm 並不是特別關心的部分,因為本身 FFMA 單元算的就不是很快,所以 Block Tiling 隨便搞搞就能夠滿足 FFMA 單元的頻寬和延遲需求了。因此,這一節的內容主要是為了有些有用到 tensor core 的同學提一個需要注意的性能提升點,其次就是有些同學可能會發現自己寫的 kernel 可能會比本文中的示例慢一點(大約 10% 左右),因此在此提一下在 sgemm 中應該怎麼隨便搞搞 Block Tiling。

Wave & L2 cache Hit Rate

首先明確一下 wave 的概念,即一個 GPU 上能夠同時運行的 Block 數量。關於 GPU 是如何決定一個 wave 由哪些 Block 組成的我並沒有找到非常明確的文件說明,但我一拍腦袋想,說不定就是樸素的按順序決定的,即 index 處於

範圍內的 Block 處在第一個 wave 中,後面的 Block 依此類推。後面試了試好像的確是這樣劃分的。

在明確了 wave 的概念後,我們便可以對 L2 cache 命中率做一個簡單的分析了。我們指定

代表一個 wave 同時運行的 Block 數量,假設一個 wave 剛好能計算 C 矩陣的整數行,那麼我們不難發現對於一個 wave 而言,它需要從 Global Memory 中讀取

個 float。但由於有 L2 cache 的存在,假設一個 wave 讀取的資料全能被 L2 cache 裝下,那麼實際只讀取了

的資料。最終 L2 cache 的命中率為:

差距越大,L2 cache 的命中就越低。那麼如果想要去最佳化 L2 cache 命中,一個比較直接的想法就是儘可能把一個 wave 的 Block 變成方的。但就算不搞,sgemm 也不在乎,因為其實對性能來講並沒有什麼區別,所以就沒搞。

SGEMM Block Tiling

SGEMM Block Tiling

而在 sgemm 的語境下,假設最壞的情況即一個 wave 都不能覆蓋目標矩陣 C 的一行,且 RTX 2080 有 46 個 SM,一個 SM 能跑兩個 Block,此時

,帶入上式可得,此時 L2 cache 命中率大概是 49.4%。這裡我們並沒有考慮訪問 C 矩陣的影響,在實踐中會把 L2 cache 的命中率拉低一點。但即便是如此,前文我們分析過只要 L2 cache 命中達到 20%,在頻寬上就不會造成性能瓶頸了。因此發現,就算我們採用樸素的 Block Tiling,Global Memory 訪問也不會成為訪存瓶頸。

但事實真的是這樣嗎?

細心的同學可能會發現,上圖所採用的 tiling 方式並不是直覺上的用 blockIdx.x 表示 Block 在 M 維上的位置,而是用 blockIdx.y 表示 Block 在 M 上的位置。而我們簡單調換一下程式碼中 blockIdx.x 與 blockIdx.y 的順序,瞬間就有了 10% 左右的性能差距。目前網上並沒有針對這個現象的分析(因為幾乎所有人都是用的 col major 的 data layout,而且李少俠直接就在程式碼裡使用了更優的 tiling 方式,所以沒有人遇到這個問題),因此我這裡提出一點個人的猜想,如果猜的不對還請指正。

L2 cache

L2 cache

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

為什麼這麼說呢?因為我採用的是行主序的方式儲存的矩陣,因此如果一個 wave 的形狀是扁平的,那麼每個 Block 在每一次循環遍歷 B 矩陣時只會有

次 cache miss。這是由於 L2 cache 的 cache line 大小為 128 bytes,因此當資料從 Global Memory 中移動到 L2 cache 後,許多 Block 就能直接從 L2 cache 中讀取資料了。然而如果一個 wave 的形狀是狹長的。那麼每個 Block 在第一次訪問 A 矩陣的每一行時都會有 cache miss 的情況出現,即產生

次 cache miss,而後 31 次的遍歷都不會有 cache miss。雖然兩種 tiling 方式最終 cache miss 的次數是一樣的,但這種短時間爆發的 cache miss 所帶來的延遲是非常難被各種最佳化手段覆蓋的,因為這種延遲不僅短時間內有很多次,同時每一次的延遲都很長,所以會造成性能損失。因此以後高性能程式碼的開發中,也要注意合理的把 cache miss 分配到 kernel 運行的各個階段。

Bank Conflict

在查看兩種 Tiling 方式的 profile 我發現,採用橫著 Tiling 方式的 kernel bank conflict 更低一些。

等等,既然我們之前已經處理過 bank conflict 了,那麼為什麼這裡還會有 bank conflict 呢?這個現象其實我也不是很清楚。但目前已知的是,在沒有加 double buffer 情況下是沒有 bank conflict 的,但加了 double buffer 之後或多或少會出現一些 bank conflict。那麼至於為什麼橫著 Tiling 方式的 bank conflict 更低,我就更不知道了,因此這裡還請各位 dalao 賜教。

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

Epilogue

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

相關文章

ChatGPT 背後的經濟賬

ChatGPT 背後的經濟賬

拿投資和賺錢之前,要先去了解下成本。 ChatGPT 能否取代 Google、百度這樣的傳統搜尋引擎?為什麼中國不能很快做出 ChatGPT...