本博文是系列課程的一部分，旨在幫助開發者學習 NVIDIA CUDA Tile 編程，掌握構建高性能 GPU 內核的方法，并以矩陣乘法作為核心示例。

在本文中，您將學習：

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

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

平鋪編程的優化實踐：通過實際代碼掌握性能調優的關鍵策略。

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

環境要求：

CUDA 13.1或更高版本

GPU 架構：NVIDIA Blackwell（例如，NVIDIA RTX 50 系列）

Python：3.10 及以上版本

安裝 cuTile Python：

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

什么是矩陣乘法？

矩陣乘法是現代技術計算中的一項基本運算，它是求解方程組的基礎，支撐著圖形處理、模擬、優化以及多數機器學習任務，并能高效地映射到 GPU 等高性能硬件上。

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

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

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

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

沿 K 維度循環，依次加載矩陣 A 和 B 對應的圖塊。

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

最終，將累積結果寫回全局內存。

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

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

GPU 內核實現

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

現在，我們來逐步分解每個關鍵部分。

1.定義 GPU 內核

在 cuTile 中，@ct.kernel裝飾器用于將普通的 Python 函數標記為 GPU 內核：

此裝飾器表示：

此函數將在 GPU 上執行。

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

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

2. 編譯時優化：常量類型的標注

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

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

執行循環展開。

優化內存訪問模式。

生成高效 Tensor Core 指令。

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

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

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

第 1 步：主機側網格劃分

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

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

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

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

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

第 2 步：在內核中獲取塊 ID

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

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

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

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

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

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

輸出結果：

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

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

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

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

在循環執行 K 維度之前，您需要先創建一個累加器以存儲中間結果：

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

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

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

6. 核心計算循環：沿 K 維遍歷

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

加載數據：

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

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

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

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

矩陣乘積累加：

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

當a和b的形狀滿足 Tensor Core 要求時，cuTile 會自動調用 GPU 的 Tensor Core 來加速此操作。

循環結束后，累加器將保存輸出圖塊的完整結果。

寫回結果：存儲到全局內存

隨后，將計算結果寫回全局內存：

首先，將 float32 累加器轉換為輸出矩陣的數據類型。

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

啟動核函數：主機側代碼

現在從主機啟動內核。首先，查看全部代碼。

在主機側啟動內核需要完成三個關鍵步驟：

第 1 步：計算網格大小

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

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

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

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

根據數據類型選擇合適的圖塊大小：

這些參數作為編譯期常量傳遞給內核：

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

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

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

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

第 3 步：調用ct.launch()啟動核函數

Stream：指定核函數在哪個 CUDA 流上執行（用于實現異步執行與多流并發）。

網格：定義要啟動的線程塊數量。

內核函數：要執行的 GPU 內核（即使用 ct.kernel 裝飾的函數）。

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

性能優化：Swizzle

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

Swizzle 如何提高性能？

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

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

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

方法 1：線性行訪問

計算結果矩陣中的一行數據（例如四個元素）時，

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

總的內存訪問量：20 個數據塊。

由于正確的矩陣數據會被頻繁加載并迅速替換，導致緩存命中率降低。

方法 2：Swizzle/ 平鋪塊訪問

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

僅需讀取左側矩陣中的 8 個相關塊和右側矩陣中的 8 個相關塊。

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

數據局部性更優，緩存命中率隨之提高。

性能基準測試

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

測試配置：

數據類型: float16

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

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

下圖展示了不同矩陣規模下的性能表現。

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

結果表明：

在大型矩陣規模下，cuTile 實現能夠充分釋放 GPU 的計算能力。

通過合理的圖塊大小配置與 swizzle 優化，cuTile 實現的性能較業界先進實現（PyTorch 調用 cuBLAS）提升90% 以上。

總結

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

關于作者

Jinman Xie 是 NVIDIA 的深度學習性能架構師。她畢業于浙江大學計算機科學與技術學院。Jinman 的主要關注領域包括深度學習模型加速、內核優化和深度學習編譯器技術。

Qiqi Xiao 畢業于北京大學計算機科學專業，并獲得卡內基梅隆大學碩士學位。Qiqi 專注于 AI 推理框架、深度學習模型優化和編譯器技術。