如何優(yōu)化 CUDA 矩陣乘內(nèi)核以獲得類似 cuBLAS 的性能: 工作日志

• 注: 本文主要是對博文 “How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog - SIBOEHM” 的翻譯, 并進(jìn)行了一定的備注和補(bǔ)充

在這篇文章中, 我將迭代優(yōu)化用 CUDA 編寫的矩陣乘法的實(shí)現(xiàn). 我的目標(biāo)不是構(gòu)建一個(gè) cuBLAS 替代品, 而是深入了解用于現(xiàn)代深度學(xué)習(xí)的 GPU 的最重要的性能特征. 這包括合并全局內(nèi)存訪問、共享內(nèi)存緩存和占用優(yōu)化等¹².

GPU 上的矩陣乘法可能是目前存在的最重要的算法, 因?yàn)樗谟?xùn)練和推理大型深度學(xué)習(xí)模型時(shí)占據(jù)了幾乎所有的浮點(diǎn)運(yùn)算(FLOPs). 那么, 從頭開始編寫一個(gè)高性能的 CUDA SGEMM³ 需要多少工作呢? 我將從一個(gè)簡單的內(nèi)核開始, 并逐步應(yīng)用優(yōu)化, 直到我們達(dá)到了 cuBLAS (NVIDIA 的官方矩陣庫) 性能的 95% (在好的情況下) :⁴

Kernel 1: 樸素實(shí)現(xiàn)

在 CUDA 編程模型中, 計(jì)算按照三級(jí)層次進(jìn)行排序. CUDA 內(nèi)核的每次調(diào)用都會(huì)創(chuàng)建一個(gè)新的網(wǎng)格(grid), 該網(wǎng)格由多個(gè)線程塊(block)組成. 每個(gè)線程塊最多由 1024 個(gè)單獨(dú)的線程(thread)組成.⁵ 同一線程塊中的線程可以訪問同一共享內(nèi)存區(qū)域(SMEM).

一個(gè)線程塊中的線程數(shù)可以使用一個(gè)通常稱為 blockDim 的變量來配置, 該變量是一個(gè)由三個(gè) int 組成的向量. 該向量的條目指定了 blockDim.x、blockDim.y 和 blockDim.z 的大小, 如下所示:

類似地, 網(wǎng)格中的線程塊數(shù)可以使用 gridDim 變量進(jìn)行配置. 當(dāng)我們從主機(jī)⁶啟動(dòng)一個(gè)新內(nèi)核時(shí), 它會(huì)創(chuàng)建一個(gè)單獨(dú)的網(wǎng)格, 包含指定的線程塊和線程⁷. 重要的是要記住, 我們剛才討論的線程層次結(jié)構(gòu)主要關(guān)注程序的正確性. 對于程序性能, 正如我們稍后將看到的, 將同一線程塊中的所有線程等同看待并不是一個(gè)好主意.

對于我們的第一個(gè)內(nèi)核, 我們將使用網(wǎng)格、線程塊和線程的層次結(jié)構(gòu), 為每個(gè)線程分配結(jié)果矩陣 C 中的一個(gè)唯一條目. 然后, 該線程將計(jì)算矩陣 A 對應(yīng)行和矩陣 B 對應(yīng)列的點(diǎn)積, 并將結(jié)果寫入矩陣 C. 由于矩陣 C 的每個(gè)位置只由一個(gè)線程寫入, 因此我們不需要進(jìn)行同步操作. 我們將這樣啟動(dòng)內(nèi)核:

// create as many blocks as necessary to map all of C
dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32), 1);
// 32 * 32 = 1024 thread per block
dim3 blockDim(32, 32, 1);
// launch the asynchronous execution of the kernel on the device
// The function call returns immediately on the host
sgemm_naive<<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);

CUDA 代碼是從單線程的角度編寫的. 在內(nèi)核代碼中, 我們訪問 blockIdx 和 threadIdx 內(nèi)置變量. 這些將包含基于訪問它們的線程返回的不同值. ⁸⁹

__global__ void sgemm_naive(int M, int N, int K, float alpha, const float *A,
                            const float *B, float beta, float *C) {
  // compute position in C that this thread is responsible for
  const uint x = blockIdx.x * blockDim.x + threadIdx.x;
  const uint y = blockIdx.y * blockDim.y + threadIdx.y;

  // `if` condition is necessary for when M or N aren't multiples of 32.
  if (x < M && y < N) {
    float tmp = 0.0;
    for (int i = 0; i < K; ++i) {
      tmp += A[x * K + i] * B[i * N + y];
    }
    // C = α*(A@B)+β*C
    C[x * N + y] = alpha * tmp + beta * C[x * N + y];
  }
}

這個(gè)簡單內(nèi)核的可視化:¹⁰

這個(gè)內(nèi)核在我的 A6000 GPU 上處理三個(gè) 4092² 的 fp32 矩陣大約需要 0.5 秒. 讓我們進(jìn)行一些非特定于實(shí)現(xiàn)的計(jì)算:

最快運(yùn)行時(shí)的下限

對于兩個(gè) 4092² 的矩陣的矩陣相乘, 然后再加上一個(gè) 4092² 的矩陣 (以構(gòu)成 GEMM):

1. 總 FLOPS:¹¹: 2 * 4092³ + 4092² = 137 GFLOPS
2. 總讀取數(shù)據(jù)(最少!): 3 * 4092² * 4B = 201MB
3. 總存儲(chǔ)數(shù)據(jù): 4092² * 4B = 67MB

假設(shè)具有足夠大的緩存¹², 那么 268MB 是任何實(shí)現(xiàn)從/到全局 GPU 內(nèi)存¹³必須傳輸?shù)慕^對最小內(nèi)存量. 讓我們對內(nèi)核性能進(jìn)行一些上界計(jì)算. GPU 宣傳具有 30TFLOPs/s 的 FP32 計(jì)算吞吐量和 768GB/s 的全局內(nèi)存帶寬. 如果我們到達(dá)了這些數(shù)據(jù)指標(biāo)¹⁴¹⁵, 那么計(jì)算需要 4.5ms, 內(nèi)存?zhèn)鬏斝枰?0.34ms. 因此, 在我們的草稿計(jì)算中, 計(jì)算時(shí)間約為內(nèi)存訪問時(shí)間的約 10 倍. 這意味著只要我們需要傳輸?shù)膬?nèi)存量小于絕對最小值 278MB 的 10 倍, 我們的最終優(yōu)化的內(nèi)核將受到計(jì)算限制(compute-bound). ¹⁶

既然我們已經(jīng)為 fp32 GEMM 計(jì)算計(jì)算出了一些下限, 讓我們回到手頭的內(nèi)核, 看看為什么它比實(shí)際速度慢得多.

樸素內(nèi)核的內(nèi)存訪問模式

在我們的內(nèi)核中, 具有 ThreadId (0,0) 和 (0,1) 的同一線程塊中的兩個(gè)線程將加載 B 的同一列但 A 的不同行. 如果我們假設(shè)零緩存是最壞的情況, 那么每個(gè)線程必須從全局內(nèi)存加載 2*4092+1 個(gè)浮點(diǎn). 由于我們總共有 4092² 個(gè)線程, 這將導(dǎo)致 548GB 的內(nèi)存流量.
以下是我們樸素內(nèi)核的內(nèi)存訪問模式的可視化, 以兩個(gè)線程 A(紅色) 和 B(綠色) 為例:

因此, 概括一下, 當(dāng)我在 A6000 GPU 上運(yùn)行這個(gè)內(nèi)核時(shí), 當(dāng)乘以兩個(gè) 4092x4092 float32 的矩陣時(shí), 它達(dá)到了約 300 GFLOP. 考慮到 A6000 被宣傳為能夠?qū)崿F(xiàn)近 30 TFLOP, 這相當(dāng)糟糕.¹⁷ 那么, 我們?nèi)绾伍_始才能讓其更快呢? 一種方法是優(yōu)化內(nèi)核的內(nèi)存訪問模式, 使全局內(nèi)存訪問可以合并(coalesced) (=combined) 為更少的訪問.

Kernel 2: 全局內(nèi)存合并

在我們進(jìn)入全局內(nèi)存合并之前, 我們需要了解 warp 的概念. 為了便于執(zhí)行, 線程塊的線程被分組為由 32 個(gè)線程組成的所謂的 warp. 然后一個(gè) warp 被分配給 warp 調(diào)度器, 該調(diào)度器是執(zhí)行指令的物理核心.¹⁸ 每個(gè)多處理器(multiprocessor)有四個(gè) warp 調(diào)度器. 分組為 warp 是基于連續(xù)的 threadId 進(jìn)行的. 如果我們將 blockDim 設(shè)置為多維, 那么 threadId 的計(jì)算如下:

threadId = threadIdx.x+blockDim.x*(threadIdx.y+blockDim.y*threadIdx.z)

然后, 具有相鄰 threadId 的線程成為同一 warp 的一部分. 下面我試著用 8 線程的小“warp”來說明這一點(diǎn)(真正的 warp 總是包含 32 個(gè)線程): ¹⁹

warp 的概念與第二個(gè)內(nèi)核相關(guān), 因?yàn)閷儆谕?warp 的線程的順序內(nèi)存訪問可以被組合并作為一個(gè)整體執(zhí)行. 這被稱為全局內(nèi)存合并. . 這是在優(yōu)化內(nèi)核的全局內(nèi)存(GMEM)訪問以達(dá)到峰值帶寬時(shí)最重要的事情.
以下是一個(gè)示例, 其中對同一 warp 中線程的連續(xù)內(nèi)存訪問進(jìn)行分組, 允許每個(gè) warp 僅使用 2 個(gè) 32B 負(fù)載執(zhí)行 8 次內(nèi)存訪問:

實(shí)際上, GPU 支持 32B、64B 和 128B 的內(nèi)存訪問. 因此, 如果每個(gè)線程都從全局內(nèi)存加載一個(gè) 32 位浮點(diǎn), 那么 warp 調(diào)度器(可能是MIO)可以將這個(gè) 32*4B=128B 的加載合并到一個(gè)事務(wù)中. 只有當(dāng)加載的浮點(diǎn)數(shù)在內(nèi)存中是連續(xù)的, 并且訪問是對齊的時(shí)候, 這才有可能.²⁰ 如果不滿足, 或者由于其他原因無法合并訪問, 那么 GPU 將執(zhí)行盡可能多的 32B 負(fù)載來獲取所有浮點(diǎn), 從而導(dǎo)致大量帶寬浪費(fèi). 通過分析我們的樸素內(nèi)核, 我們可以觀察到非合并訪問的不利影響, 因?yàn)槲覀冎粚?shí)現(xiàn)了 15GB/s 的 GMEM 吞吐量.

回顧以前的內(nèi)核, 我們?yōu)榫€程分配了矩陣 C 的條目, 如下所示:

const uint x = blockIdx.x * blockDim.x + threadIdx.x;
const uint y = blockIdx.y * blockDim.y + threadIdx.y;

因此, 相同 warp 的線程(具有連續(xù) threadIdx.x 的線程)從內(nèi)存中非連續(xù)地加載 A 的行. 樸素內(nèi)核訪問 A 內(nèi)存的模式看起來更像這樣:

為了實(shí)現(xiàn)合并, 我們可以改變結(jié)果矩陣 C 分配給線程的位置的方式. 全局內(nèi)存訪問模式的這種變化如下所示:

要實(shí)現(xiàn)這一點(diǎn), 我們只需要更改前兩行:

const int x = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE);
const int y = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);

if (x < M && y < N) {
  float tmp = 0.0;
  for (int i = 0; i < K; ++i) {
    tmp += A[x * K + i] * B[i * N + y];
  }
  C[x * N + y] = alpha * tmp + beta * C[x * N + y];
}

我們這樣調(diào)用它:²¹

// gridDim stays the same
dim3 gridDim(CEIL_DIV(M, 32), CEIL_DIV(N, 32));
// make blockDim 1-dimensional, but don't change number of threads
dim3 blockDim(32 * 32);
sgemm_coalescing<<<gridDim, blockDim>>>(M, N, K, alpha, A, B, beta, C);

• 譯者注: 在上述代碼中, BLOCKSIZE 為線程塊的寬度 32. 如果將 Kernel 1 樸素內(nèi)核的 blockDim 同樣平鋪為一維, 則 x 和 y 的代碼如下. 可以看到, 前后二者的區(qū)別在于 % 和 / 的互換.

  const uint x = blockIdx.x * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);
  const uint y = blockIdx.y * BLOCKSIZE + (threadIdx.y / BLOCKSIZE);
  

  這樣便得到了上圖的效果. 即 Kernel 2 中, 連續(xù)的 32 個(gè)線程(即同一 warp 的線程)內(nèi) x 相同, 從而訪問矩陣 A 的同一行; 而 y 彼此不同但內(nèi)存連續(xù).
  而 Kernel 1 中, 連續(xù)的 32 個(gè)線程內(nèi) y 相同, 而 x 連續(xù), 但造成訪問矩陣 A 的不同行.

全局內(nèi)存合并將內(nèi)存吞吐量從 15GB/s 提高到 110GB/s. 性能達(dá)到 2000 GFLOPS, 與第一個(gè)樸素內(nèi)核的 300 GFLOPS 相比有了很大的改進(jìn). 對于下一個(gè)內(nèi)核, 我們將使用 GPU 的快速片上內(nèi)存, 稱為共享內(nèi)存(shared memory), 來緩存將被重復(fù)使用的數(shù)據(jù).

Kernel 3: 共享內(nèi)存緩存分塊

在較大的全局內(nèi)存旁邊, GPU 有一個(gè)小得多的內(nèi)存區(qū)域, 物理上位于芯片上, 稱為共享內(nèi)存(SMEM). 在物理上, 每個(gè) SM 有一個(gè)共享內(nèi)存.²² 從邏輯上講, 這個(gè)共享內(nèi)存是在線程塊之間劃分的. 這意味著線程可以通過共享內(nèi)存塊與其線程塊中的其他線程進(jìn)行通信. 在我的 A6000 GPU 上, 每個(gè)線程塊最多可以訪問 48KB 的共享內(nèi)存.²³

由于共享內(nèi)存位于片上, 因此它比全局內(nèi)存具有更低的延遲和更高的帶寬. 我找不到 Ampere 架構(gòu)的一個(gè)好的基準(zhǔn)測試結(jié)果, 但對于 Volta 架構(gòu)(2017 年發(fā)布), 這篇論文中進(jìn)行的的基準(zhǔn)測試報(bào)告了 750GiB/s 的全局內(nèi)存帶寬和 12080GiB/s 的共享內(nèi)存帶寬.²⁴

因此, 對于下一個(gè)內(nèi)核, 我們將從全局內(nèi)存中加載 A 的一個(gè)分塊和 B 的一個(gè)分塊到共享內(nèi)存中. 然后, 我們將在這兩個(gè)分塊上執(zhí)行盡可能多的工作, 每個(gè)線程仍被分配一個(gè) C 條目. 我們將沿著 A 的列和 B 的行移動(dòng)分塊, 對 C 執(zhí)行部分求和, 直到計(jì)算出結(jié)果.
如下所示:

代碼的重要部分如下, 變量名稱與上圖相對應(yīng):²⁵:

// advance pointers to the starting positions
A += cRow * BLOCKSIZE * K;                    // row=cRow, col=0
B += cCol * BLOCKSIZE;                        // row=0, col=cCol
C += cRow * BLOCKSIZE * N + cCol * BLOCKSIZE; // row=cRow, col=cCol

float tmp = 0.0;
// the outer loop advances A along the columns and B along
// the rows until we have fully calculated the result in C.
for (int bkIdx = 0; bkIdx < K; bkIdx += BLOCKSIZE) {
  // Have each thread load one of the elements in A & B from
  // global memory into shared memory.
  // Make the threadCol (=threadIdx.x) the consecutive index
  // to allow global memory access coalescing
  As[threadRow * BLOCKSIZE + threadCol] = A[threadRow * K + threadCol];
  Bs[threadRow * BLOCKSIZE + threadCol] = B[threadRow * N + threadCol];

  // block threads in this block until cache is fully populated
  __syncthreads();

  // advance pointers onto next chunk
  A += BLOCKSIZE;
  B += BLOCKSIZE * N;

  // execute the dotproduct on the currently cached block
  for (int dotIdx = 0; dotIdx < BLOCKSIZE; ++dotIdx) {
    tmp += As[threadRow * BLOCKSIZE + dotIdx] *
            Bs[dotIdx * BLOCKSIZE + threadCol];
  }
  // need to sync again at the end, to avoid faster threads
  // fetching the next block into the cache before slower threads are done
  __syncthreads();
}
C[threadRow * N + threadCol] =
    alpha * tmp + beta * C[threadRow * N + threadCol];

• 譯者注: 結(jié)合完整的 kernel 3 代碼可知:
  • cRow = blockIdx.x, cCol = blockIdx.y. for 循環(huán)之前的 A, B, C 對應(yīng)的是當(dāng)前線程塊處理的整個(gè)部分(上圖矩陣的淺色部分)的起始位置(左上角)的索引.
  • 外層 for 循環(huán)則是對線程塊處理的部分進(jìn)行分片迭代, 即 K 維度上的循環(huán)迭代. 在每次循環(huán)中, 將當(dāng)前線程塊分片加載到共享內(nèi)存中.
  • 內(nèi)層 for 循環(huán)是對當(dāng)前線程塊分片沿 K 維度逐一進(jìn)行點(diǎn)積(A行×B列), 然后寫入矩陣 C 對應(yīng)位置.
  • 由于 kernel 的 blockDim 為 32*32, 矩陣 C 的線程塊分片大小均為 32*32, 所以線程塊的每個(gè)線程恰好處理矩陣 C 的線程塊分片對應(yīng)的一個(gè)元素. 從矩陣 A 和 B 來看, 每次內(nèi)層 for 循環(huán)時(shí), 每個(gè)線程同時(shí)對應(yīng) A 分片 dotIdx 行中一個(gè)元素和 B 分片中 dotIdx 中一個(gè)元素.

該內(nèi)核實(shí)現(xiàn)了約 2200 GFLOPS, 比上一版本提高了 50%.²⁶ 我們離 GPU 所能提供的大約 30 TFLOP 還有很長的路要走. 從下面的屋頂狀圖中可以明顯看出:²⁷

在 CHUNKSIZE 為 32 時(shí), 這將使用 2*32*32*4B=8KB 的共享內(nèi)存空間.²⁸ 我的 A6000 GPU 每個(gè)線程塊最多有 48KB 的共享內(nèi)存空間, 所以我們還遠(yuǎn)遠(yuǎn)沒有達(dá)到這個(gè)極限. 這不一定是一個(gè)問題, 因?yàn)樵黾用總€(gè)線程塊的共享內(nèi)存使用量會(huì)有負(fù)面影響. 每個(gè)多處理器(SM)最多有 100KB 的 SMEM 可用. 這意味著, 如果我們修改內(nèi)核以使用完整的 48KB SMEM, 每個(gè) SM 只能同時(shí)加載兩個(gè)線程塊. 用 CUDA 的說法, 增加每個(gè)線程塊的 SMEM 利用率會(huì)降低占用率(occupancy). 占用率定義為每個(gè) SM 的激活 warp 數(shù)與每個(gè) SM 的最大可能激活 warp 數(shù)之間的比值.

高占用率是有用的, 因?yàn)樗试S我們通過擁有更大的可發(fā)布指令池來隱藏操作的高延遲.²⁹ 保持更多激活線程塊在 SM 上加載有三個(gè)主要限制: 寄存器數(shù)量、warp 數(shù)量和 SMEM 容量. 讓我們?yōu)槲覀儺?dāng)前的內(nèi)核進(jìn)行一個(gè)示例計(jì)算.

Kernel 3 的占有率計(jì)算

以下是從 cudaGetDeviceProperties API 獲得的我的 GPU 的相關(guān)硬件統(tǒng)計(jì)數(shù)據(jù)(多處理器是我們之前討論的 SM):³⁰

以下是我們內(nèi)核的資源需求:

工作負(fù)載按線程塊粒度在 SM 上調(diào)度. 每個(gè) SM 將盡可能加載更多的線程塊, 只要它有足夠的資源來容納它們. 計(jì)算:³¹

• 共享內(nèi)存: 8192B/Block + 1024B/Block 用于 CUDA 運(yùn)行時(shí)使用 = 9216B/Block. (每 SM 102400B) / (每線程塊 9216B per Block) = 11.11 ? 上限 11 個(gè)線程塊.
• 線程: 每線程塊 1024 個(gè)線程, 每個(gè) SM 最多 1536 個(gè)線程 ? 上限 1 個(gè)線程塊.
• 寄存器: 每線程 37 個(gè)寄存器 * 每 warp 32 線程 = 每 warp 1184 個(gè)寄存器. 寄存器分配粒度在 warp 級(jí)別上為 256個(gè)寄存器, 因此每個(gè) warp 四舍五入到 1280 個(gè)寄存器. 我們有 (1024 線程 / 32) = 每個(gè)線程塊 32 個(gè) warp, 因此每個(gè) warp 1280 個(gè)寄存器 * 每個(gè)線程塊 32 warp = 每個(gè)線程塊 40960 個(gè)寄存器. 每 SM 最多 65536 個(gè)寄存器 ? 上限 1 個(gè)線程塊.³²

因此, 這個(gè)內(nèi)核受到每個(gè)線程塊的線程數(shù)和每個(gè)線程的寄存器數(shù)的限制. 我們不能為每個(gè) SM 加載一個(gè)以上的線程塊, 因此最終占用率為 32 個(gè)有效 warp/48 個(gè)最大有效 warp = 66%.

66% 的占用率還不錯(cuò), 所以這并不能解釋為什么我們的內(nèi)核運(yùn)行如此緩慢.³³ 查看監(jiān)測器(profiler)會(huì)給我們一些提示. 首先, 如果我們觀察已執(zhí)行指令的組合, 它們中的大多數(shù)都是內(nèi)存加載:³⁴

我們的內(nèi)存循環(huán)在 PTX(Godbolt 鏈接)中看起來是這樣的:

ld.shared.f32   %f91, [%r8+3456];
ld.shared.f32   %f92, [%r7+108];
fma.rn.f32      %f93, %f92, %f91, %f90;

這并不夠好, 因?yàn)閮?nèi)存負(fù)載必然比簡單的 FMA 具有更高的延遲, 并且我們知道我們的內(nèi)核應(yīng)該是計(jì)算受限的. 我們在查看 profiler 對 warp 狀態(tài)的采樣時(shí)會(huì)看到這種效果. 這量化了每個(gè)執(zhí)行指令在每個(gè)狀態(tài)下花費(fèi)的時(shí)鐘周期數(shù):³⁵

狀態(tài)的含義在《內(nèi)核評(píng)測指南》中有詳細(xì)說明. 對于 Stall MIO Throttle, 它顯示為:

warp 在等待 MIO(memory input/output, 內(nèi)存輸入/輸出)指令隊(duì)列未滿時(shí)被阻塞. 這種阻塞原因在極端利用 MIO 流水線的情況下很高, 其中包括特殊的數(shù)學(xué)指令、動(dòng)態(tài)分支以及共享內(nèi)存指令. 

我們沒有使用特殊的數(shù)學(xué)指令, 也沒有使用動(dòng)態(tài)分支, 所以很明顯, 我們正在等待 SMEM 訪問返回. 那么, 我們?nèi)绾问箖?nèi)核發(fā)出更少的 SMEM 指令呢? 一種方法是讓每個(gè)線程計(jì)算一個(gè)以上的輸出元素, 這使我們能夠在寄存器中執(zhí)行更多的工作, 并減少對 SMEM 的依賴.

Kernel 4: 用于計(jì)算每個(gè)線程的多個(gè)結(jié)果的一維塊分片

因此, 下一個(gè)內(nèi)核的工作方式與上一個(gè)內(nèi)核類似, 但添加了一個(gè)新的內(nèi)部循環(huán), 用于計(jì)算每個(gè)線程的多個(gè) C 條目. 我們現(xiàn)在使用 SMEM 緩存 BM*BK+BN*BK=64*8+64*8=1024 個(gè)浮點(diǎn)數(shù)大小, 每個(gè)線程塊總共 4KB. 可視化如下所示. 我已經(jīng)用橙色和紅色突出顯示了其中的兩個(gè)線程以及它們在內(nèi)部循環(huán)中訪問的值.

這個(gè)內(nèi)核的所有重要更改都發(fā)生在內(nèi)部循環(huán)中. GMEM 到 SMEM 的負(fù)載與以前基本相同. 讓我們看看:³⁶

// allocate thread-local cache for results in registerfile
float threadResults[TM] = {0.0};

// outer loop over block tiles
for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
  // populate the SMEM caches (same as before)
  As[innerRowA * BK + innerColA] = A[innerRowA * K + innerColA];
  Bs[innerRowB * BN + innerColB] = B[innerRowB * N + innerColB];
  __syncthreads();

  // advance blocktile for outer loop
  A += BK;
  B += BK * N;

  // calculate per-thread results
  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
    // we make the dotproduct loop the outside loop, which facilitates
    // reuse of the Bs entry, which we can cache in a tmp var.
    float Btmp = Bs[dotIdx * BN + threadCol];
    for (uint resIdx = 0; resIdx < TM; ++resIdx) {
      threadResults[resIdx] +=
          As[(threadRow * TM + resIdx) * BK + dotIdx] * Btmp;
    }
  }
  __syncthreads();
}

• 譯者注: 結(jié)合完整的 kernel 4 代碼可知:
  • BM=BN=64, BK=8, TM=8, blockDim=(BM*BN)/TM=64*64/8.
  • TM=8 為每個(gè)線程在一次 K 維度迭代中計(jì)算的元素個(gè)數(shù), 之前 Kernel 3 中相當(dāng)于 TM=1.
  • innerColA = threadIdx.x % BK, innerRowA = threadIdx.x / BK, innerColB = threadIdx.x % BN, innerRowB = threadIdx.x / BN; 且滿足 BM * BK == blockDim.x 和 BN * BK == blockDim.x.
    即這四個(gè)變量將線程塊中的線程映射到矩陣 A 和 B 的線程塊分片(上圖矩陣中的深色部分)的每個(gè)元素(由于滿足等式關(guān)系, 線程塊的線程可以恰好鋪滿矩陣 A 和 B 的分片), 用于將數(shù)據(jù)從全局內(nèi)存拷貝到共享內(nèi)存.
  • 由于矩陣 C 的分片大小為 BM*BN, 實(shí)際大小為線程塊中線程數(shù)的 TM 倍. 從另一個(gè)角度也證明了每個(gè)線程需要處理 TM 個(gè)元素, 需要 threadResults 大小為 TM.
  • 內(nèi)層 for dotIdx 循環(huán)仍然是沿 A 和 B 分片的 K 維逐一計(jì)算, 即每次 As 處理一列, Bs 處理一行. 在這里實(shí)際上矩陣乘法由之前的矩陣內(nèi)積變?yōu)榱送夥e.
  • 最內(nèi)層 for resIdx 循環(huán)用于線程內(nèi)部迭代處理中 TM 個(gè)元素.
    由于 threadCol = threadIdx.x % BN, threadRow = threadIdx.x / BN. 即相鄰線程對應(yīng) B 分片的不同列, A 分片的同一行分塊(因?yàn)槊總€(gè)線程處理 TM 個(gè)元素), 與 Kernel 3 相同.
    值得一提的是, 由于除數(shù)是 BN, 所以在一次內(nèi)部循環(huán)迭代中, Bs 分片的 dotIdx 行的不同列元素(BN 個(gè))都會(huì)被緩存到 Btmp 中; 而 As 分片一次只能處理 dotIdx 列的 blockDim/BN 行個(gè)元素.
    根據(jù) threadResults[resIdx] += As[(threadRow * TM + resIdx) * BK + dotIdx] * Btmp; 可知, 這里每個(gè)線程處理的 TM 個(gè)元素在 A 分片的 dotIdx 列中是連續(xù)的. 即最內(nèi)層循環(huán)可以理解為將 As 分片沿 BM 維度再分成了 BM/TM 個(gè)小分片(線程分片), 每個(gè)小分片由 1 個(gè)線程處理(一個(gè)小分片對應(yīng)多個(gè)線程, 它們拿 A 分片相同行元素, B 分片不同列元素).

這個(gè)內(nèi)核實(shí)現(xiàn)了約 8600 GFLOP, 比我們以前的內(nèi)核快 2.2 倍. 讓我們計(jì)算一下在我們前一個(gè)內(nèi)核中每個(gè)線程執(zhí)行的內(nèi)存訪問次數(shù), 其中每個(gè)線程計(jì)算一個(gè)結(jié)果:

• GMEM: 外部循環(huán) K/32 次迭代 * 2 次加載
• SMEM: 外部循環(huán) K/32 次迭代 * BLOCKSIZE(=32) * 2 次加載 (譯者注: 這里的 “BLOCKSIZE” 是線程塊分片 K 維度)
• 每個(gè)結(jié)果的內(nèi)存訪問: K/16 GMEM, K*2 SMEM

對于我們的新內(nèi)核, 每個(gè)線程計(jì)算八個(gè)結(jié)果:

• GMEM: 外部循環(huán) K/8 次迭代 * 2 次加載
• SMEM: 外部循環(huán) K/8 次迭代 * BK(=8) * (1 + TM(=8)) (譯者注: B 分片 1 次, A 分片 TM 次)
• 每個(gè)結(jié)果的內(nèi)存訪問: K/32 GMEM, K * 9/8 MEM (譯者注: 這里是上面兩個(gè)計(jì)算結(jié)果再除以每個(gè)線程處理的 TM=8 個(gè)元素得到的每個(gè)結(jié)果的內(nèi)存訪問均攤結(jié)果)

正如預(yù)期的那樣, 由于內(nèi)存壓力, 我們現(xiàn)在每次指令停滯所花費(fèi)的時(shí)鐘周期要少得多:³⁷

編譯優(yōu)化側(cè)記

在上面, 我們明確地將 B 的條目緩存到 Btmp 中, 并為了提高效率對兩個(gè)內(nèi)部循環(huán)進(jìn)行了重新排序. 如果我們不這樣做, 那么代碼看起來是這樣的:

for (uint resIdx = 0; resIdx < TM; ++resIdx) {
  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
    threadResults[resIdx] +=
      As[(threadRow * TM + resIdx) * BK + dotIdx] * Bs[dotIdx * BN + threadCol];
  }
}

有趣的是, 這對性能沒有負(fù)面影響. 這是令人驚訝的, 因?yàn)槲覀兊膬?nèi)部兩個(gè)循環(huán)現(xiàn)在產(chǎn)生 BK(=8) * TM(=8) = 128 次 SMEM 訪問, 而非之前的 72 次. 看看匯編(Godbolt 鏈接)就知道答案了:

// first inner-most loop
ld.shared.f32   %f45, [%r9];
ld.shared.f32   %f46, [%r8];
fma.rn.f32      %f47, %f46, %f45, %f212;
ld.shared.f32   %f48, [%r9+256];
ld.shared.f32   %f49, [%r8+4];
fma.rn.f32      %f50, %f49, %f48, %f47;
ld.shared.f32   %f51, [%r9+512];
ld.shared.f32   %f52, [%r8+8];
fma.rn.f32      %f53, %f52, %f51, %f50;
ld.shared.f32   %f54, [%r9+768];
ld.shared.f32   %f55, [%r8+12];
fma.rn.f32      %f56, %f55, %f54, %f53;
ld.shared.f32   %f57, [%r9+1024];
ld.shared.f32   %f58, [%r8+16];
fma.rn.f32      %f59, %f58, %f57, %f56;
ld.shared.f32   %f60, [%r9+1280];
ld.shared.f32   %f61, [%r8+20];
fma.rn.f32      %f62, %f61, %f60, %f59;
ld.shared.f32   %f63, [%r9+1536];
ld.shared.f32   %f64, [%r8+24];
fma.rn.f32      %f65, %f64, %f63, %f62;
ld.shared.f32   %f66, [%r9+1792];
ld.shared.f32   %f67, [%r8+28];
fma.rn.f32      %f212, %f67, %f66, %f65;
// second inner-most loop
ld.shared.f32   %f68, [%r8+32];
fma.rn.f32      %f69, %f68, %f45, %f211;
ld.shared.f32   %f70, [%r8+36];
fma.rn.f32      %f71, %f70, %f48, %f69;
ld.shared.f32   %f72, [%r8+40];
fma.rn.f32      %f73, %f72, %f51, %f71;
ld.shared.f32   %f74, [%r8+44];
fma.rn.f32      %f75, %f74, %f54, %f73;
ld.shared.f32   %f76, [%r8+48];
fma.rn.f32      %f77, %f76, %f57, %f75;
ld.shared.f32   %f78, [%r8+52];
fma.rn.f32      %f79, %f78, %f60, %f77;
ld.shared.f32   %f80, [%r8+56];
fma.rn.f32      %f81, %f80, %f63, %f79;
ld.shared.f32   %f82, [%r8+60];
fma.rn.f32      %f211, %f82, %f66, %f81;
// ... continues like this for inner-loops 3-8 ...

編譯器展開兩個(gè)循環(huán)³⁸, 然后消除 Bs 條目的重復(fù) SMEM 加載, 因此我們最終獲得的 SMEM 訪問量與優(yōu)化的 CUDA 代碼相同.

當(dāng) PTX 被編譯為 SASS 時(shí), 來自 Bs 的 SMEM 加載被向量化:³⁹

LDS     R26, [R35.X4+0x800] // a 32b load from As
LDS.128 R8,  [R2]           // a 128b load from Bs
LDS.128 R12, [R2+0x20] 
LDS     R24, [R35.X4+0x900] 
LDS.128 R20, [R2+0x60] 
LDS     R36, [R35.X4+0xb00] 
LDS.128 R16, [R2+0x40] 
LDS.128 R4,  [R2+0x80] 
LDS     R38, [R35.X4+0xd00] 

改進(jìn)方面: 計(jì)算強(qiáng)度 (Arithmetic Intensity)

我們當(dāng)前的內(nèi)核仍然存在與 Kernel 3 相同的內(nèi)存停滯(stalling-for-memory)問題, 只是程度較低. 因此, 我們將再次應(yīng)用相同的優(yōu)化: 讓每個(gè)線程計(jì)算更多的結(jié)果. 使我們的內(nèi)核運(yùn)行得更快的主要原因是它增加了計(jì)算強(qiáng)度.⁴⁰ 下面, 我試圖更清楚地說明為什么每個(gè)線程計(jì)算更多的結(jié)果會(huì)提高運(yùn)算強(qiáng)度:⁴¹

總之, 我們所有的內(nèi)核執(zhí)行相同數(shù)量的 FLOP, 但我們可以通過每個(gè)線程計(jì)算更多的結(jié)果來減少 GMEM 訪問的數(shù)量. 只要我們?nèi)匀皇莾?nèi)存限制的(memory bound), 我們就會(huì)繼續(xù)優(yōu)化計(jì)算強(qiáng)度.

Kernel 5: 通過二維塊分片增加計(jì)算強(qiáng)度

Kernel 5 的基本思想是每個(gè)線程計(jì)算 C 的 8 * 8 個(gè)元素的網(wǎng)格. 內(nèi)核的第一階段是讓所有線程一起工作來填充 SMEM 緩存. 我們將讓每個(gè)線程加載多個(gè)元素. 此代碼如下所示:⁴²

for (uint loadOffset = 0; loadOffset < BM; loadOffset += strideA) {
  As[(innerRowA + loadOffset) * BK + innerColA] =
      A[(innerRowA + loadOffset) * K + innerColA];
}
for (uint loadOffset = 0; loadOffset < BK; loadOffset += strideB) {
  Bs[(innerRowB + loadOffset) * BN + innerColB] =
      B[(innerRowB + loadOffset) * N + innerColB];
}
__syncthreads();

現(xiàn)在 SMEM 緩存已經(jīng)填充完畢, 我們讓每個(gè)線程將其相關(guān)的 SMEM 條目相乘, 并將結(jié)果累積到本地寄存器中. 下面我展示了沿著輸入矩陣的(沒有改變)外層循環(huán), 以及點(diǎn)積和 TN 和 TM 維度的三個(gè)內(nèi)層循環(huán):

代碼中有趣的部分如下所示:⁴³

// allocate thread-local cache for results in registerfile
float threadResults[TM * TN] = {0.0};
// register caches for As and Bs
float regM[TM] = {0.0};
float regN[TN] = {0.0};

// outer-most loop over block tiles
for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
  // populate the SMEM caches
  for (uint loadOffset = 0; loadOffset < BM; loadOffset += strideA) {
    As[(innerRowA + loadOffset) * BK + innerColA] =
        A[(innerRowA + loadOffset) * K + innerColA];
  }
  for (uint loadOffset = 0; loadOffset < BK; loadOffset += strideB) {
    Bs[(innerRowB + loadOffset) * BN + innerColB] =
        B[(innerRowB + loadOffset) * N + innerColB];
  }
  __syncthreads();

  // advance blocktile
  A += BK;     // move BK columns to right
  B += BK * N; // move BK rows down

  // calculate per-thread results
  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
    // load relevant As & Bs entries into registers
    for (uint i = 0; i < TM; ++i) {
      regM[i] = As[(threadRow * TM + i) * BK + dotIdx];
    }
    for (uint i = 0; i < TN; ++i) {
      regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
    }
    // perform outer product on register cache, accumulate
    // into threadResults
    for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
      for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
        threadResults[resIdxM * TN + resIdxN] +=
            regM[resIdxM] * regN[resIdxN];
      }
    }
  }
  __syncthreads();
}

• 譯者注: Kernel 5 和 Kernel 4 的整體思路一致, 核心在于將一維分片改成了二維分片. 換句話說, 從 “僅將 As 的線程塊分片進(jìn)一步劃分為線程分片” 到 “As 和 Bs 的線程塊分片均進(jìn)一步劃分為線程分片”.
  根據(jù)完整的 Kernel 5 代碼可知:
  • 由于每個(gè)線程現(xiàn)在處理 TM*TN 個(gè)元素, 因此 As 分片和 Bs 分片需要多次加載才能將其元素加載至共享內(nèi)存.
  • 由于 threadCol = threadIdx.x % (BN / TN), threadRow = threadIdx.x / (BN / TN). 即相鄰線程對應(yīng) Bs 分片的不同列線程分片、As 分片的同一行線程分塊, 與先前 Kernel 的思路是一致的.
  • for dotIdx 循環(huán)仍然是線程塊分片沿 K 維度逐一計(jì)算, 即每次 As 處理一列, Bs 處理一行.
  • 兩個(gè) for i 循環(huán)將重復(fù)使用的線程分片元素加載至寄存器. 因?yàn)?Kernel 4 是一維塊分片, 相當(dāng)于 TN=1, 這樣 regN 就退化為了 Kernel 4 中的 Btmp, 而由于 TN=1, As 的線程分片不會(huì)被重復(fù)利用, 因此就沒有 regM.
  • 內(nèi)層 for resIdxM 和 for resIdxN 循環(huán)計(jì)算線程塊分片的結(jié)果. 即依次迭代 As 分片 dotIdx 列的 TM 個(gè)元素和 Bs 分片 dotIdx 行的 TN 個(gè)元素, 計(jì)算總共 TM*TN 個(gè)值.

在內(nèi)層循環(huán)中, 我們可以通過將 dotIdx 作為外層循環(huán), 并將兩個(gè)內(nèi)層循環(huán)所需的值顯式加載到寄存器中, 來減少對 SMEM 的訪問次數(shù). 以下是隨時(shí)間推移的 dotIdx 循環(huán)的圖, 以直觀地顯示在每個(gè)步驟中哪些 SMEM 條目加載到線程本地寄存器中:⁴⁴

最終性能: 16TFLOP, 又提高了 2 倍. 讓我們重復(fù)內(nèi)存訪問的計(jì)算. 我們現(xiàn)在每個(gè)線程計(jì)算 TM*TN=8*8=64 個(gè)結(jié)果.

• GMEM: K/8 (外層循環(huán)迭代) * 2 (A+B) * 1024/256 (sizeSMEM/numThreads) 次加載 (譯者注: 每個(gè)線程塊分片的 SMEM 大小為 BM*BK = BN*BK = 128*8 = 1024, 線程塊大小為 (BM*BN)/(TM*TN) = (128*128)/(8*8)=256)
• SMEM: K/8 (外層循環(huán)迭代) * 8 (dotIdx) * 2 (A+B) * 8 次加載 (譯者注: 8 次是因?yàn)槊總€(gè)線程對應(yīng) A 和 B 的線程分片大小的 TM = TN = 8)
• 每個(gè)結(jié)果的內(nèi)存訪問: K/64 GMEM, K/4 SMEM (譯者注: 上面兩個(gè)計(jì)算結(jié)果再除以每個(gè)線程處理的 TM*TN=8*8=64 個(gè)元素)

性能慢慢地達(dá)到了可以接受的水平, 但是, 由于內(nèi)存流水線(memory pipeline)擁塞而導(dǎo)致的 warp 停滯仍然過于頻繁. 對于 Kernel 6, 我們將采取兩種措施來嘗試改進(jìn)這一點(diǎn): 轉(zhuǎn)置 As 以實(shí)現(xiàn) SMEM 加載的自動(dòng)向量化, 以及向編譯器承諾在 GMEM 訪問時(shí)的內(nèi)存對齊.

Kernel 6: 向量化 SMEM 和 GMEM 訪問

第一個(gè)優(yōu)化是我之前已經(jīng)暗示的轉(zhuǎn)置 As. 這將允許我們從 As 加載時(shí)使用向量化 SMEM 加載(SASS 中的 LDS.128). 下面是與 Kernel 5 相同的三個(gè)內(nèi)部循環(huán)的可視化, 但現(xiàn)在在內(nèi)存中轉(zhuǎn)置了 As:

查看匯編⁴⁵我們可以看到, 將 As 加載到寄存器過去是用的 32b 的 LDS 加載負(fù)載, 而現(xiàn)在是使用 128b 的 LDS.128 加載, 就像之前對 Bs 加載一樣. 這使我們的速度提高了 500GFLOP, 或約 3%.

接下來, 我們將使用向量數(shù)據(jù)類型(即 float4)對 GMEM 的所有加載和存儲(chǔ)進(jìn)行向量化.
代碼如下所示:⁴⁶

float4 tmp =
    reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
// transpose A during the GMEM to SMEM transfer
As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;

reinterpret_cast<float4 *>(&Bs[innerRowB * BN + innerColB * 4])[0] =
    reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
__syncthreads();

這導(dǎo)致 32b GMEM 加載指令( LDG.E 和 STG.E)被 128b 對應(yīng)指令(LDG.E.128 和 STG.E.128)所取代. 起初, 我很困惑為什么要運(yùn)行此指令:

reinterpret_cast<float4 *>(&Bs[innerRowB * BN + innerColB * 4])[0] =
    reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];

這將比只手動(dòng)展開訪問(或使用 pragma unroll)快嗎:

Bs[innerRowB * BN + innerColB * 4 + 0] = B[innerRowB * N + innerColB * 4 + 0];
Bs[innerRowB * BN + innerColB * 4 + 1] = B[innerRowB * N + innerColB * 4 + 1];
Bs[innerRowB * BN + innerColB * 4 + 2] = B[innerRowB * N + innerColB * 4 + 2];
Bs[innerRowB * BN + innerColB * 4 + 3] = B[innerRowB * N + innerColB * 4 + 3];

難道編譯器不應(yīng)該只合并第二個(gè)版本并生成 128b 負(fù)載嗎? 我認(rèn)為原因是編譯器無法驗(yàn)證傳遞給內(nèi)核的 float* B指針是否與 128b 對齊, 這是使用 LDG.E.128 的要求. 因此, interpret_cast 的唯一目的是向編譯器承諾 float* B指針是對齊的.⁴⁷

Kernel 6 實(shí)現(xiàn)了 19TFLOP. profiler 仍然顯示了一系列問題領(lǐng)域和優(yōu)化機(jī)會(huì): 我們遇到了共享內(nèi)存的 bank 沖突 (bank conflict)(cuBLAS 避免了這一點(diǎn)), 我們的占用率高于必要水平, 而且我們還沒有實(shí)現(xiàn)任何雙緩沖(CUTRASS 文檔似乎認(rèn)為這非常有用).
但在我們討論這些之前, 讓我們討論一些懸而未決的問題: 自動(dòng)調(diào)整內(nèi)核的參數(shù).

Kernel 7: 通過線性化解決 bank conflict (譯者補(bǔ)充)

在原作者的 Kernel 7 中, 通過對矩陣 B 的線程塊分片加載到 SMEM 時(shí)進(jìn)行了 swizzle 操作, 避免了從 Bs 加載到 regN 時(shí)的 bank conflict. (實(shí)際上只解決了讀取共享內(nèi)存的 bank conflict, 寫入共享內(nèi)存的 bank conflict 仍然存在.)

代碼的核心變動(dòng)為以下兩處:

首先是 Bs 加載時(shí):

    float4 tmp =
        reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
    As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
    As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
    As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
    As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;

    reinterpret_cast<float4 *>(&Bs[innerRowB * BN + innerColB * 4])[0] =
        reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];

改為:

    // populate the SMEM caches
    // transpose A while loading it
    float4 tmp =
        reinterpret_cast<float4 *>(&A[innerRowA * K + innerColA * 4])[0];
    As[(innerColA * 4 + 0) * BM + innerRowA] = tmp.x;
    As[(innerColA * 4 + 1) * BM + innerRowA] = tmp.y;
    As[(innerColA * 4 + 2) * BM + innerRowA] = tmp.z;
    As[(innerColA * 4 + 3) * BM + innerRowA] = tmp.w;

    // "linearize" Bs while storing it
    tmp = reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 0) * 16 + innerColB / 2] = tmp.x;
    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 1) * 16 + innerColB / 2] = tmp.y;
    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 2) * 16 + innerColB / 2] = tmp.z;
    Bs[((innerColB % 2) * 4 + innerRowB * 8 + 3) * 16 + innerColB / 2] = tmp.w;

然后:

      // block into registers
      for (uint i = 0; i < TM; ++i) {
        regM[i] = As[dotIdx * BM + threadRow * TM + i];
      }
      for (uint i = 0; i < TN; ++i) {
        regN[i] = Bs[dotIdx * BN + threadCol * TN + i];
      }

改為:

      // block into registers
      for (uint i = 0; i < TM; ++i) {
        regM[i] = As[dotIdx * BM + threadRow * TM + i];
      }
      for (uint i = 0; i < TN; ++i) {
        regN[i] = Bs[(dotIdx * 8 + i) * 16 + threadCol];
      }

在此之前, 分析一下為什么讀取共享內(nèi)存時(shí) Bs 會(huì)有 bank conflict 而 As 沒有.
關(guān)鍵在于以下兩個(gè)變量:

  const int threadCol = threadIdx.x % (BN / TN);
  const int threadRow = threadIdx.x / (BN / TN);

由于 BN/TN = 128/8 = 16, 所以在一個(gè) warp 中的 32 個(gè)線程中, 每 16 個(gè)線程有相同的 threadRow 和不同的 threadCol.

• 對于 Kernel 6 的 As[dotIdx * BM + threadRow * TM + i] 讀取, 在 dotIdx 和 i 固定時(shí), warp 中前 16 個(gè)線程訪問同一地址, 后 16 個(gè)線程訪問另一相同的地址, 兩個(gè)地址間差 TM 即 8 個(gè)元素, 也就是 8 個(gè) bank, 因此會(huì)通過廣播機(jī)制使得沒有 bank conflict 發(fā)生.
• 對于 Kernel 6 的 Bs[dotIdx * BN + threadCol * TN + i] 讀取, 情況則不同, 對于 warp 的前 16 個(gè)線程, 由于 threadCol 彼此不同, 所以每個(gè)相鄰線程訪問的元素地址差 TN, 即 8 個(gè) bank. 這樣導(dǎo)致 threadIdx 差 4 的線程訪問的元素彼此相差 32, 從而造成了 bank conflict.

作者做出的改動(dòng)即對原本的 Bs 線程塊分片進(jìn)行了 ROW=8, COL=16 的 swizzle, 可視化如下圖所示:

圖中, 每個(gè)小顏色塊對應(yīng)一個(gè)線程, 每次從 GMEM 讀取 4 個(gè)元素到 SMEM. 紅色和橙色方框?qū)?yīng)線程塊分片的一行(在后續(xù)計(jì)算時(shí), 有相同的 dotIdx). 所有尺寸的單位是 1 個(gè)元素即 float, 對應(yīng) 32 字節(jié) 1 個(gè) bank.

修改后, 讀取 Bs 變?yōu)榱?Bs[(dotIdx * 8 + i) * 16 + threadCol], 在同一 warp 中, 前 16 個(gè)線程讀取的元素彼此差 1, 即對應(yīng) 16 個(gè)不同的 bank. 而后 16 個(gè)線程與前 16 個(gè)線程訪問的地址相同. 因此, 解決了 bank conflict. (在上圖, 16 個(gè)線程讀取的 Bs 分別對應(yīng) ROW=8, COL=16 的 swizzle 分塊的每一列的 8 個(gè)元素, 即之前兩個(gè)線程讀取的元素, 對應(yīng)一列的兩個(gè)顏色塊.)

這里選擇 ROW=8, COL= 16 的原因是:
在 Bs 的讀取階段, 每個(gè)線程讀取 8 個(gè)元素, 每 BN/TN=128/8=16 個(gè)線程讀取相同的 128 個(gè)元素. 在上圖中, ROW 恰好對應(yīng)連續(xù)的兩個(gè)小顏色塊(即之前兩個(gè)相鄰線程讀取)的元素. COL=16 使得在計(jì)算的時(shí)候, 相鄰線程訪存的元素是相鄰的.
(理論上, COL 為 16 的倍數(shù)均可以避免 bank conflict, 因?yàn)?16 個(gè)線程每次迭代至多訪問 16 個(gè)元素對應(yīng) 16 bank, 但 COL 增大等會(huì)導(dǎo)致 ROW 減小, 致使相鄰線程訪問的元素不連續(xù), 影響性能.)

值得一提的是, 作者在這里并沒有解決 tmp 到 As 和 Bs 即共享內(nèi)存讀取時(shí)的 bank conflict:

• 對于 As[(innerColA * 4 + 0) * BM + innerRowA] 寫入, 同一 warp 中相鄰的兩個(gè)線程 innerRowA 相同, innerColA 差 1, 訪問的元素差 4*BM=4*128=512, 對應(yīng)同一個(gè) bank. 這樣同一 warp 中 32 個(gè)線程訪問 16 個(gè) bank 每個(gè)都有 1 個(gè) bank conflict.
• 對于 kernel 6 的 einterpret_cast<float4 *>(&Bs[innerRowB * BN + innerColB * 4])[0] 向量化寫入, 同一 warp 中 32 個(gè)線程 innerRowB 相同, innerColB 彼此差 1, 因此 threadIdx 差 8 的線程訪問元素差 32 個(gè).
  這里看似有 bank conflict, 但是并沒有! 這里譯者通過把 As 寫入的部分注釋掉再用 NVProf 檢測發(fā)現(xiàn)的.
  但如果去掉向量化寫入改為 Bs[innerRowB * BN + innerColB * 4] = B[innerRowB * N + innerColB * 4], 便會(huì)檢測出 bank conflict.
  筆者推斷應(yīng)該是 CUDA 對向量化訪問進(jìn)行了一定的優(yōu)化, 避免了 bank conflict. (我猜測是因?yàn)槿绻麤]有額外優(yōu)化, 向量化讀取數(shù)據(jù)到共享內(nèi)存必然會(huì)產(chǎn)生 bank conflict: 一個(gè)線程讀 4 個(gè)大小為 4 字節(jié)的元素, 那么 threadIdx 差 8 的線程向量化的首元素差 32, 必然有 bank conflict, 所以 CUDA 對此進(jìn)行了優(yōu)化.)
• 對于 kernel 7 的 Bs[((innerColB % 2) * 4 + innerRowB * 8 + 0) * 16 + innerColB / 2] 寫入, threadIdx 差 1 的線程訪問元素差 4*16=64 個(gè), 有 bank conflict.

以下是 kernel 6 和 kernel 7 使用 nvprof 檢測得到的 bank conflict 次數(shù), 可以看到讀取的 bank conflict 最少的情況下減少到 0; 而寫入的 bank conflict 由于 Bs 寫入操作的改變增加了 bank conflict.
但是相比與讀取 conflict 的大幅減少, 讀取增加的 conflict 相對更少, 整體上是正收益, 因此性能有一定提升.

Invocations                                Event Name         Min         Max         Avg       Total
Device "Tesla V100-SXM2-32GB (1)"
    Kernel: void sgemmVectorize<int=128, int=128, int=8, int=8, int=8>(int, int, int, float, float*, float*, float, float*)
        255                   shared_ld_bank_conflict        8192    33609822     7680052  1958413308
        255                   shared_st_bank_conflict         540     2384733      539888   137671629

Invocations                                Event Name         Min         Max         Avg       Total
Device "Tesla V100-SXM2-32GB (1)"
    Kernel: void sgemmResolveBankConflicts<int=128, int=128, int=8, int=8, int=8>(int, int, int, float, float*, float*, float, float*)
        255                   shared_ld_bank_conflict           0       33181        5656     1442492
        255                   shared_st_bank_conflict        1172     4731245     1079426   275253689

Kernel 8: 通過偏移解決 bank conflict (譯者補(bǔ)充)

Kernel 8 相比 kernel 6 的變化也是兩個(gè)地方:

  const int extraCols = 5;
  __shared__ float Bs[BK * (BN + extraCols)];
  
  // ...
  
    tmp = reinterpret_cast<float4 *>(&B[innerRowB * N + innerColB * 4])[0];
    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 0] = tmp.x;
    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 1] = tmp.y;
    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 2] = tmp.z;
    Bs[innerRowB * (BN + extraCols) + innerColB * 4 + 3] = tmp.w;

      for (uint i = 0; i < TN; ++i) {
        regN[i] = Bs[dotIdx * (BN + extraCols) + threadCol * TN + i];
      }

代碼的核心是線程塊分片 Bs 的寬度從 BN 變?yōu)榱?BN+extraCols. 譯者猜測應(yīng)該屬于通過 Padding 避免 bank conflict 的一種思路.
不過筆者并不是很理解這種方法, 因?yàn)閺拇a上來看, Bs[dotIdx * (BN + extraCols) + threadCol * TN + i] 在同一 warp 中仍然存在 threadIdx 差 4 的線程存在 bank conflict 的問題. 而且在從 GMEM 加載到 SMEM 時(shí), 由于 extraCols 的引入使得字節(jié)不對齊而不能向量化加載, 又進(jìn)一步增加了寫入 Bs 的 bank conflict. (Kernel 7 時(shí)提到過, 寫入 Bs 時(shí) threadIdx 差 8 的線程存在 bank conflict, 但 Kernel 7 的向量化訪存避免了這一點(diǎn), 在此處則無法避免了.)
以下分別是 kernel 6 和 kernel 8 使用 nvprof 檢測得到的 bank conflict 次數(shù), 可以看到, 實(shí)際上 SMEM 讀取和寫入的 bank conflict 次數(shù)都有增加.

Invocations                                Event Name         Min         Max         Avg       Total
Device "Tesla V100-SXM2-32GB (0)"
    Kernel: void sgemmVectorize<int=128, int=128, int=8, int=8, int=8>(int, int, int, float, float*, float*, float, float*)
        255                   shared_ld_bank_conflict        8192    33601588     7675626  1957284759
        255                   shared_st_bank_conflict         549     2444397      549564   140139007

Invocations                                Event Name         Min         Max         Avg       Total
Device "Tesla V100-SXM2-32GB (0)"
    Kernel: void sgemmResolveBankExtraCol<int=128, int=128, int=8, int=8, int=8>(int, int, int, float, float*, float*, float, float*)
        255                   shared_ld_bank_conflict       11776    48292001    11033844  2813630341
        255                   shared_st_bank_conflict        2226     8829639     2016935   514318658

Kernel 9: 自動(dòng)調(diào)整 (Autotuning)48

我們總共累積了五個(gè)模板參數(shù):

• BM、BN 和 BK, 它們指定了我們從 GMEM 緩存到 SMEM 的數(shù)據(jù)量.
• TM 和 TN, 它們指定了我們從 SMEM 緩存到寄存器中的數(shù)據(jù)量.

對于 Kernel 6, 這些設(shè)置為 BM＝BN＝128 和 BK＝TM＝TN＝8. 我寫了一個(gè) bash 腳本, 它搜索所有合理的組合, 并對它們的運(yùn)行時(shí)進(jìn)行基準(zhǔn)測試. 這要求我確保:

1. 我知道哪些參數(shù)組合是合理的, 跳過那些不合理的.⁴⁹
2. Kernel 實(shí)現(xiàn)對于保留的大約 400 個(gè)不同的超參數(shù)設(shè)置是正確的.

對代碼的必要修改最終花了相當(dāng)長的時(shí)間來實(shí)現(xiàn).

結(jié)果表明, 根據(jù) GPU 模型的不同, 最佳參數(shù)變化很大.⁵⁰ 在我的 A6000 上, BM=BN=128 BK=16 TM=TN=8 將性能從 19 提升到 20 TFLOP, 提高了 5%. 在 A100 SMX4 40GB 上, 相同的配置達(dá)到了 12 TFLOP, 比自動(dòng)調(diào)整器發(fā)現(xiàn)的最佳設(shè)置(BM=BN=64 BK=16 TM=TN=4)差6%, 后者達(dá)到了 12.6 TFLOPs.⁵¹

我無法解釋為什么這些特定的參數(shù)最終會(huì)產(chǎn)生最佳性能. 自動(dòng)調(diào)整是可行的, 每個(gè)高性能庫都使用它, 但它也讓人感到非常不滿意.⁵²

Kernel 10: warp 分片 (Warptiling)

目前, 我們的循環(huán)結(jié)構(gòu)如下:

現(xiàn)在, 我們將在塊分片和線程分片的循環(huán)之間添加另一個(gè)分片層次: warp 分片. warp 分片最初有點(diǎn)令人困惑, 因?yàn)榕c線程塊和線程不同, warp 不會(huì)顯式顯示在 CUDA 代碼中的任何位置. 它們是在標(biāo)量 CUDA 軟件世界中沒有直接模擬的硬件功能. 我們可以將給定線程的 warpId 計(jì)算為 warpId=threadIdx.x % warpSize, 其中 warpSize 是一個(gè)內(nèi)置變量, 在我使用過的任何 CUDA GPU 上都等于 32.

warp 與性能相關(guān), 因?yàn)?除其他原因外):

• warp 是映射到作為 SM 一部分的 warp 調(diào)度器的調(diào)度單元.⁵³
• 共享內(nèi)存 bank 沖突(我將在未來的文章中介紹這些沖突)只發(fā)生在同一 warp 中的線程之間.
• 最近的 GPU 上有一個(gè)寄存器緩存, 更緊密的線程劃分為我們提供了更多的寄存器緩存位置.

warp 分片是優(yōu)雅的, 因?yàn)槲覀儸F(xiàn)在明確了所有層次的并行性:

• 線程塊分片(blocktiling): 不同的線程塊可以在不同的 SM 上并行執(zhí)行.
• warp 分片(warptiling): 不同的 warp 可以在不同的 warp 調(diào)度器上并行執(zhí)行, 也可以在同一個(gè) warp 調(diào)度器上并發(fā)執(zhí)行.
• 線程分片(threadtiling): (數(shù)量非常有限的)指令可以在相同的 CUDA 內(nèi)核上并行執(zhí)行(=指令級(jí)并行, instruction-level parallelism, 即 ILP).

warp 分片的 CUDA 代碼如下所示:⁵⁴

// dotIdx loops over contents of SMEM
for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
  // populate registers for this thread's part of the warptile
  for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
    for (uint i = 0; i < TM; ++i) {
      regM[wSubRowIdx * TM + i] =
          As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
             threadRowInWarp * TM + i];
    }
  }
  for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
    for (uint i = 0; i < TN; ++i) {
      regN[wSubColIdx * TN + i] =
          Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
             threadColInWarp * TN + i];
    }
  }

  // execute warptile matmul. Later this will map well to
  // warp-wide matrix instructions, executed on tensor cores.
  for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
      // calculate per-thread results with register-cache locality
      for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
          threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
                        (wSubColIdx * TN) + resIdxN] +=
              regM[wSubRowIdx * TM + resIdxM] *
              regN[wSubColIdx * TN + resIdxN];
        }
      }
    }
  }
}

• 譯者注: 根據(jù) Kernel 10 完整代碼可知:
  • WN = 64, WM = 32, 表示矩陣 C warp 分片 M 和 N 維的元素個(gè)數(shù), 即 warp 分片大小;
    warpIdx = threadIdx.x / WARPSIZE,, warpRow = warpIdx / (BN / WN), warpCol = warpIdx % (BN / WN), 分片表示線程所在 warp 的 ID, 以及 warp 在線程塊分片的行列索引;
    WNITER = 2, WMITER = (WM * WN) / (WARPSIZE * TM * TN * WNITER), 分別表示 warp M 和 N 維度需要在 warp 分片內(nèi)迭代的次數(shù).
    WSUBM = WM / WMITER (32/2 = 16), WSUBN = WN / WNITER (64/2 = 32), 分別表示 warp 每次迭代時(shí), M 和 N 維度需要處理的元素?cái)?shù)
    regM[WMITER * TM] = {0.0}, regN[WNITER * TN] = {0.0}. 表示每個(gè)線程緩存的線程分片. 與 Kernel 5 不同的是, 數(shù)組大小分別乘以了 WMITER 和 WNITER, 因?yàn)?warp 分片在 M 和 N 維度上仍需要循環(huán)迭代, 而之前可以理解為 WMITER 和 WNITER 均為 1. 同樣的, 這也致使每個(gè)線程計(jì)算的矩陣 C 的結(jié)果數(shù)多了 WMITER * WNITER 倍, 因此 threadResults[WMITER * TM * WNITER * TN] = {0.0}.
    threadIdxInWarp = threadIdx.x % WARPSIZE, threadColInWarp = threadIdxInWarp % (WSUBN / TN) (i%(32/4) ), threadRowInWarp = threadIdxInWarp / (WSUBN / TN) (i/8), 分別表示 warp 內(nèi)每個(gè)線程的 lane id, 以及在 warp 中的行列索引.
  • 外層 for dotIdx 仍然是沿 A 和 B 線程塊分片的 K 維逐一計(jì)算, 即每次 As 處理一列, Bs 處理一行.
  • for wSubRowIdx 循環(huán)和 for wSubColIdx 循環(huán)分別將 warp 分片包含的元素從共享內(nèi)存中的 As 和 Bs 分片的列和行加載到寄存器中.
    以 As 線程塊分片為例, 每次 warp 迭代每個(gè)線程仍讀取 TM 個(gè)元素到寄存器數(shù)組 regM. 讀取的 As 線程塊分片的索引為 (dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM + threadRowInWarp * TM + i, 要注意, As 分片在加載時(shí)數(shù)據(jù)進(jìn)行了轉(zhuǎn)置, 可以理解為此時(shí)數(shù)據(jù)排布為列主序,
    (dotIdx * BM) 表示 As 分片的第 dotIdx 列;
    warpRow * WM 表示當(dāng)前 warp 需要處理的元素的起始索引(同一 warp 在 M 維度迭代的 WMITER 次讀取的數(shù)據(jù)是連續(xù)排布的, 構(gòu)成一個(gè) warp 分片的數(shù)據(jù), 大小為 WM, 即下圖中第 2 部分到第 3 部分的轉(zhuǎn)換);
    wSubRowIdx * WSUBM 表示 warp 在當(dāng)前 wSubRowIdx 次迭代時(shí)處理元素的起始索引(每次迭代 warp 處理 WSUBM 個(gè)元素);
    threadRowInWarp * TM 表示 warp 中當(dāng)前線程處理的 TM 個(gè)元素的起始索引;’
    i 表示 TM 個(gè)元素中的第 i 個(gè).
    相比于 Kernel 5, 此處 warptiling 帶來的變化主要有兩點(diǎn): 一是以 warp 角度進(jìn)行索引(實(shí)際上代表了將元素映射到唯一的 warp 上); 二是需要 warp 迭代 WMITER 次.
    Bs 線程塊分片與此相同, 在此不再贅述.
  • 接下來的四層嵌套循環(huán)用于執(zhí)行 warp 分片粒度的矩陣乘, 相比 Kernel 5, 多了外面 for wSubRowIdx 和 for wSubColIdx 兩層循環(huán), 分別進(jìn)行 warp 在 M 和 N 維度的迭代.
  • 關(guān)于 warp 分片的理解:
    對于只進(jìn)行線程塊分片和線程分片的 kernel 5, 在處理線程塊分片 As 的 dotIdx 列和 Bs 的 dotIdx 行時(shí), 是將整行和整列映射到所有線程上, 具體來說, 每 16 個(gè)相鄰線程對應(yīng) Bs 的整行元素和 As 的列中的 TM 個(gè)元素, 不同的 16 個(gè)線程對應(yīng) As 列中的不同元素.
    而在 Kernel 10 中, As 的 dotIdx 列的每 BM 個(gè)元素對應(yīng) 1 個(gè) warp, Bs 的 dotIdx 行中的每 BN 個(gè)元素對應(yīng) 1 個(gè) warp. 在這個(gè)大小為 BM*BN 的分片中, warp 的 32 個(gè)線程 1 次迭代處理 WSUBM*WSUBN 大小的小分片, 這個(gè)小分片的處理方式就和 Kernel 5 中處理 BM*BN 大小的線程分塊一樣了, 后續(xù)再進(jìn)行線程分片.
    譯者個(gè)人對進(jìn)行 warp 分片處理的理解是: 在這之前, 每個(gè) warp 都會(huì)讀取 Bs 整行的數(shù)據(jù)和 As 的 2*TM 個(gè)數(shù)據(jù), 不同 warp 間在 Bs 的讀取上重復(fù)性太高, warp 內(nèi) As 上讀取的重復(fù)性太高, 像是兩個(gè)極端; 而進(jìn)行了 warp 分片之后, Bs 只有一部分分給該 warp , As 分給該 warp 的數(shù)據(jù)變多, 緩和了上面的問題, 使得 warp 數(shù)據(jù)的利用更集中, 同時(shí)引入 warp 層次可以更好地解決 bank conflict. (譯者個(gè)人猜測, 不代表正確).

盡管結(jié)構(gòu)變得非常復(fù)雜, 但我還是盡了最大的努力將所有三個(gè)層次的分片都進(jìn)行了可視化.⁵⁵ 每個(gè) warp 將計(jì)算大小為 (WSUBN * WNITER) X (WSUBM * WMITER) 的塊. 每個(gè)線程計(jì)算 WNITER * WMITER 次大小為 TM*TN 的分塊.

• 譯者注:
  • Kernel 10 關(guān)于從 GMEM 加載線程塊分片到 SMEM 也有一些變化:

    for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
        const float4 tmp = reinterpret_cast<const float4 *>(
            &A[(innerRowA + offset) * K + innerColA * 4])[0];
        As[(innerColA * 4 + 0) * BM + innerRowA + offset] = tmp.x;
        As[(innerColA * 4 + 1) * BM + innerRowA + offset] = tmp.y;
        As[(innerColA * 4 + 2) * BM + innerRowA + offset] = tmp.z;
        As[(innerColA * 4 + 3) * BM + innerRowA + offset] = tmp.w;
    }
    
    for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
        reinterpret_cast<float4 *>(
            &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] =
            reinterpret_cast<const float4 *>(
                &B[(innerRowB + offset) * N + innerColB * 4])[0];
    }
    

    相比與 Kernel 6, 增加了一層 offset 的循環(huán). 其中 rowStrideA = (NUM_THREADS * 4) / BK, rowStrideB = NUM_THREADS / (BN / 4), 即增加了行步幅偏移量: 對于線程塊分片 As, 行步幅為整個(gè)線程塊讀取的元素?cái)?shù) NUM_THREADS * 4 除以線程塊分片寬度 BK; 對于線程塊分片 Bs, 行步幅同理(分子分母同乘 4). 這樣做是考慮到線程塊一次讀取的元素?cái)?shù)(NUM_THREADS * 4)小于線程塊分片大小(BM*BK=128*16)的情況, 將線程塊分片的元素按行步幅, 分次加載, 更具有通用性. Kernel 6 及之前是因?yàn)榫€程塊一次讀取的元素與線程塊分片大小是恰好相等的, 所以無需循環(huán).
  • 關(guān)于 bank conflict:
    • 從 GMEM 加載到 SMEM: 這里的代碼雖然添加了步幅, 但與 kernel 6 的代碼基本是一致的. 即:
      對于 As, threadIdx 差 1 的線程 innerColA 差 1, 訪問的元素差 4*BM=4*128=512, 對應(yīng)同一個(gè) bank, 有 bank conflict;
      對于 Bs 由于向量化合并寫入而沒有 bank conflict.
    • 從 SMEM 加載到寄存器:
      對于 As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM + threadRowInWarp * TM + i], warp 中每 4 個(gè)線程對應(yīng)同一個(gè) threadRowInWarp (取值范圍 0 ~7), 因此 i 一定時(shí), threadRowInWarp 相同的 4 個(gè)線程對應(yīng)訪問相同的 4 個(gè)元素, 8 組恰好對應(yīng) 32 個(gè)元素 32 個(gè) bank, 沒有 bank conflict.
      對于 Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN + threadColInWarp * TN + i], warp 中每 4 個(gè)線程對應(yīng) threadColInWarp 的 0~3, 因此, 共訪問 4*TN=4*4=16 個(gè)元素, 對應(yīng) 16 個(gè) bank, 也沒有 bank conflict.

在自動(dòng)調(diào)整參數(shù)后, A100 的性能從 19.7 TFLOP 提高到了 21.7 TFLOP. 下圖展示了在不斷增加的矩陣大小下, 我們的 warp 內(nèi)核與 cuBLAS 的比較:⁵⁶

在 2048 和 4096 維度上, 我們測量的 FLOP 僅比 cuBLAS 慢幾個(gè)百分點(diǎn). 然而, 對于較小的矩陣, 與英偉達(dá)的庫相比, 我們做得很差! 之所以會(huì)出現(xiàn)這種情況, 是因?yàn)?cuBLAS 不只包含了 SGEMM 的一個(gè)實(shí)現(xiàn), 而是包含了數(shù)百個(gè)實(shí)現(xiàn).⁵⁷ 在運(yùn)行時(shí), cuBLAS 將根據(jù)維度選擇運(yùn)行哪個(gè)內(nèi)核.⁵⁸ 我跟蹤了 cuBLAS 調(diào)用, 以下是它在每個(gè)大小下調(diào)用的內(nèi)核:⁵⁹

在 256 維時(shí), 它調(diào)用兩個(gè)內(nèi)核: 一個(gè)矩陣乘內(nèi)核, 后面跟著一個(gè)歸約內(nèi)核.⁶⁰ 因此, 如果我們試圖編寫一個(gè)適用于所有形狀和大小的高性能庫, 我們將針對不同的形狀進(jìn)行特化, 并在運(yùn)行時(shí)選擇最適合的內(nèi)核進(jìn)行調(diào)度.

我還想報(bào)告一個(gè)負(fù)面的結(jié)果: 對于這個(gè)內(nèi)核, 我額外實(shí)現(xiàn)了一個(gè)名為線程交錯(cuò)(thread swizzling)的優(yōu)化. 該技術(shù)假設(shè)線程塊是按照增加 blockIdx 的順序啟動(dòng)的, 并優(yōu)化 blockIdx 到 C 的分塊的映射的方式來增加 L2 局部性.⁶¹ 這篇 Nvidia 的帖子有更多的信息和可視化. 它并沒有提高性能, 大概是因?yàn)?L2 的命中率已經(jīng)相當(dāng)高, 達(dá)到了 80%, 所以我最終刪除了令人眼花繚亂的代碼.⁶²

將 BK 循環(huán)移到外部是有道理的, 因?yàn)樗衔覀兊脑瓌t: “加載一些數(shù)據(jù), 然后在該數(shù)據(jù)上盡可能多地進(jìn)行工作”. 這進(jìn)一步意味著在 BK 循環(huán)內(nèi)部發(fā)生的所有計(jì)算都是獨(dú)立的, 可以并行化處理(例如使用 ILP).

我們現(xiàn)在也可以開始預(yù)取下一次循環(huán)迭代所需的數(shù)據(jù), 這是一種稱為雙緩沖(double buffering)的技術(shù).

正在進(jìn)行的工作: Kernel 11

如果我重新回到這篇文章的工作, 下面是我接下來要看的內(nèi)容:

1. 雙緩沖, 用于更好地交錯(cuò)計(jì)算和內(nèi)存加載. 現(xiàn)在, 請參見 CUTLASS 流水線. 在 CUTLASS 中, 雙重緩沖在兩個(gè)級(jí)別上完成: GMEM ? SMEM 和 SMEM ? 寄存器文件.
   • 在 Hopper 架構(gòu)中, 為 warp 特化引入了新的指令, 例如讓一些 warp 比其他 warp 使用更少的寄存器. 這與特殊指令相結(jié)合, 可以從 GMEM 直接加載到 SMEM 而無需首先通過寄存器, 可用于降低寄存器壓力.
2. 消除 SMEM bank conflict. 這可以通過優(yōu)化 SMEM 中的數(shù)據(jù)布局來實(shí)現(xiàn).
3. 通過查看生成的 PTX, 更好地理解在 Triton 中實(shí)現(xiàn)的 GEMM 內(nèi)核.

Kernel 11: 雙緩沖 doubleBufferIdx 實(shí)現(xiàn) (譯者補(bǔ)充)

原文作者在其 GitHub 中更新了后續(xù)使用雙緩沖從全局內(nèi)存加載數(shù)據(jù)到共享內(nèi)存的 kernel 11 實(shí)現(xiàn).
在 kernel 11 中, 關(guān)鍵是通過變量 bool doubleBufferIdx = threadIdx.x >= (NUM_THREADS / 2); 將線程塊中的線程前后分為兩部分, 每一部分分別加載一個(gè)線程塊分片(包括 As 和 Bs), 然后分別計(jì)算每個(gè)線程塊分片的一部分. 下圖是一個(gè)簡單的可視化:

在圖中, 每個(gè)分塊分別用綠色代表一個(gè) load (從 GMEM 加載線程塊分片到 SMEM) 操作, 用黃色代表一個(gè) process (對線程塊分片進(jìn)行矩陣乘) 操作. 同一個(gè)分片用同一種背景紋路表示. 虛線框表示 for bIdx 循環(huán)(即最外層沿 K 維度遍歷線程塊分片的循環(huán))的一次迭代.

在圖中可以看到, doubleBufferIdx 為 0 和 為 1 的兩部分線程分別加載一個(gè)線程塊分片, 比如 B0 和 B1. 然后兩部分線程分別再處理 B0 和 B1 線程塊分片的矩陣乘計(jì)算. 這里每個(gè)線程塊分片被兩部分線程都計(jì)算的原因是, 在計(jì)算矩陣乘時(shí), 作者的實(shí)現(xiàn)中每個(gè)線程還是和 Kernel 10 一樣按照線程所在的 warp 拿取數(shù)據(jù)并計(jì)算, 換句話說, 線程塊內(nèi)每部分線程完成計(jì)算到每個(gè)線程塊分片的一半計(jì)算; 而不像加載時(shí), 每部分線程是可以將一整個(gè)完整線程塊分片從 GMEM 讀取的.

值得一提的是, 譯者對雙緩沖的實(shí)現(xiàn)存在疑問, 不是很清楚讀取與計(jì)算間的同步關(guān)系是如何保證的(比如, doubleBufferIdx=0 的線程處理 B1 時(shí)如何保證 doubleBufferIdx=1 的線程已經(jīng)加載 B1 完畢). 可能是通過 __syncthreads() 函數(shù)保證的? 但是在 CUDA C++ 編程指南 中提到, “__syncthreads() is allowed in conditional code but only if the conditional evaluates identically across the entire thread block, otherwise the code execution is likely to hang or produce unintended side effects”, 即 __syncthreads() 只有在線程塊的所有線程計(jì)算結(jié)果相同時(shí)才被允許用于條件分支, 而在作者的代碼中顯然不滿足, 因此實(shí)現(xiàn)上這里應(yīng)該是存在一定問題的. 但是代碼確實(shí)可以通過編譯且正常運(yùn)行.
根據(jù) README 來看 Kernel 11 的實(shí)際性能并不很理想, 但在譯者 V100 上進(jìn)行測試, 性能與 kernel 10 是相當(dāng)?shù)? 差距不大.

Kernel 12: 雙緩沖 cuda::memcpy_async 實(shí)現(xiàn) (譯者補(bǔ)充)

原文作者在其 GitHub 中更新了后續(xù)使用雙緩沖從全局內(nèi)存加載數(shù)據(jù)到共享內(nèi)存的 kernel 12 實(shí)現(xiàn).
在 kernel 12 中, 作者使用了 CUDA 異步操作的一些 API, 包括 cuda::barrier, cuda::memcpy_async() 等, 來異步的從 GMEM 加載數(shù)據(jù)到 SMEM, 從而達(dá)到加載到 SMEM 和進(jìn)行矩陣乘計(jì)算的 overlap. 相比于 kernel 11, 譯者認(rèn)為這應(yīng)該是更正確的實(shí)現(xiàn)方式.
Kernel 12 的整體思路與 kernel 11 是一致的, 代替 doubleBufferIdx 的是由 frontBarrierPtr 和 backBarrierPtr 指向的兩個(gè) cuda::barrier. 下圖是一個(gè)簡單的可視化:

在迭代過程中, 一個(gè) barrier 負(fù)責(zé)異步從 GMEM 加載一個(gè)線程塊分片到 SMEM, 另一個(gè) barrier 負(fù)責(zé)計(jì)算上一個(gè)加載的線程塊分片的矩陣乘結(jié)果, 在最后通過交換兩個(gè) barrier 的指針, 達(dá)到雙緩沖的目的.
進(jìn)一步來說, 在迭代線程塊分片前 frontBarrierPtr 加載一個(gè)線程塊分片 B0; 進(jìn)入迭代后, 由 backBarrierPtr 異步加載寫一個(gè)分片; 而 frontBarrierPtr 通過 arrive_and_wait() 函數(shù)等待上一個(gè)分片加載完畢, 然后指向后面的矩陣乘計(jì)算; 計(jì)算結(jié)束后, 交換兩個(gè) barrier 的指針, 這樣下一個(gè)迭代時(shí), 計(jì)算的 barrier 是上一次迭代時(shí)加載數(shù)據(jù)的 barrier, 保證了線程塊分片加載后才被計(jì)算. 在循環(huán)迭代后, 最后一個(gè)加載的線程塊分片單獨(dú)計(jì)算.
與 kernel 11 不同的是, kernel 12 的加載與計(jì)算的 overlap 不需要對線程塊中的線程進(jìn)行劃分, 用 CUDA C++ 編程指南中的話來說, 異步操作就像是另一個(gè)線程異步執(zhí)行的一樣, 因此在這里加載和計(jì)算都是整個(gè)線程塊線程對一個(gè)線程塊分片進(jìn)行的.

額外一提, 異步操作需要 7.0 及其以上算力的 GPU 設(shè)備, 且在 8.0 及其以上的設(shè)備上會(huì)有硬件加速.
譯者在 V100(7.0 算力) 上測試 kernel 12, 發(fā)現(xiàn)性能并不理想, 且又出現(xiàn)了 memory bound 的情況(隨著計(jì)算強(qiáng)度的增加, 在 kernel 5 之后內(nèi)核變?yōu)榱?compute bound), 性能甚至不如 kernel 4, 譯者推測異步操作可能需要 8.0 以上硬件加速的情況下才能取得較好的性能(比如, cuda::memcpy_async() 在 8.0 即以上算力的設(shè)備上可以直接從GMEM 拷貝數(shù)據(jù)到 SMEM 而無需寄存器過渡). 由于譯者沒有 8.0 以上算力的 GPU, 此處還待驗(yàn)證.

總結(jié)

寫這篇文章的經(jīng)歷與我上一篇關(guān)于優(yōu)化 CPU 上的 SGEMM 的文章類似: 迭代優(yōu)化 SGEMM 是深入了解硬件性能特征的最佳方法之一. 對于編寫 CUDA 程序, 我感到驚訝的是, 一旦我對希望的內(nèi)核工作方式進(jìn)行很好地可視化后, 實(shí)現(xiàn)代碼變得十分容易.

此外: 冪律定律無處不在. 我花了兩個(gè)周末的時(shí)間寫了前 6 個(gè)內(nèi)核, 它們達(dá)到了峰值 FLOP 的 80%, 然后又花了 4 個(gè)周末進(jìn)行了自動(dòng)調(diào)整和 warp 分片, 達(dá)到了 94%. 我在寫這段代碼時(shí)學(xué)到的東西也越來越少, 所以我把尋找最后 6% 的東西推遲到未來的某個(gè)時(shí)候.

我所有的代碼都可以在 Github 上找到.

最后, 非常感謝 Godbolt.org (查看 PTX 和 SASS 匯編) 和 Excalidraw (繪制 Kernel) 的創(chuàng)建者! 這兩種工具使用起來都很愉快, 幫助我更快地學(xué)習(xí).

其他資源和參考資料

• 我之所以開始寫這篇文章, 是因?yàn)槲遗既话l(fā)現(xiàn)了 wangzyon 的 Github 倉庫, 我首先嘗試了他的內(nèi)核, 然后從頭開始重寫所有內(nèi)容. 同樣相關(guān)的還有這篇關(guān)于 CUTLASS 庫的英偉達(dá)博客文章.
• 必備參考資料: 官方的 CUDA 工具包編程指南和 CUDA 最佳實(shí)踐指南. 內(nèi)核性能分析指南還包含有關(guān)低級(jí)硬件細(xì)節(jié)(如緩存和流水線)以及可收集的各種指標(biāo)的更多信息.
• Onur Mutlu 是 ETH 的教授, 他將自己的講座上傳到 Youtube 上. 與這篇文章特別相關(guān)的是計(jì)算機(jī)體系結(jié)構(gòu)和異構(gòu)系統(tǒng)的加速.
• 理解 GPU 上的延遲隱藏(Understanding Latency Hiding on GPUs,), 這是一篇深入研究如何設(shè)計(jì)工作負(fù)載以充分利用內(nèi)存帶寬和計(jì)算的博士論文. 它來自 2016 年, 因此只涵蓋較舊的 GPU 架構(gòu). 關(guān)于 warp 同步編程的章節(jié)已經(jīng)過時(shí), 請參閱使用 CUDA warp 級(jí)原語.
• 毛磊(英偉達(dá)的工程師)在他的博客上有很好的 CUDA 內(nèi)容, 包括正確處理 CUDA 錯(cuò)誤.
• 似乎沒有任何好的官方資源來理解 SASS. 這是關(guān)于 CUDA 二進(jìn)制實(shí)用程序的英偉達(dá)文檔. 更有用的可能是查看開源 SASS 匯編程序, 比如 Da Yan 的 turingas.
• 我正在收集可讀但經(jīng)過優(yōu)化的 CUDA 代碼示例, 以供學(xué)習(xí):
  • ONNX Runtime’s CUDA provider, 例如他們對 softmax 的實(shí)現(xiàn).
  • 英偉達(dá)的開源 CUTLASS 庫, 例如他們的 GEMM 實(shí)現(xiàn), 它使用雙緩沖來預(yù)取最內(nèi)層的維度, 這在我的內(nèi)核中仍然缺失.

到了這里，關(guān)于[CUDA 學(xué)習(xí)筆記] 如何優(yōu)化 CUDA 矩陣乘內(nèi)核以獲得類似 cuBLAS 的性能: 工作日志的文章就介紹完了。如果您還想了解更多內(nèi)容，請?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！