Hotdry.
归档
systems-engineering

用 cuTile Python 写 GPU Tile 并行内核：共享内存排布与可扩展参数实战

围绕 cuTile 的 tile-based 并行范式，给出共享内存自动排布机制、可扩展 grid/tile 参数表与矩阵乘实战要点。

cuTile Python 把 “写内核” 这件事压缩成了三步：加载 tile、算 tile、写回 tile。它把线程块、共享内存排布、Tensor Core 调度全部藏进编译器，只留下一套 Python 级 DSL。本文用 16×16 矩阵乘做主线，把 “自动共享内存布局” 与 “可扩展参数” 这两个最容易踩坑的点拆给你看。

从 SIMT 到 tile-based：cuTile 的设计动机

传统 CUDA C 的 SIMT 模型要求开发者手动划分线程索引、计算偏移、申请 __shared__ 并处理 bank conflict。cuTile 把 “线程” 这个概念弱化，改用 “tile”—— 一个形状在编译期就确定的子矩阵 —— 作为最小并行单元。运行时，你只需要告诉 cuTile：

  • 全局数组多大
  • 每个 tile 多大
  • 要做什么运算

编译器会自动把 tile 映射到线程块、把共享内存排成 swizzle 模式、把矩阵乘拆到 Tensor Core。宿主代码通过 ct.launch() 把 kernel 扔进 CUDA 流即可。

tile 与 array 的两级内存抽象

cuTile 的数据层只有两级：

  1. array：住在 global memory，可变、有 stride，宿主端可以是 CuPy、PyTorch 甚至 NumPy（通过 DLPack）。
  2. tile：住在寄存器或共享内存，不可变值对象，维度必须是 2 的幂且 ≤512。

关键函数只有四个：

a_tile = ct.load(array, index=(block_id,), shape=(TILE_SIZE,))
...  # 运算
ct.store(array, index=(block_id,), tile=a_tile)

共享内存排布完全由编译器决定。以 16×16 的 float16 tile 为例，后端会选用 32-byte swizzle，保证相邻线程访问同一行时落在不同 bank，天然免疫冲突。如果你想手动 double-buffer，只需再申请一个同名 tile，编译器会自动做阶段划分与预取。

实战：16×16 矩阵乘 kernel 逐行拆解

下面给出 “C = A × B” 最精简写法，tile 大小 16×16，grid 按输出矩阵的 tile 数计算：

import cuda.tile as ct
import cupy as cp

@ct.kernel
def gemm_16x16(A, B, C, M, N, K):
    # 当前 tile 在全局输出矩阵中的坐标
    tile_x = ct.bid(0)  # 列方向
    tile_y = ct.bid(1)  # 行方向

    # 申请三个 tile
    a_frag = ct.tile((16, 16), dtype=ct.float16)
    b_frag = ct.tile((16, 16), dtype=ct.float16)
    c_frag = ct.tile((16, 16), dtype=ct.float32)

    # K 维按 16 步长累加
    for k_tile in range(0, K, 16):
        ct.load(a_frag, A, index=(tile_y, k_tile // 16))
        ct.load(b_frag, B, index=(k_tile // 16, tile_x))
        c_frag += ct.matmul(a_frag, b_frag)   # 自动启用 Tensor Core

    ct.store(C, index=(tile_y, tile_x), tile=c_frag)

# 宿主端调用
M, N, K = 4096, 4096, 4096
grid = (M // 16, N // 16, 1)
ct.launch(cp.cuda.get_current_stream(), grid, gemm_16x16,
          (A, B, C, M, N, K))

注意三点：

  1. ct.matmul 会在编译期识别到 float16×float16→float32 的 Tensor Core 路径，无需 wmma 或 mma 指令。
  2. 共享内存大小由编译器根据 tile 形状与双缓冲策略自动计算，默认每线程块 48 KiB 以内都能放下。
  3. 如果 M/N/K 不是 16 的整数倍，只需在 kernel 里加边界判断，或在外层用 ct.cdiv 向上取整。

可扩展参数清单

参数 推荐值 上限 备注
tile 行 / 列 16, 32, 64, 128 512 必须是 2 的幂；越大共享内存压力越高
每 SM 同时驻留块 — 32 由寄存器 + 共享内存同时决定，Nsight Compute 看 Achieved Active Warps Per SM
grid 维度 任意 2^31-1 按输出 tile 数计算即可，cuTile 自动映射到 hardware blockIdx
占用率经验公式 min(1, 32 / (reg_tiles + 2)) — reg_tiles 为 kernel 里同时 live 的 tile 数，>0.5 即算良好

Nsight Compute 必看指标：

  • sm__sass_average_data_ports_busy：接近 100% 说明 Tensor Core 吃满。
  • shared_load_conflict：cuTile 编译后应为 0，若不为 0 说明 tile 形状或 dtype 触发 fallback。
  • launch__occupancy：与上述经验公式互验，低于 0.35 考虑减 tile 大小或减少 live tile。

局限与回退策略

  1. 动态形状：tile 维度必须编译期已知。若你的业务在运行时才知道矩阵宽高，需要在外层做 padding 到 2 的幂，或退回普通 CUDA kernel。
  2. 非 2 的幂维度：cuTile 会直接报错。可用 ct.cdiv 向上对齐，再在 kernel 里用掩码忽略多余元素。
  3. 平台支持：目前仅 Linux x86_64/aarch64 与 Windows x86_64，macOS 无计划。CI 环境可用官方 Docker：nvidia/cuda:13.1-devel-ubuntu22.04。

小结

cuTile Python 把 “共享内存排布” 与 “Tensor Core 调度” 这两件最费手的活收进编译器，留给开发者的只剩三件事：选 tile 大小、写运算、调 grid。只要记住 “tile 必须是 2 的幂” 与 “grid 按输出 tile 数算” 这两条铁律，你就能在 30 行代码内拿到一块可扩展、可剖面、跨 GPU 架构的 GEMM 内核。

资料来源
[1] NVIDIA/cutile-python GitHub 仓库
[2] cuTile Python 官方文档

查看归档