Rust 与 GPU 编程的现状与前景探究

话说，程序员三大浪漫，操作系统、编译器和图形处理。Rust 语言已经攻陷了其中两大浪漫，操作系统和编译器，那么图形处理呢？Rust 语言还能“浪”起来吗?

这周在 Phoronix[1] 网站上得知用 Rust 重写的 NAK 编译器已经被 Merge 到了 Mesa 24.0 版本中，用于 Nouveau Gallium3D 驱动程序和 NVK Vulkan 驱动程序。这激起了我的好奇心。因为 GPU 编程是 Rust 语言进入图形处理的关键，所以我想彻底了解一下 Rust 目前在 GPU 编程生态方面的现状和前景。

这就是本文的出发点。

Mesa 相关背景介绍

Mesa （或称 Mesa3D） 是一个 OpenGL/Vulkan 的实现，以及为所有开源图形驱动提供各种图形库（GL）的入口点。

Mesa有两个作用：

1. 对接各种 GPU 硬件，将应用层对 GL API 的调用转换到对硬件 GPU 的调用上；
2. 各种 GL API 的纯软实现，当没有可用的硬件时，它可以提供纯软件的 GL API 的实现；

它可以用于 Linux/Windows/Mac 等系统平台。在 Windows 上运行时它提供 OpenGL API over DirectX 的转换。AMD 和 Intel 都提供了对 Mesa 支持的驱动程序。

NVIDIA 有两种类型的显卡驱动：一个是开源的 Nouveau ，另一个是闭源的 NVIDIA 官方驱动。Mesa 使用了开源的 Nouveau 驱动。而 Gallium3D 是 Mesa 提出的用于简化 GPU 驱动开发的框架。

NVK，是由 Collabora 推出的一个新的 Mesa 开源驱动程序， 为 NVIDIA 显卡实现 Vulkan 图形 API。该驱动程序是使用 NVIDIA 发布的官方头文件，以及开放的数据中心 GPU 和消费级 GPU（GTX/RTX）的 GPU 内核模块，从头开始编写的。它的目标是成为新的主流显卡驱动。

NVK 与其他的 Nouveau 驱动非常不同，因为它是从头开始编写的。

nouveau 是一个主要的 NVIDIA 显卡的开源驱动程序，已经年久失修了，试图在它的基础上构建是一个很多人都无法承担的任务。当然，它是由有很多才华的工程师开发的，但是缺乏公司的支持和贡献者的影响了它的发展。NVK 旨在克服这些问题，同时专注于对 Turing 系列及更高版本 GPU 的支持。

由于内核的开发方式，对于 Kepler、Maxwell 和 Pascal 等较旧的 GPU 的支持可能不会很容易地加入 NVK。它也许极大地依赖于新内核，从而只支持较新的 GPU。同时，nouveau 内核接口与 Vulkan 不兼容，阻碍了对较旧 GPU 的支持。

而用 Rust 实现的新的 NAK （Nvidia awesome kompile）编译器，是一个为 NVIDIA GPU 设计的后端编译器，专门为处理 GPU 任务优化。NAK 使用静态单赋值形式（SSA）和中间表示（NIR），这有助于优化代码并提高 GPU 的执行效率。

静态单赋值形式（SSA）是一种编译器中间表示，使每个变量只被赋值一次。这简化了许多编译器优化，因为变量的值在它们的生命周期内保持不变。在转换为 SSA 形式时，编译器会重写代码，使得每个变量的每个赋值操作都有一个唯一的变量名。使得数据流分析更加直接和高效，因为每个变量的定义和使用都是显而易见的。

“顺便说一下，Rust 编译器（rustc）在编译过程的 MIR 和 LLVM IR 这两个阶段也使用了静态单赋值形式（SSA）。

NIR 是 Mesa 中提供的一种更加具体的中间语言，它是为了优化和简化驱动编译器的工作流程而设计的。NIR 设计上更接近硬件，旨在作为多个不同前端（如 GLSL、SPIR-V）和多个不同后端（如不同的 GPU 驱动）之间的桥梁。NIR 也支持 SSA 等各种优化技术。NIR 使 Mesa 能够更有效地处理来自不同源的图形和计算着色器代码，为最终在 GPU 上执行的代码生成和优化奠定基础。

借用网上的一张旧图来说明 Mesa 中的 IR 架构：

此次 Rust 实现的 NAK 编译后端大约研发了半年，长期计划是成为 NVK 的编译器。这并不意味着可以直接使用 Rust 来编写着色器程序，因为它只是一个编译后端。然而，这也算是朝 Rust 直接进行 GPU 编程更进了一步。

Rust 作为 GPU 着色器语言的前景

图形渲染机制简单来说是这样的：

图形库（比如 OpenGL）将渲染计算任务实时派发给 GPU，具体由一种用图形库提供的着色语言（GLSL或 WLSL等）或 SPIR-V（着色语言中间语言标准）编写的称为着色器（sharder）的小程序，在 GPU 上编译运行。从基本意义上来说，着色器只是一种把输入转化为输出的程序。着色器也是一种非常独立的程序，因为它们之间不能相互通信。着色语言一般包含一些针对向量和矩阵操作的有用特性。

常见的着色语言有 :

• DirectX 使用 HLSL（High Level Shading Language）
• Metal 使用 MSL（Metal Shading Language）
• OpenGL 使用 GLSL（OpenGL Shading Language）
• Vulkan 使用的着色器必须以 SPIR-V 这种二进制字节码的格式提供。

所以一般来说，在 Rust 生态中，需要使用 GPU 进行图形渲染则多半是需要直接使用 GLSL 这类着色语言。这就带来了极大的不便。幸好，Rust 生态有一些开源项目，正在致力于改变这一状况。

rust-gpu 项目

国外知名游戏工作室 EmbarkStudios 一直在维护一个 rust-gpu[2] 项目。它旨在为 Rust 编译器打造一个 spir-v 的编译后端。

SPIR-V 是一个为 Vulkan 和 OpenCL 设计的中间语言（IL）标准。它是一个低层次、与硬件无关的 IR，用于表达着色器和计算核心。SPIR-V 设计上更靠近硬件执行层次，它直接被 GPU 驱动所接受，并转换为特定硬件的机器代码。

正如前面 Mesa 的 IR 图所示，SPIR-V 通常作为着色器语言（如 GLSL）的编译输出，然后可以被转换为 NIR 进行进一步的优化和处理。这就意味着，如果 rust-gpu 成熟了，可以直接用 Rust 语言作为着色语言来编写着色器程序，这样就可以通过 SPIR-V 转换为 NIR ，进一步让 NAK 编译后端来处理了。

从这个角度来看，rust-gpu 这个项目对于 Rust GPU 图形编程渲染生态还是非常重要的。

wgpu

另外一个项目是 wgpu[3] ，它提供了一个安全、跨平台的图形和计算 API，基于 WebGPU 规范。WebGPU 是一种新的图形标准，旨在为现代图形硬件提供统一的低层次访问。它被设计为更安全、更高效，特别是在 Web 应用程序中。

wgpu 是基于 Rust 实现的，所以它利用 Rust 的安全特性来帮助避免常见的内存错误和并发问题，这在处理复杂的图形任务时尤其重要。wgpu 也充分利用了现代 GPU 的能力，提供高效的图形和计算性能。它支持最新的图形技术，如计算着色器和高效的资源管理。wgpu 提供了 Rust 风格的 API，相比于直接使用 Vulkan 或 Direct3D，它提供了更高级别的抽象，简化了图形编程的复杂性。

wgpu 不仅可以在 Web 环境运行，还可以在 macOS / iOS、Android、Window 和 Linux 等系统上原生运行。随着 WebGPU 标准的发展和成熟，wgpu 可能会成为 Web 和非 Web 应用程序中 GPU 编程的首选解决方案。wgpu 对于游戏开发、图形设计和可视化、科学计算和机器学习等领域非常适合。目前被用于 Firefox、Servo 和 Deno 中 WebGPU 整合的核心。

“wgpu 实际上也提供了 C 语言绑定 (wgpu-native[4])，你可以写 C/C++ 或其他能与 C 互通的语言来使用它。

wgpu 还有另一个重要的优势，那就是可以利用各种强大的桌面端 GPU 调试工具。在开发大型 2D/3D 应用时，通过使用命令记录/回放、帧捕捉、Buffer 视图等功能，可以快速定位 GPU 层代码/数据的性能瓶颈和程序缺陷。相较于仅依靠浏览器提供的有限调试能力，这些工具能够事半功倍，帮助开发者更快地解决问题。

WebGPU 使用的着色语言是 WGSL，它的目标不是要与 GLSL 兼容，它是对现代着色器语言的重新设计。详情可以从 WGSL 规范[5] 了解。wgpu 里使用的 WGSL 转译工具叫 naga[6]，性能相比于其他转译工具快十倍。如果你学过 Rust 语言，你会发现 WGSL 的语法和 Rust 语言十分相近。

当前 Rust UI 框架如何使用 GPU 渲染

当前 Rust 生态已经涌现出一些比较优秀的自带 GPU 渲染的 GUI 框架，比如 Makepad[7] 、slint[8]和 egui[9] ，甚至还有一些 Rust 实现的终端也利用 GPU 来加速渲染，比如 Wezterm[10] 和 rio[11]。

Makepad 实现了自己的着色 DSL 语言。因为 Makepad 想要实现的目标是让设计和代码分离，它专门设计了一套 Live System 来有效地对界面实时热更新而不需要重新编译 Rust 代码。着色语言 DSL 需要直接嵌入到这套 Live System 来使用。这套着色 DSL 语言底层绑定了多个 GPU 硬件平台的图形接口，并且对这些 GPU 硬件平台做了优先级支持。现在 Makepad 中 Tier 1 级别支持的是 Quest VR 头显，Tier 2级别支持的是 Qualcomm（骁龙，移动和 AI优势），Tier 3 支持的是 Intel 等。

Slint 通过 Rust 第三方库 femtovg（基于 grow 库，一个 GL 接口绑定库）来支持 GPU 渲染。为什么不考虑使用 wgpu 呢？原因之一是因为可能会失去在iOS/iPadOS/macOS上使用Safari运行演示的能力，因为它仅支持WebGL1；原因之二是因为 slint 在实现之初采用的渲染机制是通过一种遍历树来对 GL 命令进行遍历，类似于 Qt Quick 的机制。也许在未来会考虑迁移到 wgpu。

egui 目前通过 egui-wgpu 和 egui-grow 来支持 WebGPU 和 OpenGL 等多个后端，这方面还是比较成熟的。

Wezterm 和 rio 则是利用 WebGPU 来加速渲染的终端，使用了 wgpu。但是 Wezterm 也允许你通过 lua 脚本来配置使用 OpenGL 。

以上就是这些 Rust 生态中提供渲染机制的 UI 框架或终端使用 GPU 渲染的方式。基本上 wgpu 还是比较常用的。除非有自身的特殊需求，像 Makepad 那样，就需要自己实现着色语言了。当然，我们最希望的还是直接用 Rust 来编写着色语言，目前最接近这个目标的是使用 wgpu(WebGL)。

大模型与 Rust GPU 编程

除了图形处理之外，深度学习和大模型训练领域也是非常依赖于 GPU 的。如果 Rust 能够方便地支持 GPU 编程，那对于 Rust 在人工智能领域的应用也将打开一片天地。

目前深度学习使用 GPU 主要是用 CUDA（Compute Unified Device Architecture）来利用 NVIDIA GPU 进行高性能并行计算。因为深度学习依赖于大量的矩阵和向量运算，这些运算可以在 GPU 上高效地并行处理。在训练大型神经网络模型时，CUDA 可以显著加速计算过程。它通过优化数据传输和执行大量的并行数学运算来减少模型训练所需的时间。CUDA 被广泛支持于各种深度学习框架，如 TensorFlow、PyTorch 和 MXNet。这些框架利用 CUDA 加速后端来提高训练和推理的性能。

在机器学习中，传统着色器语言（如 GLSL 或 HLSL）通常不直接用于模型训练。这些语言主要设计用于图形渲染，而非通用计算。但是在计算机视觉方面可能会有一些应用。

Rust 语言生态中有一些 CUDA 绑定库，比如 Rust-CUDA[12] ，该库提供了rustc_codegen_nvvm 这样一个 rustc 后端，针对 NVVM IR（LLVM IR 的一个子集）进行编译。它生成可以由 CUDA 驱动 API 加载并在 GPU 上执行的高度优化的 PTX 代码。然而，该库已经停止维护超过一年了。

今年大模型 ChatGPT 火了之后，Rust 生态出现了一个完全用 Rust 实现的深度学习框架 burn[13]，创建这个新框架的动机是为了构建一个适应各种用户的多功能框架，包括研究人员、机器学习工程师和低级软件工程师。

Burn 现在支持 wgpu 和 torch-gpu，来进行深度学习的 GPU 并行计算。利用 wgpu 计算着色器来高效处理不同类型的 GPU 的操作，而不像 CUDA 只适用于 Nvidia 的 GPU。但 Burn 并不完全依赖于 wgpu，依然计划在某个时候添加一个仅支持 CUDA 的后端，以在 Nvidia GPU上实现绝对的性能。当前 wgpu 的性能相比 CUDA 还是差一些。Wgpu 无法直接利用供应商特定的功能，如张量计算核心（Tensor Cores）。

另一个机器学习框架是由 Hugging Face 推出的 极简机器学习框架 candle[14] ，它是一个专注于性能（包括GPU支持）和易用性的 Rust 最小化机器学习框架。Candle 旨在支持无服务器推理（Serverless），这是一种在不需要管理任何基础设施的情况下运行机器学习（ML）模型的方式。

candle-core 中通过 cudarc[15] 第三方库来支持 CUDA ，该库是对 CUDA API 的安全 Rust 绑定，看起来是最近才发布的新库。candle 当前不支持 wgpu，但是看起来 wgpu 的支持正在讨论中，见 #344[16]。该 issue 中有人评论到：“通过对不同的 GPGPU 性能和使用 GLSL 的Vulkan 进行了一些初步测试，发现在相同的优化技巧下，Vulkan 的性能可以与 CUDA 相媲美，而使用 WGSL 的 WGPU 很早就达到了瓶颈。此外，WebGPU 也支持 GLSL，所以我们不仅可以有一个 WebGPU 后端，还可以有一个 Vulkan 后端”。

后记

通过以上对 Rust 生态中 GPU 编程的现状的探索，我认为 rust-gpu 中实现的 SPIR-V 编译后端对于 Rust 占据 GPU 编程生态位一席是非常重要的。另外一个安全且稳定维护的 CUDA Rust 库也是非常重要的。希望借助EmbarkStudios 和 HuggingFace 两家商业公司的力量，来打通 Rust 和 GPU 。

感谢阅读。

参考

https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24998

https://jinleili.github.io/learn-wgpu-zh/

参考资料

[1]

Phoronix: https://www.phoronix.com/news/NAK-Merged-Mesa-24.0

[2]

rust-gpu: https://github.com/EmbarkStudios/rust-gpu

[3]

wgpu: https://github.com/gfx-rs/wgpu

[4]

wgpu-native: https://github.com/gfx-rs/wgpu-native

[5]

WGSL 规范: https://www.w3.org/TR/WGSL/

[6]

naga: https://github.com/gfx-rs/naga

[7]

Makepad: https://github.com/makepad/makepad

[8]

slint: https://slint.rs/

[9]

egui: https://github.com/emilk/egui

[10]

Wezterm: https://github.com/wez/wezterm

[11]

rio: https://github.com/raphamorim/rio

[12]

Rust-CUDA: https://github.com/Rust-GPU/Rust-CUDA

[13]

burn: https://github.com/Tracel-AI/burn

[14]

candle: https://github.com/huggingface/candle

[15]

cudarc: https://github.com/coreylowman/cudarc

[16]

#344: https://github.com/huggingface/candle/issues/344