GPU 上的 Async/Await
Hacker News 摘要VectorWare 团队宣布在 GPU 编程领域取得重大突破，成功在 GPU 上实现了 Rust 语言的 async/await 异步编程模型。这项成果标志着开发者可以使用熟悉的 Rust 抽象编写复杂且高性能的 GPU 应用程序。
GPU 并发编程的演进
传统的 GPU 编程主要关注数据并行。开发者编写单个操作，GPU 将其并行运行在数据的不同部分。随着程序变得复杂，开发者开始使用 Warp Specialization（线程束专业化）来引入复杂的控制流，让 GPU 的不同部分并发执行不同的任务。例如，一个线程束负责从内存加载数据，另一个负责计算。
这种模式虽然提高了硬件利用率，但代价巨大。开发者必须手动管理并发和同步，这极具挑战性且容易出错，类似于在 CPU 上手动处理多线程同步。
现有解决方案及其局限性
目前已有多个项目试图解决这一痛点：
• JAX：将 GPU 程序建模为计算图，通过编译器分析依赖关系并优化执行，提供基于 Python 的高级编程模型。
• Triton：通过独立执行的块来表达计算，使用多级 MLIR 方言管道进行管理和优化。
• CUDA Tile：由 NVIDIA 推出，引入了“平铺”作为一等公民数据单元，使数据依赖关系更加显白。
然而，这些方法要求开发者以特定的新方式组织代码，很难适配所有应用。此外，新范式构成了采用障碍，大多数人仅在机器学习任务中使用这些工具，而非编写整个应用。由于缺乏语言运行时的支持，现有的 CPU 库也无法直接复用。
为什么选择 Rust 的 Future 和 Async/Await
VectorWare 认为 Rust 的异步模型是理想的抽象，原因如下：
• 结构化并发：Future 特性在不绑定特定执行模型的情况下编码了并发结构。它不关心运行在线程、核心还是 GPU 的线程束上。
• 组合性：异步任务像 JAX 的计算图一样可以延迟执行和组合，允许编译器在执行前分析依赖关系。
• 数据安全：Rust 的所有权模型、Pin 以及 Send 和 Sync 等约束，明确了数据在并发单元间共享和转移的规则。
• 自动状态机：Future 会被编译器编译成状态机，而 Warp Specialization 本质上也是手动编写的任务状态机。
在 GPU 上运行异步代码的实现
VectorWare 通过修复多个编译器后端的漏洞，并解决了 NVIDIA ptxas 工具中的问题，成功让 async/await 在 GPU 上运行。
在演示中，简单的异步函数、链式调用、条件分支、多步骤工作流以及第三方库中的组合器（如 futures-util）都能在 GPU 内核中正确执行。代码语法与在 CPU 上编写时完全一致。
GPU 上的执行器
由于 Rust 的 Future 是惰性的，需要执行器（Executor）来驱动。
1. 初始阶段：使用简单的 block_on 执行器证明了可行性。
2. 进阶阶段：引入了 Embassy 执行器。它专为嵌入式系统设计，适用于没有标准库和操作系统的 GPU 环境。
通过 Embassy，开发者可以在 GPU 上构建多个独立的异步任务，这些任务可以无限循环并在共享状态中增加计数。执行器能够交替推进任务进度，实现真正的并发调度。
挑战与不足
尽管取得了成功，但在 GPU 上使用异步模型仍面临挑战：
• 协作式任务：如果某个任务不主动出让执行权，会造成其他任务饥饿。
• 缺乏中断机制：GPU 不提供中断，执行器必须通过轮询来确定任务是否可以继续，这比中断驱动的执行效率低。
• 寄存器压力：维护调度状态会增加寄存器压力，可能降低 GPU 的占用率并影响性能。
• 函数染色问题：依然存在异步函数与同步函数互不调用的固有问题。
未来展望
VectorWare 正在实验专门针对 GPU 硬件特性的原生执行器，利用 CUDA Graphs 或共享内存来优化任务调度和通信。同时，由于他们已经在 GPU 上启用了 Rust 标准库（std），这将为更丰富的运行时和现有异步库的集成打开大门。
虽然该团队目前主要利用 Rust 的强大抽象能力，但他们表示未来的产品将支持多种编程语言和运行时。