cuTile.jl 为 Julia 带来基于 NVIDIA CUDA Tile 的编程

NVIDIA CUDA Tile 是 NVIDIA CUDA 编程的一项重要新增功能，可自动访问 Tensor Core 和其他专用硬件。今年早些时候，NVIDIA 推出了适用于 Python 的 cuTile，为 Python 开发者提供了一种编写高性能 GPU 内核的自然方式。

现在，Julia 中通过 cuTile.jl 提供了相同的编程模型。在本文中，我们将探讨 cuTile.jl 如何简化高性能 CUDA 核函数的开发，展示其符合 Julia 惯用语法的设计，并讨论其与现有的 cuTile Python 实现在性能上的相当表现。

什么是基于图块的 GPU 编程？

使用 CUDA 进行传统 GPU 编程时，开发者需要考虑线程、线程束以及内存层次结构。这种方法虽然功能强大，但要求程序员将算法高效地映射到硬件上。借助 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 后，相同的运算将在图块级别上进行表达，从而隐藏索引计算、越界检查等细节：

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 函数，经过增强后可用于处理图块。这些点（.^、.-、./）采用标准的 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 代码如何映射到平铺操作至关重要。

cuTile.jl 的当前状态

cuTile.jl 是一个实验性开源软件包，由 JuliaGPU/cuTile.jl 积极开发。它支持多种平铺运算，包括内存访问、算术运算、归约、扫描、矩阵乘法、形状操作以及原子操作。此外，还提供了向量加法、矩阵乘法、转置、批量矩阵乘法、层归一化和 FFT 的示例实现。

也就是说，这是处于早期阶段的软件，并且：

• 并非所有 cuTile 功能都已实现。
• 内核不支持部分 Julia 语言功能（尤其是基于迭代器的“for”循环），且可能生成低效代码。
• 与 CUDA.jl 的集成仍需改进，以更好地支持与 SIMT 内核的共存。
• API 可能随时变更，恕不另行通知。

该项目基于 Julia 现有的 GPU 生态系统，通过集成 CUDA.jl 实现数组管理和核函数启动。对于已在 Julia 中使用 CUDA.jl 编写 GPU 代码的用户而言，过渡到基于图块的编程将十分简便。

开始使用

与 cuTile Python 相同，cuTile.jl 需要配备 NVIDIA Blackwell GPU，并安装支持 CUDA 13 或更高版本的 NVIDIA 驱动程序。该软件包还要求使用 Julia 1.11 或更高版本。

启动 Julia，然后在 REPL 中按下 `】` 进入集成包管理器，安装 cuTile.jl：

pkg> add cuTile

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

The GitHub 提供了受支持操作的完整列表，以及关于 cuTile.jl 与 cuTile Python 和标准 Julia 之间差异的详细文档。