話說，程序員三大浪漫，操作系統、編譯器和圖形處理。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有兩個作用：

對接各種 GPU 硬件，將應用層對 GL API 的調用轉換到對硬件 GPU 的調用上；

各種 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 。

編輯：黃飛