cuTile.jl 将 NVIDIA CUDA 基于 Tile 的编程引入 Julia

NVIDIA CUDA Tile 是 NVIDIA CUDA 编程中最重大的新增功能之一，它实现了对张量核心及其他专用硬件的自动访问。今年早些时候，NVIDIA 发布了 Python 版 cuTile，为 Python 开发者提供了一种编写高性能 GPU 内核的自然方式。

如今，通过 cuTile.jl，Julia 语言也已支持相同的编程模型。本文将探讨 cuTile.jl 如何简化高性能 CUDA 内核的开发，展示其符合 Julia 语言习惯的语法，并讨论其与现有 Python 版 cuTile 实现的性能持平情况。

什么是基于分块（tile）的 GPU 编程？

使用 CUDA 进行传统的 GPU 编程要求开发者考虑线程、线程束（warps）和内存层次结构。尽管该方法功能强大，但要求程序员能够高效地将算法映射到硬件上。借助 CUDA Tile，开发者只需描述针对数据分块的操作，编译器将自动处理到底层硬件的映射。

以向量加法为例。在传统的 GPU 编程模型中，使用 CUDA.jl 时，程序员必须显式地管理各个线程：

using CUDA

function vadd(a, b, c, n)
    i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    if i <= n
        @inbounds c[i] = a[i] + b[i]
    end
    return
end

threads = 512
blocks = cld(vector_size, threads)
@cuda threads blocks vadd(a, b, c, vector_size)

通过 cuTile.jl 使用 CUDA Tile 后，相同的操作现可在 tile 级别表达，从而隐藏了索引计算或越界检查等细节：

import cuTile as ct

function vadd(a, b, c, tile_size)
    pid = ct.bid(1)
    tile_a = ct.load(a, pid, (tile_size,))
    tile_b = ct.load(b, pid, (tile_size,))
    ct.store(c, pid, tile_a + tile_b)
    return
end

tile_size = 1024
grid = cld(vector_size, tile_size)
ct.launch(vadd, grid, a, b, c, ct.Constant(tile_size))

将其与等效的 Python 代码进行对比：

@ct.kernel
def vadd(a, b, c, tile_size: ct.Constant[int]):
    pid = ct.bid(0)
    tile_a = ct.load(a, index=(pid,), shape=(tile_size,))
    tile_b = ct.load(b, index=(pid,), shape=(tile_size,))
    ct.store(c, index=(pid,), tile=tile_a + tile_b)

tile_size = 1024
grid = ceil(vector_size / tile_size)
ct.launch(stream, grid, vadd, (a, b, c, tile_size))

两者极为相似，且这是有意为之。cuTile.jl 保持了内核的抽象层级与 cuTile Python 编写的版本完全一致，便于代码移植或查阅 cuTile Python 文档进行学习。同时，它在可能的情况下尽可能采用 Julia 惯用法，以使该包对 Julia 程序员更加直观易用，包括基于 1 的索引和用于逐元素操作的广播表达式。

Julia 惯用内核

这一优势在处理超越简单加载与存储的内核时尤为突出。以下是一个行归一化内核——它是层归一化的核心，此处不包含权重和偏置：

function normalize_rows(X, Y, tile_n)
    bid = ct.bid(1)
    tile = ct.load(X, (bid, 1), (1, tile_n))
    mean = sum(tile; dims=2) / size(X, 2)
    centered = tile .- mean
    var = sum(centered .^ 2.0f0; dims=2) / size(X, 2)
    ct.store(Y, (bid, 1), centered ./ sqrt.(var .+ 1f-5))
    return
end

在此示例中，sum、size 和 sqrt 是标准的 Julia 函数，已进行扩展以支持在 tile 上运行。点号（.^、.-、./）是标准的 Julia 广播语法，表示操作将按元素执行。该内核的代码读起来就像常规的 Julia 数组代码。cuTile.jl 内核越贴近普通 Julia 代码，在 CPU 和 GPU 之间共享与复用代码就越容易。

cuTile.jl 的性能

cuTile.jl 与 cuTile Python 一样，均针对相同的 NVIDIA Tile IR 后端，因此两个包生成的 GPU 机器代码类型相同。在 NVIDIA GeForce RTX 5080（计算能力 12.0，NVIDIA Blackwell 架构）上，计算密集型内核的性能可与 Python 实现持平：

由于 cuTile.jl 编译器仍在不断完善中，部分控制流更复杂的内核（如层归一化或 FFT）尚未实现完全的性能对等。这些问题已作为已知问题进行跟踪，并正在积极处理中。

cuTile.jl 的工作原理

cuTile.jl 使用自定义的 Julia 编译器，拦截标准库调用（如 + 、 sum 、 reshape ），并将其路由至 Tile IR 操作。生成的 IR 随后被降级为 Tile IR 字节码，该二进制格式与 cuTile Python 生成的格式相同。此后，NVIDIA 的 tileiras 编译器负责完成向 GPU 机器码的最终编译。

可检查任意内核生成的 Tile IR：

julia> ct.@device_code_tiled ct.launch(vadd, grid, a, b, c, ct.Constant(16))
cuda_tile.module @kernels {
  entry @vadd(%arg0: tile<ptr<f32>>, %arg1: tile<i32>, ...) {
    ...
    return
  }
}

这种透明性对于调试以及理解高级 Julia 代码如何映射为 tile 操作具有重要价值。

cuTile.jl 的当前状态

cuTile.jl 是一个实验性开源软件包，目前正在 JuliaGPU/cuTile.jl 积极开发中。它支持广泛的 tile 操作，包括内存访问、算术运算、规约、扫描、矩阵乘法、形状变换以及原子操作。此外，它还提供了向量加法、矩阵乘法、转置、批量矩阵乘法、层归一化和 FFT 的可运行示例。

尽管如此，这仍是一款处于早期阶段的软件，并且：

• 并非所有 cuTile 功能都已实现。
• 部分 Julia 语言特性（尤其是基于迭代器的 for 循环）在内核中不受支持，或会生成效率较低的代码
• 与 CUDA.jl 的集成仍需改进，以便更好地实现与 SIMT 内核的共存。
• API 可能会在不另行通知的情况下发生变更。

该项目构建于 Julia 现有的 GPU 生态系统之上，并与 CUDA.jl 集成以实现数组管理与内核启动。已使用 CUDA.jl 在 Julia 中编写 GPU 代码的用户会发现，转向基于 tile 的编程十分简便。

快速入门

与 cuTile Python 相同，cuTile.jl 需要 NVIDIA Ada、NVIDIA Ampere 或 NVIDIA Blackwell GPU，以及适用于 CUDA 13.1 或更高版本的 NVIDIA 驱动程序。该程序包还需 Julia 1.11 或更高版本。

启动 Julia，在 REPL 中按下 `]` 键以进入内置包管理器来安装 cuTile.jl：

pkg> add cuTile

pkg> # if you want, run the test suite
     test cuTile

该 GitHub 仓库包含支持操作的完整列表，以及详细说明 cuTile.jl 与 cuTile Python 和标准 Julia 差异的文档。