本文是帮助开发者学习NVIDIA CUDA Tile编程以构建高性能GPU核函数系列文章的一部分，以矩阵乘法作为核心示例。 在这篇文章中，您将学习到：
- 如何使用NVIDIA cuTile实现高性能矩阵乘法：理解Tile的加载、计算和存储流程。
- 块级并行编程思维：从线程级思维转向块级思维。
- Tile编程的最佳实践：从代码中学习性能优化策略。
在开始之前，请确保您的环境满足以下要求（更多信息请参阅快速入门指南）： 环境要求：
- CUDA 13.1 或更高版本
- GPU架构：某机构 Blackwell（例如，某机构 RTX 50系列）
- Python：3.10 或更高版本
安装cuTile Python包：
pip install cuda-tile
注意：cuTile 是某机构的下一代GPU编程框架。虽然它目前仅支持针对Blackwell架构（计算能力10.x和12.x）的优化，但对更多架构的支持将在后续CUDA Toolkit版本中提供。
什么是矩阵乘法？
矩阵乘法是现代技术计算中的基本运算。它是求解方程组的基础，支撑着图形学、模拟、优化以及大部分机器学习领域，并且非常适合GPU等高性能硬件。
给定输入矩阵 A (MxK) 和 B (KxN)，结果矩阵 C (MxN) 中元素的计算公式如下：
Ci,j = Σk=1K Ai,k Bk,j
从公式可以看出，C中的一个元素是通过取A的一行与B的一列的点积计算得出的。
Tile编程通过将输出矩阵划分为多个Tile来简化实现，同时实现卓越的性能。每个Block负责一个输出Tile的计算，cuTile会自动处理内存访问和线程同步。具体来说：
- 每个Block处理输出矩阵 C 的一个 (tm × tn) Tile。
- 遍历K维度，逐个加载 A 和 B 对应的Tile。
- 使用
ct.mma()执行矩阵乘加运算（自动调用Tensor Core）。 - 最后，将累加结果存回全局内存。
下图展示了计算过程，类似于逐个元素的算法，但此处由Tile取代了单个元素。
图1. 分解为Tile的矩阵乘法 (A * B = C) 示意图
GPU核函数实现
在阐述了核心思想之后，我们来看完整的实现代码。代码分为两部分：在GPU上运行的核函数和在CPU上的启动代码，如下所示。
import cuda.tile as ct
from math import ceil
import torch
# 编译时常量的类型别名
ConstInt = ct.Constant[int]
# 步骤 1: 定义核函数
@ct.kernel
def matmul_kernel(A, B, C, tm: ConstInt, tn: ConstInt, tk: ConstInt):
# 1.1 获取块ID并映射到输出Tile位置
# 在 swizzle_2d 内部，我们访问 ct.bid(0) 并输出 bidx 和 bidy
bidx, bidy = swizzle_2d(M, N, tm, tn, GROUP_SIZE_M)
# 1.2 计算沿K维度的Tile数量
num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk))
# 1.3 初始化累加器
accumulator = ct.full((tm, tn), 0, dtype=ct.float32)
# 1.4 遍历K维度
for k in range(num_tiles_k):
# 从A和B加载Tile
a = ct.load(A, index=(bidx, k), shape=(tm, tk))
b = ct.load(B, index=(k, bidy), shape=(tk, tn))
# 矩阵乘加运算
accumulator = ct.mma(a, b, accumulator)
# 1.5 存储结果
ct.store(C, index=(bidx, bidy), tile=accumulator)
# 步骤 2: 启动核函数
def cutile_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
# 选择Tile尺寸
tm, tn, tk = 128, 256, 64 # 适用于 float16
# 计算网格维度
grid_x = ceil(m / tm)
grid_y = ceil(n / tn)
grid = (grid_x * grid_y, 1, 1)
# 创建输出并启动
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
现在，让我们逐步解析每个关键部分。
1. 定义GPU核函数
在cuTile中，@ct.kernel 装饰器用于将一个标准的Python函数标记为GPU核函数：
@ct.kernel
def matmul_kernel(A, B, C, tm: ConstInt, tn: ConstInt, tk: ConstInt):
# 核函数代码
该装饰器表示：
- 该函数将在GPU上执行。
- 每个Block将独立运行该函数的一个实例。
- 它不能被直接调用，必须使用
ct.launch()启动。
2. 编译时优化：常量类型注解
注意参数 tm、tn 和 tk 使用了特殊的类型注解 ct.Constant[int]：
ConstInt = ct.Constant[int] # 定义类型别名
def matmul_kernel(A, B, C,
tm: ConstInt, # 沿M维度的Tile尺寸
tn: ConstInt, # 沿N维度的Tile尺寸
tk: ConstInt): # 沿K维度的Tile尺寸
这表明它们是编译时常量。cuTile会为不同的Tile尺寸值生成专门的机器代码，使编译器能够：
- 进行循环展开。
- 优化内存访问模式。
- 生成最优的Tensor Core指令。
3. 确定工作范围：块ID映射
每个Block计算输出矩阵的一个特定Tile。通过 swizzle_2d() 函数，我们获取当前正在处理的Block的索引：
def swizzle_2d(M, N, tm, tn, GROUP_SIZE_M):
# 在一维网格中获取当前CUDA块（CTA）的全局ID。
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)
这段代码的功能是确定当前Block应处理输出矩阵的哪个Tile。为了理解这个过程，我们从主机端的网格划分开始。
步骤 1：主机端网格划分 在主机端启动核函数时，计算所需Block的数量：
grid_x = ceil(m / tm) # M维度所需的Block数
grid_y = ceil(n / tn) # N维度所需的Block数
grid_size = grid_x * grid_y # 总Block数
grid = (grid_size, 1, 1) # 定义为一维网格
m和n：输出矩阵 C 的行数和列数。tm：每个Block处理的行方向（M维度）上的输出Tile尺寸。tn：每个Block处理的列方向（N维度）上的输出Tile尺寸。
逻辑上，启动 grid_x * grid_y 个Block，并将它们展平为一维网格：grid = (grid_size, 1, 1)。
步骤 2：在核函数中获取块ID
在核函数内部，每个Block通过 ct.bid(0) 获取其唯一标识符：
bid = ct.bid(0) # 返回值范围：[0, grid_size-1]
ct.bid(0)查询当前Block在x轴维度上的ID。- 参数
0表示第一个维度（x轴），对应于网格定义(grid_size, 1, 1)中的第一个元素。 - 每个Block获得唯一的一维坐标：
bid = 0, 1, 2, …, grid_size-1。
步骤 3：将一维块ID映射到二维Tile坐标
现在的问题是，块ID (bid) 是一维的，但输出矩阵是二维的。我们需要知道这个Block应该处理哪个行和列的Tile。swizzle_2d_from_bid() 函数决定了该Block负责处理哪个行和列的Tile。
bidx, bidy = swizzle_2d_from_bid(M, N, tm, tn, GROUP_SIZE_M, bid)
输出结果：
bidx：当前Block负责的输出Tile的行索引（M维度）。范围：[0, grid_x-1]。bidy：当前Block负责的输出Tile的列索引（N维度）。范围：[0, grid_y-1]。
具体的映射逻辑涉及Swizzle（用于提高内存访问效率），我们将在第4节详细解释。目前，只需理解它将一维Block ID转换为二维Tile坐标即可。
4. 准备累加器：初始化输出Tile
在遍历K维度之前，需要创建一个累加器来存储中间结果：
num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk))
accumulator = ct.full((tm, tn), 0, dtype=ct.float32)
num_tiles_k：计算在K维度上需要处理多少个Tile。accumulator：形状为(tm, tn)的零矩阵，用于累加结果。- 使用
float32可以确保数值精度并避免累加误差。
5. 核心计算循环：遍历K维度
这是矩阵乘法的核心。现在循环遍历K维度上的每个Tile并累加结果：
for k in range(num_tiles_k):
# 加载Tile
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)
# 累加
accumulator = ct.mma(a, b, accumulator)
加载数据：
ct.load(A, index=(bidx, k), shape=(tm, tk))：从矩阵 A 加载一个Tile。index=(bidx, k)：在Tile空间中指定要加载的Tile坐标。shape=(tm, tk)：Tile的尺寸。padding_mode=zero_pad：如果加载数据越界，则用零填充。
矩阵乘加：
ct.mma(a, b, accumulator)：执行a * b，加到accumulator上，并将结果存回accumulator（mma代表矩阵乘加）。- 当
a和b的形状满足Tensor Core要求时，cuTile会自动调用GPU的Tensor Core来加速此操作。
循环结束后，accumulator 存储了输出Tile的完整结果。
6. 写回结果：存储到全局内存
最后，将计算出的结果写回全局内存：
accumulator = ct.astype(accumulator, C.dtype)
ct.store(C, index=(bidx, bidy), tile=accumulator)
- 首先，将
float32类型的累加器转换为输出矩阵的数据类型。 - 使用
ct.store()将Tile写回全局内存中的对应位置。
启动核函数：主机端代码
现在从主机端启动核函数。首先，查看完整的代码。
def cutile_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
# 根据数据类型确定Tile尺寸
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
# 计算网格维度
grid_x = ceil(m / tm)
grid_y = ceil(n / tn)
grid_size = grid_x * grid_y
grid = (grid_size, 1, 1)
# 创建输出张量
C = torch.empty((m, n), device=A.device, dtype=A.dtype)
# 启动核函数
ct.launch(torch.cuda.current_stream(), grid, matmul_kernel,
(A, B, C, tm, tn, tk))
return C
在主机端启动核函数需要三个关键步骤： 步骤 1：计算网格尺寸 根据输入矩阵的维度和Tile尺寸，计算需要多少个Block：
m, k = A.shape # 矩阵A维度：m行，k列
_, n = B.shape # 矩阵B维度：k行，n列
# 计算所需Block数量
grid_x = ceil(m / tm) # M维度需要的Tile数量
grid_y = ceil(n / tn) # N维度需要的Tile数量
grid_size = grid_x * grid_y # 总Block数
grid = (grid_size, 1, 1) # 定义为一维网格
ceil()向上取整以确保覆盖所有元素（即使矩阵维度不能被Tile尺寸整除）。- 将二维Block布局展平为一维网格简化了启动逻辑。
步骤 2：设置Tile尺寸（编译时常量） 根据数据类型选择合适的Tile维度：
if A.dtype.itemsize == 2: # float16/bfloat16 (每个元素2字节)
tm, tn, tk = 128, 256, 64
else: # float32 (每个元素4字节)
tm, tn, tk = 32, 32, 32
这些参数作为编译时常量传递给核函数：
tm：输出Tile的行数（M维度）。tn：输出Tile的列数（N维度）。tk：每次在K维度加载的Tile尺寸。
注意：此处展示的Tile尺寸配置仅为示例。实际上，不同的GPU架构需要不同的参数配置才能达到最佳性能。最优配置取决于M/N/K的尺寸、GPU架构、共享内存大小、寄存器数量、流多处理器（SM）数量等。在开发中，建议使用性能分析工具（如某机构 Nsight Compute）来寻找最佳参数。TileGym 提供了一个自动调优器来自动获取最优参数。
步骤 3：调用 ct.launch() 启动核函数
C = torch.empty((m, n), device=A.device, dtype=A.dtype) # 创建输出张量
ct.launch(
torch.cuda.current_stream(), # CUDA流
grid, # 网格维度: (grid_size, 1, 1)
matmul_kernel, # 核函数
(A, B, C, tm, tn, tk) # 传递给核函数的参数
)
- 流：指定核函数在哪个CUDA流上执行（用于异步执行和多流并发）。
- 网格：定义要启动多少个Block。
- 核函数：要执行的GPU核函数（用
@ct.kernel装饰的函数）。 - 参数元组：传递给核函数的所有参数；
tm、tn和tk将被编译器识别为常量。
性能优化：Swizzle
前文介绍了Swizzle（交织）以提高性能。swizzle_2d_from_bid 的代码如下所示。
def swizzle_2d_from_bid(M, N, tm, tn, GROUP_SIZE_M, bid):
# 在一维网格中获取给定CUDA块的全局ID。
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
在矩阵乘法的例子中，我们使用 GROUP_SIZE_M=8。
Swizzle 如何提升性能？ 它通过分组和交织的方式重新映射块ID到Tile索引，从而更有效地利用缓存。 以下图为例，使用输出矩阵的四个元素（阴影区域）来比较线性访问和Swizzled访问。
图2. 线性行访问与平铺块访问的视觉比较
方法 1：线性行访问
- 计算结果矩阵中的一行数据（例如，四个元素）。
- 需要从左矩阵读取四个块 + 从右矩阵读取所有16个块。
- 总内存访问：20个数据块。
- 右矩阵数据被频繁加载并快速替换，导致缓存命中率低。
方法 2：Swizzling / 平铺块访问
- 将计算重组为2×2的局部块。
- 只需从左矩阵读取八个相关块 + 从右矩阵读取八个相关块。
- 总内存访问：16个数据块（减少了20%）。
- 更好的数据局部性带来了更高的缓存命中率。
性能基准测试
为了验证所实现的矩阵乘法核函数的性能，我们在某机构 GeForce RTX 5080（计算能力12.0）上进行了测试。完整的基准测试代码可以在TileGym仓库中找到。请确保遵循安装说明，然后您可以按照快速入门指南运行此测试和其他测试。
测试配置：
- 数据类型：float16
- 矩阵形状：标准方阵 (N×N)
- 测试尺寸：N = 1024, 2048, 4096, 8192, 16384 (即 2^10 到 2^14)
下图展示了在不同矩阵尺寸下的性能表现。
图3. 在某机构 GeForce RTX 5080 上，cuTile 和 PyTorch 的 TFLOP/s 性能数据随矩阵尺寸变化的对比图
结果表明：
- 在大矩阵规模下，cuTile实现能够充分利用GPU的计算能力。
- 通过合适的Tile尺寸配置和Swizzle优化，cuTile实现的性能达到了业界先进实现（PyTorch调用cuBLAS）的90%以上。
总结
这个经典的矩阵乘法示例展示了使用cuTile实现GPU核函数的完整过程。虽然矩阵乘法很简单，但它包含了Tile编程的核心思想。掌握这些概念将使您能够使用cuTile实现各种高性能GPU核函数。请查看TileGym仓库中的完整矩阵乘法示例以及更多内容，并立即开始编写高性能的Tile代码。FINISHED
