cuTile.jl为Julia语言带来CUDA分块编程的全新高效解决方案

NVIDIA CUDA Tile 是 CUDA 编程领域中一项极具价值的新特性，其核心价值在于帮助开发者轻松挖掘张量核心等专用硬件的运算潜力。今年早些时候，NVIDIA 面向 Python 推出了 cuTile，为 Python 社区提供了一种编写高性能 GPU 内核的“原生”途径。

如今，好消息接踵而至：借助 cuTile.jl，Julia 开发者同样能够享受这一编程模型带来的红利。本文将带你深入了解 cuTile.jl 如何简化高性能 CUDA 内核的开发流程，其语法如何与 Julia 的语言习惯无缝契合，以及在性能表现上与 Python 版本究竟有多接近。

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

在传统的 CUDA 编程中，开发者需要时刻关注线程、线程束以及复杂的内存层次结构。这种方法虽然威力强大，但代价不菲——程序员必须精心将算法高效地“映射”到硬件上。CUDA Tile 的思路则截然不同：它允许开发者直接在数据分块上描述操作，而将底层硬件的映射细节统统交给编译器处理。

以向量加法为例。在传统的 CUDA.jl 模型下，你需要显式管理每个线程，代码如下：

```julia 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=threads blocks=blocks vadd(a, b, c, vector_size) ```

而通过 cuTile.jl 使用 CUDA Tile，同样的操作可以在分块级别表达，隐藏在背后的索引计算、边界检查等细节全部消失：

```julia 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 版本对比：

```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 保持了与 Python 内核完全一致的抽象层级，无论是移植代码还是查阅 Python 文档学习，都毫无障碍。与此同时，它也尽可能地向 Julia 的惯用语法靠拢，比如基于 1 的索引、逐元素操作的广播表达式，这些特性让 Julia 程序员倍感亲切。

符合 Julia 语言习惯的内核

这一优势在稍微复杂的内核中体现得尤为明显。来看一个行归一化内核，这是层归一化（不包括权重和偏置）的核心部分：

```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，Blackwell 架构）上进行的测试表明，计算密集型内核的性能与 Python 版本旗鼓相当：

内核	cuTile.jl	cuTile Python	cuTile.jl 与 cuTile Python 之比
向量加法	838 GB/s	843 GB/s	99%
矩阵转置	797 GB/s	812 GB/s	98%
矩阵乘法	50.9 TFLOPS	50.5 TFLOPS	100%
批量矩阵乘法	43.0 TFLOPS	47.5 TFLOPS	91%

当然，一些控制流更复杂的内核（比如层归一化或 FFT），目前尚未达到完全的性能持平，这主要归因于 cuTile.jl 的编译器仍在持续演进。这些已知问题已经被追踪，开发团队正在积极解决。

cuTile.jl 的工作原理

cuTile.jl 利用一个自定义的 Julia 编译器，它会拦截像 +、sum、reshape 这样的标准库调用，并将它们路由到 Tile IR 操作上去。生成的 IR 接着被降级为 Tile IR 字节码，这个二进制格式与 cuTile Python 生成的内容完全一致。最后，NVIDIA 的 tileiras 编译器负责完成到 GPU 机器码的最终编译。

你还可以检查任何内核生成的 Tile IR，这对于理解高级 Julia 代码是如何映射到分块操作的，非常有价值：

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

这种透明度在调试和深入理解底层细节时，堪称救星。

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 中按下 ] 进入包管理器模式，然后执行：

```julia pkg> add cuTile pkg> # 如果愿意，可以跑一下测试套件 pkg> test cuTile ```

关于支持操作的完整列表，以及 cuTile.jl 与 cuTile Python 和标准 Julia 的详细差异对比，都可以在 GitHub 仓库的文档中找到。