如何在NVIDIA CUDA Tile中編寫高性能矩陣乘法
本博文是系列課程的一部分，旨在幫助開發(fā)者學(xué)習(xí) NVIDIA CUDA Tile 編程，掌握構(gòu)建高性能 GPU 內(nèi)核的方法，并以矩陣乘法作為核心示例。
在本文中，您將學(xué)習(xí)：
如何使用 NVIDIAcuTile實(shí)現(xiàn)高性能矩陣乘法：深入理解平鋪加載、計(jì)算與存儲(chǔ)的執(zhí)行流程。
塊級(jí)并行編程思維的轉(zhuǎn)變：從線程級(jí)思考逐步過渡到以線程塊為核心的編程模式。
平鋪編程的優(yōu)化實(shí)踐：通過實(shí)際代碼掌握性能調(diào)優(yōu)的關(guān)鍵策略。
開始之前，請(qǐng)確認(rèn)您的環(huán)境符合以下要求（更多詳情請(qǐng)參閱快速入門）：
環(huán)境要求：
CUDA 13.1或更高版本
GPU 架構(gòu)：NVIDIA Blackwell（例如，NVIDIA RTX 50 系列）
Python：3.10 及以上版本
安裝 cuTile Python：
| pip install cuda-tile |
注意：cuTile 是 NVIDIA 推出的新一代 GPU 編程框架。盡管目前僅支持針對(duì) Blackwell（計(jì)算能力 10.x 和 12.x）架構(gòu)的優(yōu)化，但即將發(fā)布的 CUDA 工具包版本將擴(kuò)展對(duì)更多架構(gòu)的支持。
什么是矩陣乘法？
矩陣乘法是現(xiàn)代技術(shù)計(jì)算中的一項(xiàng)基本運(yùn)算，它是求解方程組的基礎(chǔ)，支撐著圖形處理、模擬、優(yōu)化以及多數(shù)機(jī)器學(xué)習(xí)任務(wù)，并能高效地映射到 GPU 等高性能硬件上。
給定輸入矩陣 A (M×K) 和 B (K×N)，計(jì)算結(jié)果矩陣 C (M×N) 中各元素的公式如下。
從公式可以看出，矩陣 C 的元素是通過計(jì)算矩陣 A 的行與矩陣 B 的列的點(diǎn)積得到的。
圖塊編程可以通過將輸出矩陣劃分為多個(gè)圖塊，既能簡化實(shí)現(xiàn)，又能實(shí)現(xiàn)優(yōu)異的性能。每個(gè)圖塊負(fù)責(zé)計(jì)算輸出矩陣的一個(gè)子塊，cuTile 會(huì)自動(dòng)處理內(nèi)存訪問和線程同步。具體而言：
每個(gè)塊處理輸出矩陣 C 的 (tm×tn) 圖塊。
沿 K 維度循環(huán)，依次加載矩陣 A 和 B 對(duì)應(yīng)的圖塊。
調(diào)用ct.mma()執(zhí)行矩陣乘積累加運(yùn)算（自動(dòng)啟用 Tensor Core）。
最終，將累積結(jié)果寫回全局內(nèi)存。
圖 1 展示了計(jì)算過程，其方式類似于逐個(gè)元素的算法，但在本例中，圖塊取代了單個(gè)元素。
圖 1。矩陣乘法 (A + B + C) 分解為圖塊的示意圖
GPU 內(nèi)核實(shí)現(xiàn)
在介紹完核心理念之后，我們來看一下完整的實(shí)現(xiàn)代碼。代碼分為兩部分：一部分是在 GPU 上運(yùn)行的內(nèi)核，另一部分是在 CPU 上啟動(dòng)的代碼，如下所示。
|
import cuda.tile as ct from math import ceil import torch # Type alias for compile-time constants ConstInt = ct.Constant[int] # Step 1: Define the kernel @ct.kernel def matmul_kernel(A, B, C, tm: ConstInt, tn: ConstInt, tk: ConstInt): # 1.1 Get block ID and map to output tile position # inside swizzle_2d, we access ct.bid(0) and output bidx and bidy bidx, bidy = swizzle_2d(M, N, tm, tn, GROUP_SIZE_M) # 1.2 Calculate the number of tiles along the K dimension num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) # 1.3 Initialize accumulator accumulator = ct.full((tm, tn), 0, dtype=ct.float32) # 1.4 Loop over K dimension for k in range(num_tiles_k): # Load tiles from A and B a = ct.load(A, index=(bidx, k), shape=(tm, tk)) b = ct.load(B, index=(k, bidy), shape=(tk, tn)) # Matrix multiply-accumulate accumulator = ct.mma(a, b, accumulator) # 1.5 Store result ct.store(C, index=(bidx, bidy), tile=accumulator) # Step 2: Launch the kernel def cutile_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: # Choose tile sizes tm, tn, tk = 128, 256, 64 # for float16 # Calculate grid dimensions grid_x = ceil(m / tm) grid_y = ceil(n / tn) grid = (grid_x * grid_y, 1, 1) # Create output and launch C = torch.empty((m, n), device=A.device, dtype=A.dtype) ct.launch(stream, grid, matmul_kernel, (A, B, C, tm, tn, tk)) return C |
現(xiàn)在，我們來逐步分解每個(gè)關(guān)鍵部分。
1.定義 GPU 內(nèi)核
在 cuTile 中，@ct.kernel裝飾器用于將普通的 Python 函數(shù)標(biāo)記為 GPU 內(nèi)核：
|
@ct.kernel def matmul_kernel(A, B, C, tm: ConstInt, tn: ConstInt, tk: ConstInt): # Kernel code here |
此裝飾器表示：
此函數(shù)將在 GPU 上執(zhí)行。
每個(gè)線程塊將運(yùn)行該函數(shù)的一個(gè)獨(dú)立實(shí)例。
它無法被直接調(diào)用，必須通過ct.launch()來啟動(dòng)。
2. 編譯時(shí)優(yōu)化：常量類型的標(biāo)注
請(qǐng)注意，參數(shù)tm、tn和tk采用特殊類型標(biāo)注ct.Constant[int]：
|
ConstInt = ct.Constant[int] # Define type alias def matmul_kernel(A, B, C, tm: ConstInt, # Tile size along M dimension tn: ConstInt, # Tile size along N dimension tk: ConstInt): # Tile size along K dimension |
這表明它們是編譯時(shí)常量。cuTile 會(huì)針對(duì)不同的圖塊大小值生成專用的機(jī)器代碼，從而使編譯器能夠：
執(zhí)行循環(huán)展開。
優(yōu)化內(nèi)存訪問模式。
生成高效 Tensor Core 指令。
3.確定工作范圍：塊 ID 映射
每個(gè)塊負(fù)責(zé)計(jì)算輸出矩陣的特定圖塊。通過swizzle_2d()函數(shù)，我們獲取當(dāng)前正在處理的塊的索引：
|
def swizzle_2d(M, N, tm, tn, GROUP_SIZE_M): # Get the global IDs of the current CUDA block (CTA) in a 1D grid. bid = ct.bid(0) return swizzle_2d_from_bid(M, N, tm, tn, GROUP_SIZE_M, bid) bidx, bidy = swizzle_2d(M, N, tm, tn, GROUP_SIZE_M) |
此代碼的功能是確定當(dāng)前塊應(yīng)處理的輸出矩陣中的哪個(gè)圖塊。為了理解該過程，我們首先從主機(jī)端的網(wǎng)格劃分開始。
第 1 步：主機(jī)側(cè)網(wǎng)格劃分
在主機(jī)端啟動(dòng)核函數(shù)時(shí)（如第 3 節(jié)所述），計(jì)算所需的任務(wù)塊數(shù)量：
|
grid_x = ceil(m / tm) # Number of Blocks needed for M dimension grid_y = ceil(n / tn) # Number of Blocks needed for N dimension grid_size = grid_x * grid_y # Total Blocks grid = (grid_size, 1, 1) # Defined as a 1D grid |
m和n: 輸出矩陣 C 的行和列。
tm: 輸出圖塊大小行方向 (M 維)由每個(gè)塊處理。
tn: 按列方向( N 維)輸出每個(gè)塊處理的圖塊大小。
從邏輯上講，啟動(dòng)grid_x * grid_y塊并將其展平為一維網(wǎng)格：grid = (grid_size, 1, 1)。
第 2 步：在內(nèi)核中獲取塊 ID
在內(nèi)核內(nèi)部，每個(gè)塊通過ct.bid(0)獲取其唯一的標(biāo)識(shí)符：
| bid = ct.bid(0) # Return value range: [0, grid_size-1] |
ct.bid(0)在 x 軸維度中查詢當(dāng)前塊的 ID。
參數(shù) 0 表示第一個(gè)維度 ( x 軸) ，對(duì)應(yīng)網(wǎng)格定義中的第一個(gè)元素(grid_size, 1, 1).
每個(gè)塊都有一個(gè)唯一的一維坐標(biāo)： bid = 0, 1, 2, …, grid_size-1.
第 3 步：將 1D 塊 ID 映射到 2D 圖塊坐標(biāo)
現(xiàn)在的問題是塊 ID (bid) 為一維，而輸出矩陣是二維。需要明確該塊應(yīng)處理的行和列圖塊。swizzle_2d_from_bid()函數(shù)可用于確定該塊所負(fù)責(zé)的行和列圖塊。
| bidx, bidy = swizzle_2d_from_bid(M, N, tm, tn, GROUP_SIZE_M, bid) |
輸出結(jié)果：
bidx：當(dāng)前塊負(fù)責(zé)的輸出圖塊在 M 維度上的行索引。取值范圍：【0，grid_x -1】。
bidy：當(dāng)前塊負(fù)責(zé)的輸出圖塊在 N 維度上的列索引。取值范圍：【0，grid_y -1】。
特定的映射邏輯涉及 Swizzling（用于提升內(nèi)存訪問效率），我們將在第 4 節(jié)中詳細(xì)解釋這一點(diǎn)。目前，只需理解它將 1D 塊 ID 轉(zhuǎn)換為 2D 圖塊坐標(biāo)即可。
5. 準(zhǔn)備累加器：初始化輸出圖塊
在循環(huán)執(zhí)行 K 維度之前，您需要先創(chuàng)建一個(gè)累加器以存儲(chǔ)中間結(jié)果：
|
num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) accumulator = ct.full((tm, tn), 0, dtype=ct.float32) |
num_tiles_k: 計(jì)算在 K 維度中需要處理的圖塊數(shù)量。
accumulator: 用于累加結(jié)果的形狀 (tm，tn) 零矩陣。
使用 float32 可確保數(shù)值精度并避免累積錯(cuò)誤。
6. 核心計(jì)算循環(huán)：沿 K 維遍歷
這是矩陣乘法的核心。接下來，循環(huán)遍歷 K 維度中的每個(gè)圖塊，并累加結(jié)果：
|
for k in range(num_tiles_k): # Load tiles a = ct.load(A, index=(bidx, k), shape=(tm, tk), padding_mode=zero_pad) b = ct.load(B, index=(k, bidy), shape=(tk, tn), padding_mode=zero_pad) # Accumulate accumulator = ct.mma(a, b, accumulator) |
加載數(shù)據(jù)：
ct.load(A, index=(bidx, k), shape=(tm, tk)): 從矩陣 A 中加載圖塊。
index=(bidx, k): 指定要在圖塊空間中加載的圖塊坐標(biāo)。
shape=(tm, tk): 圖塊的大小。
padding_mode=zero_pad: 如果負(fù)載數(shù)據(jù)超出范圍，則用 0 填充。
矩陣乘積累加：
ct.mma(a, b, accumulator): 乘a*b, 加到accumulator,然后把結(jié)果保存至accumulator(mma表示矩陣乘積累加)
當(dāng)a和b的形狀滿足 Tensor Core 要求時(shí)，cuTile 會(huì)自動(dòng)調(diào)用 GPU 的 Tensor Core 來加速此操作。
循環(huán)結(jié)束后，累加器將保存輸出圖塊的完整結(jié)果。
寫回結(jié)果：存儲(chǔ)到全局內(nèi)存
隨后，將計(jì)算結(jié)果寫回全局內(nèi)存：
|
accumulator = ct.astype(accumulator, C.dtype) ct.store(C, index=(bidx, bidy), tile=accumulator) |
首先，將 float32 累加器轉(zhuǎn)換為輸出矩陣的數(shù)據(jù)類型。
使用ct.store()將圖塊寫回到全局內(nèi)存中的對(duì)應(yīng)位置。
啟動(dòng)核函數(shù)：主機(jī)側(cè)代碼
現(xiàn)在從主機(jī)啟動(dòng)內(nèi)核。首先，查看全部代碼。
|
def cutile_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: # Determine tile sizes based on dtype if A.dtype.itemsize == 2: # float16/bfloat16 tm, tn, tk = 128, 256, 64 else: # float32 tm, tn, tk = 32, 32, 32 m, k = A.shape _, n = B.shape # Calculate grid dimensions grid_x = ceil(m / tm) grid_y = ceil(n / tn) grid_size = grid_x * grid_y grid = (grid_size, 1, 1) # Create output tensor C = torch.empty((m, n), device=A.device, dtype=A.dtype) # Launch kernel ct.launch(torch.cuda.current_stream(), grid, matmul_kernel, (A, B, C, tm, tn, tk)) return C |
在主機(jī)側(cè)啟動(dòng)內(nèi)核需要完成三個(gè)關(guān)鍵步驟：
第 1 步：計(jì)算網(wǎng)格大小
根據(jù)輸入矩陣的維度和圖塊大小，計(jì)算所需塊的數(shù)量：
|
m, k = A.shape # Matrix A dimensions: m rows, k columns _, n = B.shape # Matrix B dimensions: k rows, n columns # Calculate number of Blocks needed grid_x = ceil(m / tm) # How many tiles needed for M dimension grid_y = ceil(n / tn) # How many tiles needed for N dimension grid_size = grid_x * grid_y # Total Blocks grid = (grid_size, 1, 1) # Defined as 1D grid |
ceil()向上取整，確保覆蓋所有元素 (即使矩陣維度無法被圖塊大小整除) 。
將 2D 塊布局扁平化為 1D 網(wǎng)格可簡化啟動(dòng)邏輯。
第 2 步：設(shè)置圖塊大小 (編譯時(shí)常量)
根據(jù)數(shù)據(jù)類型選擇合適的圖塊大小：
|
if A.dtype.itemsize == 2: # float16/bfloat16 (2 bytes per element) tm, tn, tk = 128, 256, 64 else: # float32 (4 bytes per element) tm, tn, tk = 32, 32, 32 |
這些參數(shù)作為編譯期常量傳遞給內(nèi)核：
tm: 輸出圖塊行 ( M 維) 。
tn: 輸出圖塊列 ( N 個(gè)維度) 。
tk: 每次以 K 維加載的圖塊大小。
注意：此處的圖塊大小配置僅為示例。在實(shí)踐中，不同的 GPU 架構(gòu)需要相應(yīng)的參數(shù)配置以達(dá)到理想性能。合適的配置取決于 M/ N/ K 大小、GPU 架構(gòu)、共享內(nèi)存大小、寄存器數(shù)量、SM 數(shù)量等因素。在開發(fā)過程中，建議使用性能分析工具（如 NVIDIA Nsight Compute）確定較優(yōu)參數(shù)。TileGym 提供了一個(gè)自動(dòng)調(diào)整程序，可用于自動(dòng)獲取較優(yōu)參數(shù)。
第 3 步：調(diào)用ct.launch()啟動(dòng)核函數(shù)
|
C = torch.empty((m, n), device=A.device, dtype=A.dtype) # Create output tensor ct.launch( torch.cuda.current_stream(), # CUDA stream grid, # Grid dimensions: (grid_size, 1, 1) matmul_kernel, # Kernel function (A, B, C, tm, tn, tk) # Arguments passed to kernel ) |
Stream：指定核函數(shù)在哪個(gè) CUDA 流上執(zhí)行（用于實(shí)現(xiàn)異步執(zhí)行與多流并發(fā)）。
網(wǎng)格：定義要啟動(dòng)的線程塊數(shù)量。
內(nèi)核函數(shù)：要執(zhí)行的 GPU 內(nèi)核（即使用 ct.kernel 裝飾的函數(shù)）。
參數(shù)元組: 傳遞給內(nèi)核的所有參數(shù)；其中tm、tn和tk將被編譯器識(shí)別為常量。
性能優(yōu)化：Swizzle
為了提升性能，我們引入了早期的 Swizzle。如swizzle_2d_from_bid的代碼所示。
|
def swizzle_2d_from_bid(M, N, tm, tn, GROUP_SIZE_M, bid): # Get the global IDs of a given CUDA block in a 1D grid. num_bid_m = ct.cdiv(M, tm) num_bid_n = ct.cdiv(N, tn) num_bid_in_group = GROUP_SIZE_M * num_bid_n group_id = bid // num_bid_in_group first_bid_m = group_id * GROUP_SIZE_M group_size_m = min(num_bid_m - first_bid_m, GROUP_SIZE_M) bid_m = first_bid_m + (bid % group_size_m) bid_n = (bid % num_bid_in_group) // group_size_m return bid_m, bid_n |
Swizzle 如何提高性能？
它通過分組與交錯(cuò)的方式，將塊 ID 重新映射到平鋪索引，以更高效地利用緩存。
本圖以輸出矩陣的四個(gè)元素（著色區(qū)域）為例，對(duì)比了線性內(nèi)存訪問與 Swizzled 內(nèi)存訪問方式。
圖 2。線性行訪問與分塊平鋪訪問的直觀對(duì)比
方法 1：線性行訪問
計(jì)算結(jié)果矩陣中的一行數(shù)據(jù)（例如四個(gè)元素）時(shí)，
需要讀取左側(cè)矩陣的四個(gè)塊以及右側(cè)矩陣的全部 16 個(gè)塊。
總的內(nèi)存訪問量：20 個(gè)數(shù)據(jù)塊。
由于正確的矩陣數(shù)據(jù)會(huì)被頻繁加載并迅速替換，導(dǎo)致緩存命中率降低。
方法 2：Swizzle/ 平鋪塊訪問
將計(jì)算重新組織為 2 × 2 的本地塊。
僅需讀取左側(cè)矩陣中的 8 個(gè)相關(guān)塊和右側(cè)矩陣中的 8 個(gè)相關(guān)塊。
總顯存訪問量: 16 個(gè)數(shù)據(jù)塊 (減少 20%).
數(shù)據(jù)局部性更優(yōu)，緩存命中率隨之提高。
性能基準(zhǔn)測試
為了驗(yàn)證已實(shí)現(xiàn)的矩陣乘法內(nèi)核的性能，測試在NVIDIA GeForce RTX 5080（計(jì)算能力 12.0）上進(jìn)行。完整的基準(zhǔn)測試代碼可在TileGym資源庫中找到。 請(qǐng)按照安裝說明完成配置后，參照快速入門指南運(yùn)行本測試及其他相關(guān)測試。
測試配置：
數(shù)據(jù)類型: float16
矩陣形狀: 標(biāo)準(zhǔn)方形矩陣（N × N）
測試規(guī)模: N = 1024、2048、4096、8192、16384（即 21? 到 21?）
下圖展示了不同矩陣規(guī)模下的性能表現(xiàn)。
圖 3. NVIDIA GeForce RTX 5080 上 cuTile 與 PyTorch 的 TFLOP/s 性能隨矩陣大小變化的對(duì)比
結(jié)果表明：
在大型矩陣規(guī)模下，cuTile 實(shí)現(xiàn)能夠充分釋放 GPU 的計(jì)算能力。
通過合理的圖塊大小配置與 swizzle 優(yōu)化，cuTile 實(shí)現(xiàn)的性能較業(yè)界先進(jìn)實(shí)現(xiàn)（PyTorch 調(diào)用 cuBLAS）提升90% 以上。
總結(jié)
這個(gè)經(jīng)典的矩陣乘法示例展示了使用 cuTile 實(shí)現(xiàn) GPU 內(nèi)核的完整過程。盡管矩陣乘法較為簡單，但它涵蓋了 Tile 編程的核心理念。掌握這些概念后，您將能夠運(yùn)用 cuTile 實(shí)現(xiàn)多種高性能 GPU 內(nèi)核。請(qǐng)?jiān)赥ileGym 庫中查看完整的矩陣乘法示例及其他相關(guān)內(nèi)容，立即開始編寫高效的圖塊代碼。
關(guān)于作者
Jinman Xie 是 NVIDIA 的深度學(xué)習(xí)性能架構(gòu)師。她畢業(yè)于浙江大學(xué)計(jì)算機(jī)科學(xué)與技術(shù)學(xué)院。Jinman 的主要關(guān)注領(lǐng)域包括深度學(xué)習(xí)模型加速、內(nèi)核優(yōu)化和深度學(xué)習(xí)編譯器技術(shù)。
Qiqi Xiao 畢業(yè)于北京大學(xué)計(jì)算機(jī)科學(xué)專業(yè)，并獲得卡內(nèi)基梅隆大學(xué)碩士學(xué)位。Qiqi 專注于 AI 推理框架、深度學(xué)習(xí)模型優(yōu)化和編譯器技術(shù)。
-
內(nèi)核
+關(guān)注
關(guān)注
4文章
1470瀏覽量
42986 -
NVIDIA
+關(guān)注
關(guān)注
14文章
5662瀏覽量
109942 -
gpu
+關(guān)注
關(guān)注
28文章
5226瀏覽量
135787 -
編程
+關(guān)注
關(guān)注
90文章
3718瀏覽量
97287
原文標(biāo)題：如何在 NVIDIA CUDA Tile 中編寫高性能矩陣乘法
文章出處：【微信號(hào)：NVIDIA-Enterprise，微信公眾號(hào)：NVIDIA英偉達(dá)企業(yè)解決方案】歡迎添加關(guān)注！文章轉(zhuǎn)載請(qǐng)注明出處。
發(fā)布評(píng)論請(qǐng)先 登錄
使用CUDA并行化矩陣乘法加速Blender Python
NVIDIA火熱招聘深度學(xué)習(xí)/高性能計(jì)算解決方案架構(gòu)師
NVIDIA火熱招聘GPU高性能計(jì)算架構(gòu)師
NVIDIA Grid SERIES K2卡兼容CUDA？
探求NVIDIA GPU極限性能的利器
Adreno GPU 矩陣乘法——第1講：OpenCL優(yōu)化
如何使用Warp在Python環(huán)境中編寫CUDA內(nèi)核
NVIDIA cuSPARSELt v0.2.0提高激活函數(shù)
NVIDIA CUDA工具包的概念及主要功能
CUDA矩陣乘法優(yōu)化手段詳解
NVIDIA Hopper GPU上的新cuBLAS12.0功能和矩陣乘法性能
在Python中借助NVIDIA CUDA Tile簡化GPU編程
NVIDIA CUDA Tile的創(chuàng)新之處、工作原理以及使用方法
工商網(wǎng)監(jiān)
評(píng)論