話說，程序員三大浪漫，操作系統(tǒng)、編譯器和圖形處理。Rust 語(yǔ)言已經(jīng)攻陷了其中兩大浪漫，操作系統(tǒng)和編譯器，那么圖形處理呢？Rust 語(yǔ)言還能“浪”起來嗎?

這周在 Phoronix[1] 網(wǎng)站上得知用 Rust 重寫的 NAK 編譯器已經(jīng)被 Merge 到了 Mesa 24.0 版本中，用于 Nouveau Gallium3D 驅(qū)動(dòng)程序和 NVK Vulkan 驅(qū)動(dòng)程序。這激起了我的好奇心。因?yàn)?GPU 編程是 Rust 語(yǔ)言進(jìn)入圖形處理的關(guān)鍵，所以我想徹底了解一下 Rust 目前在 GPU 編程生態(tài)方面的現(xiàn)狀和前景。

這就是本文的出發(fā)點(diǎn)。

Mesa 相關(guān)背景介紹

Mesa （或稱 Mesa3D） 是一個(gè) OpenGL/Vulkan 的實(shí)現(xiàn)，以及為所有開源圖形驅(qū)動(dòng)提供各種圖形庫(kù)（GL）的入口點(diǎn)。

Mesa有兩個(gè)作用：

對(duì)接各種 GPU 硬件，將應(yīng)用層對(duì) GL API 的調(diào)用轉(zhuǎn)換到對(duì)硬件 GPU 的調(diào)用上；

各種 GL API 的純軟實(shí)現(xiàn)，當(dāng)沒有可用的硬件時(shí)，它可以提供純軟件的 GL API 的實(shí)現(xiàn)；

它可以用于 Linux/Windows/Mac 等系統(tǒng)平臺(tái)。在 Windows 上運(yùn)行時(shí)它提供 OpenGL API over DirectX 的轉(zhuǎn)換。AMD 和 Intel 都提供了對(duì) Mesa 支持的驅(qū)動(dòng)程序。

NVIDIA 有兩種類型的顯卡驅(qū)動(dòng)：一個(gè)是開源的 Nouveau ，另一個(gè)是閉源的 NVIDIA 官方驅(qū)動(dòng)。Mesa 使用了開源的 Nouveau 驅(qū)動(dòng)。而 Gallium3D 是 Mesa 提出的用于簡(jiǎn)化 GPU 驅(qū)動(dòng)開發(fā)的框架。

NVK，是由 Collabora 推出的一個(gè)新的 Mesa 開源驅(qū)動(dòng)程序，為 NVIDIA 顯卡實(shí)現(xiàn) Vulkan 圖形 API。該驅(qū)動(dòng)程序是使用 NVIDIA 發(fā)布的官方頭文件，以及開放的數(shù)據(jù)中心 GPU 和消費(fèi)級(jí) GPU（GTX/RTX）的 GPU 內(nèi)核模塊，從頭開始編寫的。它的目標(biāo)是成為新的主流顯卡驅(qū)動(dòng)。

NVK 與其他的 Nouveau 驅(qū)動(dòng)非常不同，因?yàn)樗菑念^開始編寫的。

nouveau 是一個(gè)主要的 NVIDIA 顯卡的開源驅(qū)動(dòng)程序，已經(jīng)年久失修了，試圖在它的基礎(chǔ)上構(gòu)建是一個(gè)很多人都無法承擔(dān)的任務(wù)。當(dāng)然，它是由有很多才華的工程師開發(fā)的，但是缺乏公司的支持和貢獻(xiàn)者的影響了它的發(fā)展。NVK 旨在克服這些問題，同時(shí)專注于對(duì) Turing 系列及更高版本 GPU 的支持。

由于內(nèi)核的開發(fā)方式，對(duì)于 Kepler、Maxwell 和 Pascal 等較舊的 GPU 的支持可能不會(huì)很容易地加入 NVK。它也許極大地依賴于新內(nèi)核，從而只支持較新的 GPU。同時(shí)，nouveau 內(nèi)核接口與 Vulkan 不兼容，阻礙了對(duì)較舊 GPU 的支持。

而用 Rust 實(shí)現(xiàn)的新的 NAK （Nvidia awesome kompile）編譯器，是一個(gè)為 NVIDIA GPU 設(shè)計(jì)的后端編譯器，專門為處理 GPU 任務(wù)優(yōu)化。NAK 使用靜態(tài)單賦值形式（SSA）和中間表示（NIR），這有助于優(yōu)化代碼并提高 GPU 的執(zhí)行效率。

靜態(tài)單賦值形式（SSA）是一種編譯器中間表示，使每個(gè)變量只被賦值一次。這簡(jiǎn)化了許多編譯器優(yōu)化，因?yàn)樽兞康闹翟谒鼈兊纳芷趦?nèi)保持不變。在轉(zhuǎn)換為 SSA 形式時(shí)，編譯器會(huì)重寫代碼，使得每個(gè)變量的每個(gè)賦值操作都有一個(gè)唯一的變量名。使得數(shù)據(jù)流分析更加直接和高效，因?yàn)槊總€(gè)變量的定義和使用都是顯而易見的。

“

順便說一下，Rust 編譯器（rustc）在編譯過程的 MIR 和 LLVM IR 這兩個(gè)階段也使用了靜態(tài)單賦值形式（SSA）。

NIR 是 Mesa 中提供的一種更加具體的中間語(yǔ)言，它是為了優(yōu)化和簡(jiǎn)化驅(qū)動(dòng)編譯器的工作流程而設(shè)計(jì)的。NIR 設(shè)計(jì)上更接近硬件，旨在作為多個(gè)不同前端（如 GLSL、SPIR-V）和多個(gè)不同后端（如不同的 GPU 驅(qū)動(dòng)）之間的橋梁。NIR 也支持 SSA 等各種優(yōu)化技術(shù)。NIR 使 Mesa 能夠更有效地處理來自不同源的圖形和計(jì)算著色器代碼，為最終在 GPU 上執(zhí)行的代碼生成和優(yōu)化奠定基礎(chǔ)。

借用網(wǎng)上的一張舊圖來說明 Mesa 中的 IR 架構(gòu)：

此次 Rust 實(shí)現(xiàn)的 NAK 編譯后端大約研發(fā)了半年，長(zhǎng)期計(jì)劃是成為 NVK 的編譯器。這并不意味著可以直接使用 Rust 來編寫著色器程序，因?yàn)樗皇且粋€(gè)編譯后端。然而，這也算是朝 Rust 直接進(jìn)行 GPU 編程更進(jìn)了一步。

Rust 作為 GPU 著色器語(yǔ)言的前景

圖形渲染機(jī)制簡(jiǎn)單來說是這樣的：

圖形庫(kù)（比如 OpenGL）將渲染計(jì)算任務(wù)實(shí)時(shí)派發(fā)給 GPU，具體由一種用圖形庫(kù)提供的著色語(yǔ)言（GLSL或 WLSL等）或 SPIR-V（著色語(yǔ)言中間語(yǔ)言標(biāo)準(zhǔn)）編寫的稱為著色器（sharder）的小程序，在 GPU 上編譯運(yùn)行。從基本意義上來說，著色器只是一種把輸入轉(zhuǎn)化為輸出的程序。著色器也是一種非常獨(dú)立的程序，因?yàn)樗鼈冎g不能相互通信。著色語(yǔ)言一般包含一些針對(duì)向量和矩陣操作的有用特性。

常見的著色語(yǔ)言有 :

DirectX 使用HLSL（High Level Shading Language）

Metal 使用MSL（Metal Shading Language）

OpenGL 使用GLSL（OpenGL Shading Language）

Vulkan 使用的著色器必須以 SPIR-V 這種二進(jìn)制字節(jié)碼的格式提供。

所以一般來說，在 Rust 生態(tài)中，需要使用 GPU 進(jìn)行圖形渲染則多半是需要直接使用 GLSL 這類著色語(yǔ)言。這就帶來了極大的不便。幸好，Rust 生態(tài)有一些開源項(xiàng)目，正在致力于改變這一狀況。

rust-gpu 項(xiàng)目

國(guó)外知名游戲工作室 EmbarkStudios 一直在維護(hù)一個(gè) rust-gpu[2] 項(xiàng)目。它旨在為 Rust 編譯器打造一個(gè) spir-v 的編譯后端。

SPIR-V 是一個(gè)為 Vulkan 和 OpenCL 設(shè)計(jì)的中間語(yǔ)言（IL）標(biāo)準(zhǔn)。它是一個(gè)低層次、與硬件無關(guān)的 IR，用于表達(dá)著色器和計(jì)算核心。SPIR-V 設(shè)計(jì)上更靠近硬件執(zhí)行層次，它直接被 GPU 驅(qū)動(dòng)所接受，并轉(zhuǎn)換為特定硬件的機(jī)器代碼。

正如前面 Mesa 的 IR 圖所示，SPIR-V 通常作為著色器語(yǔ)言（如 GLSL）的編譯輸出，然后可以被轉(zhuǎn)換為 NIR 進(jìn)行進(jìn)一步的優(yōu)化和處理。這就意味著，如果 rust-gpu 成熟了，可以直接用 Rust 語(yǔ)言作為著色語(yǔ)言來編寫著色器程序，這樣就可以通過 SPIR-V 轉(zhuǎn)換為 NIR ，進(jìn)一步讓 NAK 編譯后端來處理了。

從這個(gè)角度來看，rust-gpu 這個(gè)項(xiàng)目對(duì)于 Rust GPU 圖形編程渲染生態(tài)還是非常重要的。

wgpu

另外一個(gè)項(xiàng)目是 wgpu[3] ，它提供了一個(gè)安全、跨平臺(tái)的圖形和計(jì)算 API，基于 WebGPU 規(guī)范。WebGPU 是一種新的圖形標(biāo)準(zhǔn)，旨在為現(xiàn)代圖形硬件提供統(tǒng)一的低層次訪問。它被設(shè)計(jì)為更安全、更高效，特別是在 Web 應(yīng)用程序中。

wgpu 是基于 Rust 實(shí)現(xiàn)的，所以它利用 Rust 的安全特性來幫助避免常見的內(nèi)存錯(cuò)誤和并發(fā)問題，這在處理復(fù)雜的圖形任務(wù)時(shí)尤其重要。wgpu 也充分利用了現(xiàn)代 GPU 的能力，提供高效的圖形和計(jì)算性能。它支持最新的圖形技術(shù)，如計(jì)算著色器和高效的資源管理。wgpu 提供了 Rust 風(fēng)格的 API，相比于直接使用 Vulkan 或 Direct3D，它提供了更高級(jí)別的抽象，簡(jiǎn)化了圖形編程的復(fù)雜性。

wgpu不僅可以在 Web 環(huán)境運(yùn)行，還可以在 macOS / iOS、Android、Window 和 Linux 等系統(tǒng)上原生運(yùn)行。隨著 WebGPU 標(biāo)準(zhǔn)的發(fā)展和成熟，wgpu 可能會(huì)成為 Web 和非 Web 應(yīng)用程序中 GPU 編程的首選解決方案。wgpu 對(duì)于游戲開發(fā)、圖形設(shè)計(jì)和可視化、科學(xué)計(jì)算和機(jī)器學(xué)習(xí)等領(lǐng)域非常適合。目前被用于 Firefox、Servo 和 Deno 中 WebGPU 整合的核心。

wgpu 實(shí)際上也提供了 C 語(yǔ)言綁定 (wgpu-native[4])，你可以寫 C/C++ 或其他能與 C 互通的語(yǔ)言來使用它。

wgpu 還有另一個(gè)重要的優(yōu)勢(shì)，那就是可以利用各種強(qiáng)大的桌面端 GPU 調(diào)試工具。在開發(fā)大型 2D/3D 應(yīng)用時(shí)，通過使用命令記錄/回放、幀捕捉、Buffer 視圖等功能，可以快速定位 GPU 層代碼/數(shù)據(jù)的性能瓶頸和程序缺陷。相較于僅依靠瀏覽器提供的有限調(diào)試能力，這些工具能夠事半功倍，幫助開發(fā)者更快地解決問題。

WebGPU 使用的著色語(yǔ)言是 WGSL，它的目標(biāo)不是要與GLSL兼容，它是對(duì)現(xiàn)代著色器語(yǔ)言的重新設(shè)計(jì)。詳情可以從WGSL 規(guī)范[5] 了解。wgpu 里使用的 WGSL 轉(zhuǎn)譯工具叫naga[6]，性能相比于其他轉(zhuǎn)譯工具快十倍。如果你學(xué)過 Rust 語(yǔ)言，你會(huì)發(fā)現(xiàn) WGSL 的語(yǔ)法和 Rust 語(yǔ)言十分相近。

當(dāng)前 Rust UI 框架如何使用 GPU 渲染

當(dāng)前 Rust 生態(tài)已經(jīng)涌現(xiàn)出一些比較優(yōu)秀的自帶 GPU 渲染的 GUI 框架，比如 Makepad[7] 、slint[8]和 egui[9] ，甚至還有一些 Rust 實(shí)現(xiàn)的終端也利用 GPU 來加速渲染，比如 Wezterm[10] 和 rio[11]。

Makepad 實(shí)現(xiàn)了自己的著色 DSL 語(yǔ)言。因?yàn)?Makepad 想要實(shí)現(xiàn)的目標(biāo)是讓設(shè)計(jì)和代碼分離，它專門設(shè)計(jì)了一套 Live System 來有效地對(duì)界面實(shí)時(shí)熱更新而不需要重新編譯 Rust 代碼。著色語(yǔ)言 DSL 需要直接嵌入到這套 Live System 來使用。這套著色 DSL 語(yǔ)言底層綁定了多個(gè) GPU 硬件平臺(tái)的圖形接口，并且對(duì)這些 GPU 硬件平臺(tái)做了優(yōu)先級(jí)支持?，F(xiàn)在 Makepad 中 Tier 1 級(jí)別支持的是 Quest VR 頭顯，Tier 2級(jí)別支持的是 Qualcomm（驍龍，移動(dòng)和 AI優(yōu)勢(shì)），Tier 3 支持的是 Intel 等。

Slint 通過 Rust 第三方庫(kù) femtovg（基于 grow 庫(kù)，一個(gè) GL 接口綁定庫(kù)）來支持 GPU 渲染。為什么不考慮使用 wgpu 呢？原因之一是因?yàn)榭赡軙?huì)失去在iOS/iPadOS/macOS上使用Safari運(yùn)行演示的能力，因?yàn)樗鼉H支持WebGL1；原因之二是因?yàn)?slint 在實(shí)現(xiàn)之初采用的渲染機(jī)制是通過一種遍歷樹來對(duì) GL 命令進(jìn)行遍歷，類似于 Qt Quick 的機(jī)制。也許在未來會(huì)考慮遷移到 wgpu。

egui 目前通過 egui-wgpu 和 egui-grow 來支持 WebGPU 和 OpenGL 等多個(gè)后端，這方面還是比較成熟的。

Wezterm 和 rio 則是利用 WebGPU 來加速渲染的終端，使用了 wgpu。但是 Wezterm 也允許你通過 lua 腳本來配置使用 OpenGL 。

以上就是這些 Rust 生態(tài)中提供渲染機(jī)制的 UI 框架或終端使用 GPU 渲染的方式?；旧?wgpu 還是比較常用的。除非有自身的特殊需求，像 Makepad 那樣，就需要自己實(shí)現(xiàn)著色語(yǔ)言了。當(dāng)然，我們最希望的還是直接用 Rust 來編寫著色語(yǔ)言，目前最接近這個(gè)目標(biāo)的是使用 wgpu(WebGL)。

大模型與 Rust GPU 編程

除了圖形處理之外，深度學(xué)習(xí)和大模型訓(xùn)練領(lǐng)域也是非常依賴于 GPU 的。如果 Rust 能夠方便地支持 GPU 編程，那對(duì)于 Rust 在人工智能領(lǐng)域的應(yīng)用也將打開一片天地。

目前深度學(xué)習(xí)使用 GPU 主要是用 CUDA（Compute Unified Device Architecture）來利用 NVIDIA GPU 進(jìn)行高性能并行計(jì)算。因?yàn)樯疃葘W(xué)習(xí)依賴于大量的矩陣和向量運(yùn)算，這些運(yùn)算可以在 GPU 上高效地并行處理。在訓(xùn)練大型神經(jīng)網(wǎng)絡(luò)模型時(shí)，CUDA 可以顯著加速計(jì)算過程。它通過優(yōu)化數(shù)據(jù)傳輸和執(zhí)行大量的并行數(shù)學(xué)運(yùn)算來減少模型訓(xùn)練所需的時(shí)間。CUDA 被廣泛支持于各種深度學(xué)習(xí)框架，如 TensorFlow、PyTorch 和 MXNet。這些框架利用 CUDA 加速后端來提高訓(xùn)練和推理的性能。

在機(jī)器學(xué)習(xí)中，傳統(tǒng)著色器語(yǔ)言（如 GLSL 或 HLSL）通常不直接用于模型訓(xùn)練。這些語(yǔ)言主要設(shè)計(jì)用于圖形渲染，而非通用計(jì)算。但是在計(jì)算機(jī)視覺方面可能會(huì)有一些應(yīng)用。

Rust 語(yǔ)言生態(tài)中有一些 CUDA 綁定庫(kù)，比如 Rust-CUDA[12] ，該庫(kù)提供了rustc_codegen_nvvm 這樣一個(gè) rustc 后端，針對(duì) NVVM IR（LLVM IR 的一個(gè)子集）進(jìn)行編譯。它生成可以由 CUDA 驅(qū)動(dòng) API 加載并在 GPU 上執(zhí)行的高度優(yōu)化的 PTX 代碼。然而，該庫(kù)已經(jīng)停止維護(hù)超過一年了。

今年大模型 ChatGPT 火了之后，Rust 生態(tài)出現(xiàn)了一個(gè)完全用 Rust 實(shí)現(xiàn)的深度學(xué)習(xí)框架 burn[13]，創(chuàng)建這個(gè)新框架的動(dòng)機(jī)是為了構(gòu)建一個(gè)適應(yīng)各種用戶的多功能框架，包括研究人員、機(jī)器學(xué)習(xí)工程師和低級(jí)軟件工程師。

Burn 現(xiàn)在支持 wgpu 和 torch-gpu，來進(jìn)行深度學(xué)習(xí)的 GPU 并行計(jì)算。利用 wgpu 計(jì)算著色器來高效處理不同類型的 GPU 的操作，而不像 CUDA 只適用于 Nvidia 的 GPU。但 Burn 并不完全依賴于 wgpu，依然計(jì)劃在某個(gè)時(shí)候添加一個(gè)僅支持 CUDA 的后端，以在 Nvidia GPU上實(shí)現(xiàn)絕對(duì)的性能。當(dāng)前 wgpu 的性能相比 CUDA 還是差一些。Wgpu 無法直接利用供應(yīng)商特定的功能，如張量計(jì)算核心（Tensor Cores）。

另一個(gè)機(jī)器學(xué)習(xí)框架是由 Hugging Face 推出的 極簡(jiǎn)機(jī)器學(xué)習(xí)框架 candle[14] ，它是一個(gè)專注于性能（包括GPU支持）和易用性的 Rust 最小化機(jī)器學(xué)習(xí)框架。Candle 旨在支持無服務(wù)器推理（Serverless），這是一種在不需要管理任何基礎(chǔ)設(shè)施的情況下運(yùn)行機(jī)器學(xué)習(xí)（ML）模型的方式。

candle-core 中通過 cudarc[15] 第三方庫(kù)來支持 CUDA ，該庫(kù)是對(duì) CUDA API 的安全 Rust 綁定，看起來是最近才發(fā)布的新庫(kù)。candle 當(dāng)前不支持 wgpu，但是看起來 wgpu 的支持正在討論中，見 #344[16]。該 issue 中有人評(píng)論到：“通過對(duì)不同的 GPGPU 性能和使用 GLSL 的Vulkan 進(jìn)行了一些初步測(cè)試，發(fā)現(xiàn)在相同的優(yōu)化技巧下，Vulkan 的性能可以與 CUDA 相媲美，而使用 WGSL 的 WGPU 很早就達(dá)到了瓶頸。此外，WebGPU 也支持 GLSL，所以我們不僅可以有一個(gè) WebGPU 后端，還可以有一個(gè) Vulkan 后端”。

后記

通過以上對(duì) Rust 生態(tài)中 GPU 編程的現(xiàn)狀的探索，我認(rèn)為 rust-gpu 中實(shí)現(xiàn)的 SPIR-V 編譯后端對(duì)于 Rust 占據(jù) GPU 編程生態(tài)位一席是非常重要的。另外一個(gè)安全且穩(wěn)定維護(hù)的 CUDA Rust 庫(kù)也是非常重要的。希望借助EmbarkStudios 和 HuggingFace 兩家商業(yè)公司的力量，來打通 Rust 和 GPU 。

編輯：黃飛