Rust目前在GPU編程生態(tài)方面的現(xiàn)狀和前景

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

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

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

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

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

Mesa有兩個作用：

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

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

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

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

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

NVK 與其他的 Nouveau 驅(qū)動非常不同，因為它是從頭開始編寫的。

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

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

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

靜態(tài)單賦值形式（SSA）是一種編譯器中間表示，使每個變量只被賦值一次。這簡化了許多編譯器優(yōu)化，因為變量的值在它們的生命周期內(nèi)保持不變。在轉(zhuǎn)換為 SSA 形式時，編譯器會重寫代碼，使得每個變量的每個賦值操作都有一個唯一的變量名。使得數(shù)據(jù)流分析更加直接和高效，因為每個變量的定義和使用都是顯而易見的。

“

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

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

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

此次 Rust 實現(xiàn)的 NAK 編譯后端大約研發(fā)了半年，長期計劃是成為 NVK 的編譯器。這并不意味著可以直接使用 Rust 來編寫著色器程序，因為它只是一個編譯后端。然而，這也算是朝 Rust 直接進行 GPU 編程更進了一步。

Rust 作為 GPU 著色器語言的前景

圖形渲染機制簡單來說是這樣的：

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

常見的著色語言有 :

DirectX 使用HLSL（High Level Shading Language）

Metal 使用MSL（Metal Shading Language）

OpenGL 使用GLSL（OpenGL Shading Language）

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

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

rust-gpu 項目

國外知名游戲工作室 EmbarkStudios 一直在維護一個 rust-gpu[2] 項目。它旨在為 Rust 編譯器打造一個 spir-v 的編譯后端。

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

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

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

wgpu

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

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

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

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

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

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

當(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 實現(xiàn)的終端也利用 GPU 來加速渲染，比如 Wezterm[10] 和 rio[11]。

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

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

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

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

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

大模型與 Rust GPU 編程

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

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

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

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

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

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

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

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

后記

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

編輯：黃飛