本備忘錄將回歸GPU基礎知識，重點介紹深度學習的核心運算：矩陣乘法。首先，我將描述其基本實作和理論方法。然後，我將探討如何優化GPU效能，使其發揮最大效用。
讓我們從一個你在教科書中經常看到的簡單例子開始。
/**
* Kernel to perform matrix multiplication C = A × B
* A is of size N x M, B is M x N, and C is N x N
* Each thread computes one element of matrix C.
*/
__global__ void naive_matrix_multiply(float* A, float* B, float* C) {
// Calculate global row and column for this thread
int row = blockIdx.y /* 0...16 */ * blockDim.y + threadIdx.y;
int col = blockIdx.x /* 0...16 */ * blockDim.x + threadIdx.x;
if (row < N && col < N) {
float sum = 0.0f;
for (int k = 0; k < M; k++) {
// C[row][col] = sum(A[row][i] * B[i][col])
sum += A[row * N + k] * B[k * N + col];
}
C[row * N + col] = sum;
}
}
乍一看，這似乎很簡單。你只需要計算輸出矩陣 C 的每個單元格。這段內核程式碼看起來確實能完成這項任務。下面是該內核的更直觀的表示。
實際情況要複雜一些。如何理解 blockIdx、blockDim 和 threadIdx 呢？要理解這一點，我們需要回顧一下 GPU 的組織結構，它圍繞著執行緒、執行緒束和串流多處理器。以下是英偉達 H100「Hopper」GPU 微架構的主觀描述：
從宏觀層面來看：
1 個 GPU = 10 到 100 個串流多處理器
1 流式多處理器 = 4 ~ 64 個封裝
1圈 = 32~64根線
那麼，計算 sum 的程式碼片段是如何與執行緒、執行緒束，甚至是串流多處理器關聯的呢？記憶體分配和映射決策是在哪裡做出的？首先，我們需要了解核心是如何配置的：
int launch_kernel(float *d_A, *d_B, *d_C) {
dim3 block_size(16, 16); // 16x16 threads per block -> 256 threads
dim3 grid_size(4, 4); // 4x4 blocks per grid = 64 blocks
// 2. Launch the kernel with the specified dimensions
naive_matrix_multiply<<<grid_size, block_size>>>(d_A, d_B, d_C);
}
這種分層配置映射到GPU硬體的方式如下：
在這種配置下，可以對大小為 N 的矩陣進行乘法運算，其值為 16*4 = 64。如果 N 為 1000 呢？在這種情況下，只需將網格大小設為 (1000+15)/16 即可。
那麼，如何映射到實際硬體呢？這是否意味著將資料塊分配給線程束？如果資料塊中的執行緒數大於硬體執行緒數該怎麼辦？
這就需要引進GPU調度機制了。簡單來說，如果需要的執行緒數超過了實體硬體上的執行緒數，那麼GPU就會分配多個執行緒束（warp）：
下一個問題是，每個資料區塊所需的 8 個執行緒束如何與串流多處理器關聯起來，尤其是在串流多處理器擁有的執行緒束數量少於資料區塊所需數量的情況下。答案非常巧妙：一個串流多處理器可以擁有超過 4 個執行緒束，並且會交替調度這些執行緒束：每當一個執行緒束因記憶體存取而停滯時，串流多處理器會立即切換到另一個執行緒束。這被稱為延遲隱藏，它是 GPU 效能提升的核心策略。如果沒有這項技術，GPU 的效率就不會如此高。
現在考慮 grid_size。 grid_size 是否直接對應到硬體資源？並非如此。網格只是一個邏輯概念，它確保有足夠的塊來覆蓋矩陣。 grid_size 並不對應硬體。區塊是在運行時從佇列中分配給每個 SM 的。
分配依據是區塊所需的資源（每個包裝器都有有限數量的暫存器，並且需要確保每個區塊都有自己靜態分配的一組暫存器）。
當一個計算區塊完成計算後，SM 會從調度佇列中取出一個新的計算區塊來取代已完成的計算區塊。一旦分配完成，該計算區塊就會被綁定到該 SM，並且其暫存器會被指派以供執行。
最後一個問題：記憶體區塊是綁定到服務管理員 (SM) 還是綁定到線程束 (warp)？線程塊中線程使用的記憶體是否分配在線程束中？如果不是，服務管理員就需要在不同線程束之間移動記憶體。是的，在服務管理員內部，記憶體區塊的執行緒束是綁定的。如果記憶體區塊 0 分配到線程束 0，記憶體區塊 1 分配到線程束 1，則始終如此。
現在我們了解並發線程執行，考慮一個簡單的矩陣乘法內核。所有線程都競爭對矩陣 A 和 B 的記憶體存取。每個線程都從全域記憶體中讀取 A 的一行和 B 的一列。記憶體無法重複使用。 B 是按列存取的（B[i * N + col]），因此突發存取和記憶體快取都無法有效利用。
最簡單的解決方案是使用共享記憶體。首先，將矩陣的一部分以大小為 T*T 的區塊形式載入到線程束記憶體中。然後，讓線程僅從這部分本地記憶體讀取資料。這樣，線程之間就不會爭用記憶體。透過線程束交錯，SM 可以確保當一個線程束將記憶體載入到快取中時，另一個線程束可以計算總和。以下是增強後的核心：
__global__ void matmul_tiled(float* A, float* B, float* C) {
__shared__ float A_shared[T][T];
__shared__ float B_shared[T][T];
int tx = threadIdx.x, ty = threadIdx.y;
int row = blockIdx.y * T + ty; // global row in C
int col = blockIdx.x * T + tx; // global col in C
float accumulator = 0.0f;
for (int t = 0; t < M; t += T) {
// Read [T,T] subtile from matrix A and B.
// Assume that M and N are multiples of T (no bound check)
A_shared[ty][tx] = A[row * M + t + tx];
B_shared[tx][ty] = B[(t + ty) * N + col];
__syncthreads(); // wait all threads to finish loading
for (int k = 0; k < TILE; ++k)
accumulator += A_shared[ty][k] * B_shared[k][tx];
__syncthreads(); // all threads are ready to load again
}
C[row * N + col] = accumulator;
}
為簡單起見，我假設 M 和 N 是 T 的倍數，因此在將全域矩陣載入到 A_shared 和 B_shared 時無需進行邊界檢查。
這段程式碼裡藏著一些玄機。程式碼區塊內的所有執行緒都執行相同的記憶體讀取操作：A[row * M + t + tx]。由於 t 是連續的，線程從同一記憶體區域讀取資料。這使得記憶體突發讀取和快取成為可能。第一個執行緒發起突發讀取，其他執行緒則使用由該突發讀取填充的快取。由於 B 沒有被轉置，所以這個技巧對 B 無效。此外，__syncthreads() 函數會同步程式碼區塊內的線程，確保它們運行在同一個記憶體區塊上。
使用之前的簡單乘法器時，區塊大小是任意的。而使用分塊方法時，區塊大小必須與圖塊大小 T*T 相匹配，如圖所示：
分塊方法允許部分資料區塊加載，而其他資料區塊則進行計算。在這種情況下，計算量可能很小，耗時比記憶體載入時間更短。為了解決這個問題，可以在計算期間增加每個執行緒的工作量。每個執行緒可以計算一個 2x2 的元素集，而不僅僅是一個。當 K=2 且 T 是 K 的倍數時，這是可行的。
_global__ void matmul_tiled_and_coarsened(float* A, float* B, float* C) {
__shared__ float A_shared[T][T];
__shared__ float B_shared[T][T];
int tx = threadIdx.x, ty = threadIdx.y;
int rowBase = blockIdx.y * T * K + ty;
int colBase = blockIdx.x * T * K + tx;
float accumulators[K][K]; // That's register memory
for (int t = 0; t < M; t += TILE) {
A_shared[ty][tx] = A[row * M + t + tx];
B_shared[tx][ty] = B[(t + ty) * N + col];
__syncthreads();
for (int k = 0; k < TILE; ++k)
for (int wr = 0; wr < K; ++wr)
for (int wc = 0; wc < K; ++wc)
acc[wr][wc] += A_shared[ty + wr * (T/K)][k] * B_shared[k][tx + wc * (T/K)];
__syncthreads();
}
// Write out the T×T accumulated result
for (int wr = 0; wr < T; ++wr)
for (int wc = 0; wc < T; ++wc) {
int r = rowBase + wr * (T/K);
int c = colBase + wc * (T/K);
C[r * N + c] = acc[wr][wc];
}
}
＃ 結論
瞧！這就是通用矩陣乘法（GEMM）單元的基本運作原理！
接下來，我應該介紹一些特殊硬件，例如 Tensor Core（NVIDIA）或 Matrix Fused Multiply-Add（AMD）。這些是串流多處理器中的專用單元，用於計算 T*T*K 次累積乘法。我稍後會在備忘錄中詳細說明！
References:
DrawIO diagrams used in this memo: