cuTile.jl为Julia带来CUDA分块编程

cuTile.jl 为 Julia 带来 NVIDIA CUDA 分块编程

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

现在，通过 cuTile.jl，Julia 也能使用相同的编程模型。本文将探讨 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 实现的性能相当：

一些具有更复杂控制流的内核（如层归一化或 FFT）尚未达到完全性能相当，因为 cuTile.jl 编译器仍在成熟过程中。这些已被追踪为已知问题，并正在积极解决中。

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 Ada、NVIDIA Ampere 或 NVIDIA Blackwell GPU，以及支持 CUDA 13.1 或更高版本的 NVIDIA 驱动程序。该包还需要 Julia 1.11 或更高版本。

启动 Julia，在 REPL 中按 ] 键进入集成包管理器以安装 cuTile.jl：

pkg> add cuTile
pkg> # 如果需要，可以运行测试套件
     test cuTile

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