本博文是系列課程的一部分，旨在幫助開發(fā)者學(xué)習(xí) NVIDIA CUDA Tile 編程，掌握構(gòu)建高性能 GPU 內(nèi)核的方法，并以矩陣乘法作為核心示例。

在本文中，您將學(xué)習(xí)：

如何使用 NVIDIAcuTile實現(xiàn)高性能矩陣乘法：深入理解平鋪加載、計算與存儲的執(zhí)行流程。

塊級并行編程思維的轉(zhuǎn)變：從線程級思考逐步過渡到以線程塊為核心的編程模式。

平鋪編程的優(yōu)化實踐：通過實際代碼掌握性能調(diào)優(yōu)的關(guān)鍵策略。

開始之前，請確認您的環(huán)境符合以下要求（更多詳情請參閱快速入門）：

環(huán)境要求：

CUDA 13.1或更高版本

GPU 架構(gòu)：NVIDIA Blackwell（例如，NVIDIA RTX 50 系列）

Python：3.10 及以上版本

安裝 cuTile Python：

注意：cuTile 是 NVIDIA 推出的新一代 GPU 編程框架。盡管目前僅支持針對 Blackwell（計算能力 10.x 和 12.x）架構(gòu)的優(yōu)化，但即將發(fā)布的 CUDA 工具包版本將擴展對更多架構(gòu)的支持。

什么是矩陣乘法？

矩陣乘法是現(xiàn)代技術(shù)計算中的一項基本運算，它是求解方程組的基礎(chǔ)，支撐著圖形處理、模擬、優(yōu)化以及多數(shù)機器學(xué)習(xí)任務(wù)，并能高效地映射到 GPU 等高性能硬件上。

給定輸入矩陣 A (M×K) 和 B (K×N)，計算結(jié)果矩陣 C (M×N) 中各元素的公式如下。

從公式可以看出，矩陣 C 的元素是通過計算矩陣 A 的行與矩陣 B 的列的點積得到的。

圖塊編程可以通過將輸出矩陣劃分為多個圖塊，既能簡化實現(xiàn)，又能實現(xiàn)優(yōu)異的性能。每個圖塊負責(zé)計算輸出矩陣的一個子塊，cuTile 會自動處理內(nèi)存訪問和線程同步。具體而言：

每個塊處理輸出矩陣 C 的 (tm×tn) 圖塊。

沿 K 維度循環(huán)，依次加載矩陣 A 和 B 對應(yīng)的圖塊。

調(diào)用ct.mma()執(zhí)行矩陣乘積累加運算（自動啟用 Tensor Core）。

最終，將累積結(jié)果寫回全局內(nèi)存。

圖 1 展示了計算過程，其方式類似于逐個元素的算法，但在本例中，圖塊取代了單個元素。

圖 1。矩陣乘法 (A + B + C) 分解為圖塊的示意圖

GPU 內(nèi)核實現(xiàn)

在介紹完核心理念之后，我們來看一下完整的實現(xiàn)代碼。代碼分為兩部分：一部分是在 GPU 上運行的內(nèi)核，另一部分是在 CPU 上啟動的代碼，如下所示。

現(xiàn)在，我們來逐步分解每個關(guān)鍵部分。

1.定義 GPU 內(nèi)核

在 cuTile 中，@ct.kernel裝飾器用于將普通的 Python 函數(shù)標記為 GPU 內(nèi)核：

此裝飾器表示：

此函數(shù)將在 GPU 上執(zhí)行。

每個線程塊將運行該函數(shù)的一個獨立實例。

它無法被直接調(diào)用，必須通過ct.launch()來啟動。

2. 編譯時優(yōu)化：常量類型的標注

請注意，參數(shù)tm、tn和tk采用特殊類型標注ct.Constant[int]：

這表明它們是編譯時常量。cuTile 會針對不同的圖塊大小值生成專用的機器代碼，從而使編譯器能夠：

執(zhí)行循環(huán)展開。

優(yōu)化內(nèi)存訪問模式。

生成高效 Tensor Core 指令。

3.確定工作范圍：塊 ID 映射

每個塊負責(zé)計算輸出矩陣的特定圖塊。通過swizzle_2d()函數(shù)，我們獲取當(dāng)前正在處理的塊的索引：

此代碼的功能是確定當(dāng)前塊應(yīng)處理的輸出矩陣中的哪個圖塊。為了理解該過程，我們首先從主機端的網(wǎng)格劃分開始。

第 1 步：主機側(cè)網(wǎng)格劃分

在主機端啟動核函數(shù)時（如第 3 節(jié)所述），計算所需的任務(wù)塊數(shù)量：

m和n: 輸出矩陣 C 的行和列。

tm: 輸出圖塊大小行方向 (M 維)由每個塊處理。

tn: 按列方向( N 維)輸出每個塊處理的圖塊大小。

從邏輯上講，啟動grid_x * grid_y塊并將其展平為一維網(wǎng)格：grid = (grid_size, 1, 1)。

第 2 步：在內(nèi)核中獲取塊 ID

在內(nèi)核內(nèi)部，每個塊通過ct.bid(0)獲取其唯一的標識符：

ct.bid(0)在 x 軸維度中查詢當(dāng)前塊的 ID。

參數(shù) 0 表示第一個維度 ( x 軸) ，對應(yīng)網(wǎng)格定義中的第一個元素(grid_size, 1, 1).

每個塊都有一個唯一的一維坐標： bid = 0, 1, 2, …, grid_size-1.

第 3 步：將 1D 塊 ID 映射到 2D 圖塊坐標

現(xiàn)在的問題是塊 ID (bid) 為一維，而輸出矩陣是二維。需要明確該塊應(yīng)處理的行和列圖塊。swizzle_2d_from_bid()函數(shù)可用于確定該塊所負責(zé)的行和列圖塊。

輸出結(jié)果：

bidx：當(dāng)前塊負責(zé)的輸出圖塊在 M 維度上的行索引。取值范圍：【0，grid_x -1】。

bidy：當(dāng)前塊負責(zé)的輸出圖塊在 N 維度上的列索引。取值范圍：【0，grid_y -1】。

特定的映射邏輯涉及 Swizzling（用于提升內(nèi)存訪問效率），我們將在第 4 節(jié)中詳細解釋這一點。目前，只需理解它將 1D 塊 ID 轉(zhuǎn)換為 2D 圖塊坐標即可。

5. 準備累加器：初始化輸出圖塊

在循環(huán)執(zhí)行 K 維度之前，您需要先創(chuàng)建一個累加器以存儲中間結(jié)果：

num_tiles_k: 計算在 K 維度中需要處理的圖塊數(shù)量。

accumulator: 用于累加結(jié)果的形狀 (tm，tn) 零矩陣。

使用 float32 可確保數(shù)值精度并避免累積錯誤。

6. 核心計算循環(huán)：沿 K 維遍歷

這是矩陣乘法的核心。接下來，循環(huán)遍歷 K 維度中的每個圖塊，并累加結(jié)果：

加載數(shù)據(jù)：

ct.load(A, index=(bidx, k), shape=(tm, tk)): 從矩陣 A 中加載圖塊。

index=(bidx, k): 指定要在圖塊空間中加載的圖塊坐標。

shape=(tm, tk): 圖塊的大小。

padding_mode=zero_pad: 如果負載數(shù)據(jù)超出范圍，則用 0 填充。

矩陣乘積累加：

ct.mma(a, b, accumulator): 乘a*b, 加到accumulator,然后把結(jié)果保存至accumulator(mma表示矩陣乘積累加)

當(dāng)a和b的形狀滿足 Tensor Core 要求時，cuTile 會自動調(diào)用 GPU 的 Tensor Core 來加速此操作。

循環(huán)結(jié)束后，累加器將保存輸出圖塊的完整結(jié)果。

寫回結(jié)果：存儲到全局內(nèi)存

隨后，將計算結(jié)果寫回全局內(nèi)存：

首先，將 float32 累加器轉(zhuǎn)換為輸出矩陣的數(shù)據(jù)類型。

使用ct.store()將圖塊寫回到全局內(nèi)存中的對應(yīng)位置。

啟動核函數(shù)：主機側(cè)代碼

現(xiàn)在從主機啟動內(nèi)核。首先，查看全部代碼。

在主機側(cè)啟動內(nèi)核需要完成三個關(guān)鍵步驟：

第 1 步：計算網(wǎng)格大小

根據(jù)輸入矩陣的維度和圖塊大小，計算所需塊的數(shù)量：

ceil()向上取整，確保覆蓋所有元素 (即使矩陣維度無法被圖塊大小整除) 。

將 2D 塊布局扁平化為 1D 網(wǎng)格可簡化啟動邏輯。

第 2 步：設(shè)置圖塊大小 (編譯時常量)

根據(jù)數(shù)據(jù)類型選擇合適的圖塊大?。?/p>

這些參數(shù)作為編譯期常量傳遞給內(nèi)核：

tm: 輸出圖塊行 ( M 維) 。

tn: 輸出圖塊列 ( N 個維度) 。

tk: 每次以 K 維加載的圖塊大小。

注意：此處的圖塊大小配置僅為示例。在實踐中，不同的 GPU 架構(gòu)需要相應(yīng)的參數(shù)配置以達到理想性能。合適的配置取決于 M/ N/ K 大小、GPU 架構(gòu)、共享內(nèi)存大小、寄存器數(shù)量、SM 數(shù)量等因素。在開發(fā)過程中，建議使用性能分析工具（如 NVIDIA Nsight Compute）確定較優(yōu)參數(shù)。TileGym 提供了一個自動調(diào)整程序，可用于自動獲取較優(yōu)參數(shù)。

第 3 步：調(diào)用ct.launch()啟動核函數(shù)

Stream：指定核函數(shù)在哪個 CUDA 流上執(zhí)行（用于實現(xiàn)異步執(zhí)行與多流并發(fā)）。

網(wǎng)格：定義要啟動的線程塊數(shù)量。

內(nèi)核函數(shù)：要執(zhí)行的 GPU 內(nèi)核（即使用 ct.kernel 裝飾的函數(shù)）。

參數(shù)元組: 傳遞給內(nèi)核的所有參數(shù)；其中tm、tn和tk將被編譯器識別為常量。

性能優(yōu)化：Swizzle

為了提升性能，我們引入了早期的 Swizzle。如swizzle_2d_from_bid的代碼所示。

Swizzle 如何提高性能？

它通過分組與交錯的方式，將塊 ID 重新映射到平鋪索引，以更高效地利用緩存。

本圖以輸出矩陣的四個元素（著色區(qū)域）為例，對比了線性內(nèi)存訪問與 Swizzled 內(nèi)存訪問方式。

圖 2。線性行訪問與分塊平鋪訪問的直觀對比

方法 1：線性行訪問

計算結(jié)果矩陣中的一行數(shù)據(jù)（例如四個元素）時，

需要讀取左側(cè)矩陣的四個塊以及右側(cè)矩陣的全部 16 個塊。

總的內(nèi)存訪問量：20 個數(shù)據(jù)塊。

由于正確的矩陣數(shù)據(jù)會被頻繁加載并迅速替換，導(dǎo)致緩存命中率降低。

方法 2：Swizzle/ 平鋪塊訪問

將計算重新組織為 2 × 2 的本地塊。

僅需讀取左側(cè)矩陣中的 8 個相關(guān)塊和右側(cè)矩陣中的 8 個相關(guān)塊。

總顯存訪問量: 16 個數(shù)據(jù)塊 (減少 20%).

數(shù)據(jù)局部性更優(yōu)，緩存命中率隨之提高。

性能基準測試

為了驗證已實現(xiàn)的矩陣乘法內(nèi)核的性能，測試在NVIDIA GeForce RTX 5080（計算能力 12.0）上進行。完整的基準測試代碼可在TileGym資源庫中找到。 請按照安裝說明完成配置后，參照快速入門指南運行本測試及其他相關(guān)測試。

測試配置：

數(shù)據(jù)類型: float16

矩陣形狀: 標準方形矩陣（N × N）

測試規(guī)模: N = 1024、2048、4096、8192、16384（即 21? 到 21?）

下圖展示了不同矩陣規(guī)模下的性能表現(xiàn)。

圖 3. NVIDIA GeForce RTX 5080 上 cuTile 與 PyTorch 的 TFLOP/s 性能隨矩陣大小變化的對比

結(jié)果表明：

在大型矩陣規(guī)模下，cuTile 實現(xiàn)能夠充分釋放 GPU 的計算能力。

通過合理的圖塊大小配置與 swizzle 優(yōu)化，cuTile 實現(xiàn)的性能較業(yè)界先進實現(xiàn)（PyTorch 調(diào)用 cuBLAS）提升90% 以上。

總結(jié)

這個經(jīng)典的矩陣乘法示例展示了使用 cuTile 實現(xiàn) GPU 內(nèi)核的完整過程。盡管矩陣乘法較為簡單，但它涵蓋了 Tile 編程的核心理念。掌握這些概念后，您將能夠運用 cuTile 實現(xiàn)多種高性能 GPU 內(nèi)核。請在TileGym 庫中查看完整的矩陣乘法示例及其他相關(guān)內(nèi)容，立即開始編寫高效的圖塊代碼。

關(guān)于作者

Jinman Xie 是 NVIDIA 的深度學(xué)習(xí)性能架構(gòu)師。她畢業(yè)于浙江大學(xué)計算機科學(xué)與技術(shù)學(xué)院。Jinman 的主要關(guān)注領(lǐng)域包括深度學(xué)習(xí)模型加速、內(nèi)核優(yōu)化和深度學(xué)習(xí)編譯器技術(shù)。

Qiqi Xiao 畢業(yè)于北京大學(xué)計算機科學(xué)專業(yè)，并獲得卡內(nèi)基梅隆大學(xué)碩士學(xué)位。Qiqi 專注于 AI 推理框架、深度學(xué)習(xí)模型優(yōu)化和編譯器技術(shù)。