WebGPU 計算管線、計算着色器(通用計算)入門案例:2D 物理模擬


原文譯名:WebGPU - 專注於處理核心(GPU Cores),而不是繪圖畫布(Canvas)

原文發佈於 2022年3月8日,傳送門 https://surma.dev/things/webgpu

這篇東西非常長,不計代碼字符也有1w字,能比較好理解 WebGPU 的計算管線中的各個概念,並使用一個簡單的 2D 物理模擬程序來理解它,本篇重點是在計算管線和計算着色器,繪圖部分使用 Canvas2D 來完成。


WebGPU 是即將推出的 WebAPI,你可以用它訪問圖形處理器(GPU),它是一種底層接口。

原作者對圖形編程沒有多少經驗,他是通過研究 OpenGL 構建遊戲引擎的教程來學習 WebGL 的,還在 ShaderToy 上學習 Inigo Quilez 的例子來研究着色器。因此,他能在 PROXX 中創建背景動畫之類的效果,但是他表示對 WebGL 並不太滿意。別急,下文馬上會解釋。

當作者開始注意 WebGPU 後,大多數人告訴他 WebGPU 這東西比 WebGL 多很多條條框框。他沒考慮這些,已經預見了最壞的情況,他儘可能找了一些教程和規範文檔來看,雖然彼時並不是很多,因爲他找的時候 WebGPU 還在早期制定階段。不過,他深入之後發現 WebGPU 並沒有比 WebGL 多所謂的“條條框框”,反而是像見到了一位老朋友一樣熟悉。

所以,這篇文章就是來分享學到的東西的。

作者明確指出,他 不會 在這裏介紹如何使用 WebGPU 繪製圖形,而是要介紹 WebGPU 如何調用 GPU 進行它本身最原始的計算(譯者注:也就是通用計算)。

他覺得已經有很多資料介紹如何用 WebGPU 進行繪圖了,例如 austin 的例子,或許他考慮之後也寫一些繪圖方面的文章。

他在這裏會討論得比較深入,希望讀者能正確、有效地使用 WebGPU,但是他不保證你讀完就能成爲 GPU 性能專家。

絮絮叨叨結束後,準備發車。

1. WebGL

WebGL 是 2011 年發佈的,迄今爲止,它是唯一能在 Web 訪問 GPU 的底層 API,實際上它是 OpenGL ES 2.0 的簡易封裝版以便能在 Web 中使用。WebGL 和 OpenGL 都是科納斯組標準化的,這個工作組是圖形界的 W3C,可以這麼理解。

OpenGL 本身是一個頗具歷史的 API,按今天的標準看,它不算是一個很好的 API,它以內部全局狀態對象爲中心。這種設計可以最大限度減少特定調用的 GPU 的 IO 數據量。但是,這種設計有很多額外的開銷成本。

image

上圖:WebGL 內部全局狀態對象的可視化,源自 WebGL Fundamentals

內部狀態對象,說白了,大多數都是一些指針。調用 OpenGL API 會改變這些指針的指向,所以改變狀態的順序相當重要,這導致了抽象和寫庫的困難程度大大增加。你必須非常清楚地知道你現在馬上要進行的 API 調用需要準備什麼狀態,調用完了還得恢復到之前的值。

他說,他經常會看到一個黑色的畫布(因爲 WebGL 報錯大多數時候就這樣),然後得狂躁地找沒調用哪些 API 沒有正確設置全局狀態。

他承認,他不知道 ThreeJS 是如何做到狀態管理架構的,但是的確做的不錯,所以大多數人會使用 ThreeJS 而不是原生 WebGL,這是主要的原因了。

“不能很好認同 WebGL”這只是對原作者他自己說的,而不是讀者們。他表示,比他聰明的人用 WebGL 和 OpenGL 已經做了不少 nice 的東西,但是他一直不滿意罷了。

隨着機器學習、神經網絡以及加密貨幣的出現,GPU 證明了它可以幹除了畫三角形之外的事情。使用 GPU 進行任意數據的計算,這種被稱爲 GPGPU,但是 WebGL 1.0 的目的並不在於此。如果你在 WebGL 1.0 想做這件事,你得把數據編碼成紋理,然後在着色器中對數據紋理進行解碼、計算,然後重新編碼成紋理。WebGL 2.0 通過 轉移反饋 讓這攤子事情更容易了一些,但是直到 2021 年 9 月,Safari 瀏覽器才支持 WebGL 2.0(大多數瀏覽器 2017 年 1 月就支持了),所以 WebGL 2.0 不算是好的選擇。

儘管如此,WebGL 2.0 仍然沒有改變 WebGL 的本質,就是全局狀態。

2. WebGPU

在 Web 領域外,新的圖形 API 已經逐漸成型。它們向外部暴露了一套訪問顯卡的更底層的接口。這些新的 API 改良了 OpenGL 的侷促性。

主要就是指 DirectX 12、Vulkan、Metal

一方面來說,現在 GPU 哪裏都有,甚至移動設備都有不錯的 GPU 了。所以,現代圖形編程(3D渲染、光追)和 GPGPU 會越來越普遍。

另一方面來看,大多數設備都有多核處理器,如何優化多線程與 GPU 進行交互,是一個重要的課題。

WebGPU 標準制定者注意到了這些現狀,在預加載 GPU 之前要做好驗證工作,這樣才能給 WebGPU 開發者以更多精力專注於壓榨 GPU 的性能。

下一代最受歡迎的 GPU API 是:

  • 科納斯組的 Vulkan
  • 蘋果的 Metal
  • 微軟的 DirectX 12

爲了把這些技術融合並帶到 Web,WebGPU 就誕生了。

WebGL 是 OpenGL 的一個淺層封裝,但是 WebGPU 並沒這麼做。它引入了自己的抽象概念體系,汲取上述 GPU API 的優點,而不是繼承自這些更底層的 API.

原因很簡單,這三個 API 並不是全部都是全平臺通用的,而且有一些他們自己的非常底層的概念,對於 Web 這個領域來說顯得不那麼合理。

相反,WebGPU 的設計讓人感覺“哇,這就是給 Web 設計的”,但是它的的確確又基於你當前機器的 GPU API,抽象出來的概念被 W3C 標準化,所有的瀏覽器都得實現。由於 WebGPU 相對來說比較底層,它的學習曲線會比較陡峭,但是作者表示會盡可能地分解。

2.1. 適配器(Adapter)和設備(Device)

最開始接觸到的 WebGPU 抽象概念是適配器(Adapter)和設備(Device)。

image

上圖:抽象層,從物理 GPU 到邏輯設備。

物理設備就是 GPU 本身,有內置的 GPU(核芯顯卡)和外部 GPU(獨立顯卡)兩種。通常,某個設備一般只有一個 GPU,但是也有兩個或者多個的情況。例如,微軟的 Surface 筆記本就具備雙顯卡,以便操作系統在不同的情況進行切換。

操作系統使用顯卡廠商提供的驅動程序來訪問 GPU;反過來,操作系統也可以用特定的 API(例如 Vulkan 或者 Metal)向外暴露 GPU 的功能。

GPU 是共享資源,它不僅要被各種程序調用,還要負責向顯示器上輸出。這看起來需要一個東西來讓多個進程同時使用 GPU,以便每個進程把自己的東西畫在屏幕上。

對於每個進程來說,似乎看起裏他們對 GPU 有唯一的控制權,但是那只是表象,實際上這些複雜邏輯是驅動程序和操作系統來完成調度的。

適配器(Adapter)是特定操作系統的 API 與 WebGPU 之間的中介。

但是,由於瀏覽器又是一個可以運行多個 Web 程序的“迷你操作系統”,因此,在瀏覽器層面仍需要共享適配器,以便每個 Web 程序感覺上就像唯一控制 GPU 一樣,所以,每個 Web 程序就獲得了再次抽象的概念:邏輯設備(Logical Device)

要訪問適配器對象,請調用 navigator.gpu.requestAdapter(),在寫本文時,這個方法的參數比較少,能讓你選請求的是高性能的適配器(通常是高性能獨顯)還是低功耗適配器(通常是核顯)。

譯者注:本篇討論 WebGPU 的代碼,沒特殊指明,均爲瀏覽器端的 WebGPU JavaScript API.

軟渲染:一些操作系統(諸如小衆 Linux)可能沒有 GPU 或者 GPU 的能力不足,會提供“後備適配器(Fallback Adapter)”,實際上這種適配器是純軟件模擬出來的,它可能不是很快,可能是 CPU 模擬出來的,但是能基本滿足系統運作。

若能請求到非空的適配器對象,那麼你可以繼續異步調用 adapter.requestDevice() 來請求邏輯設備對象。下面是示例代碼:

if (!navigator.gpu) throw Error("WebGPU not supported.");

const adapter = await navigator.gpu.requestAdapter();
if (!adapter) throw Error("Couldn’t request WebGPU adapter.");

const device = await adapter.requestDevice();
if (!device) throw Error("Couldn’t request WebGPU logical device.");

如果沒有任何請求設備的參數,那麼 requestDevice() 會返回一個不匹配任何設備功能要求的設備,即 WebGPU 團隊認爲是合理且對於所有 GPU 都通用的設備對象。

請求設備對象過程中的“限制”見 規範

舉個例子,即使我的 GPU 可以輕易處理 4GB 的數據,返回的設備對象也只允許最大 1GB 的數據,你請求再多也只會返回最大允許 1GB,這樣就算你切換到別的機器上跑代碼,就不會有太多問題。

你可以訪問 adapter.limits 查看物理 GPU 的實際限制情況。也可以在請求設備對象時,傳遞你所需要檢驗的更高限制參數。

2.2. 着色器(Shaders)

如果你用過 WebGL,那麼你應該熟悉頂點着色器和片元(片段)着色器。其實也沒多複雜,常規技術路線就是上載三角形緩衝數據到 GPU,告訴 GPU 緩衝數據是如何構成三角形的。頂點緩衝的每個頂點數據描述了頂點的位置,當然還包括顏色、紋理座標、法線等其它輔助內容。每個頂點都要經過頂點着色器處理,以完成平移、旋轉、透視變形等操作。

讓原作者感到困惑的是“着色器”這個詞,因爲它除了着色之外還有別的作用。但是在很久以前(1980年代後期)來看,這個詞非常合適,它在 GPU 上的功能就是計算出像素的顏色值。而如今,它泛指在 GPU 上運行的任何程序。

GPU 會對三角形進行光柵化處理,計算出每個三角形在屏幕上佔據的像素。每個像素,則交由片段着色器處理,它能獲取像素座標,當然也可以加入一些輔助數據來決定該像素的最終着色。如果使用得當,就能繪製出令人驚歎的 3D 效果。

將緩衝數據傳遞到頂點着色器,然後繼續傳送到片段着色器,最終輸出到屏幕上這一過程,可以簡單的稱之爲管道(或管線,Pipeline),在 WebGPU 中,必須明確定義 Pipeline.

2.3. 管線(Pipeline)

目前,WebGPU 支持兩大管線:

  • 渲染管線
  • 計算管線

顧名思義,渲染管線繪製某些東西,它結果是 2D 圖像,這個圖像不一定要繪製到屏幕上,可以直接渲染到內存中(被稱作幀緩衝)。計算管線則更加通用,它返回的是一個緩衝數據對象,意味着可以輸出任意數據。

在本文的其它部分會專注於計算管線的介紹,因爲作者認爲渲染管線算是計算管線的一種特殊情況。

現在開始算開歷史倒車,計算管線原來其實是爲了創建渲染管線而先做出來的“基礎”,這些所謂的管線在 GPU 中其實就是不同的物理電路罷了。

基於上述理解,倘若未來向 WebGPU 中添加更多類型的管線,例如“光追管線”,就顯得理所當然了。

使用 WebGPU API,管線由一個或多個可編程階段組成,每個階段由一個着色器模塊和一個入口函數定義。計算管線擁有一個計算着色階段,渲染管線有一個頂點着色階段和一個片段着色階段,如下所示是一個計算着色模塊與計算管線:

const module = device.createShaderModule({
  code: `
    @stage(compute) @workgroup_size(64)
    fn main() {
      // ...
    }
  `,
})

const pipeline = device.createComputePipeline({
  compute: {
    module,
    entryPoint: "main",
  },
})

這是 WebGPU 的着色語言(WGSL,發音 /wig-sal/)的首次登場。

WGSL 給作者的初印象是 Rust + GLSL,它有很多類似 Rust 的語法,也有類似 GLSL 一樣的全局函數(如 dot()norm()len() 等),以及類型(vec2mat4x4 等),還有 swizzling 語法(例如 some_vec.xxy)。

瀏覽器會把 WGSL 源碼編譯成底層系統的着色器目標程序,可能是 D3D12HLSL,也可能是 MetalMSL,或者 VulkanSPIR-V.

SPIR-V:是科納斯組標準化出來的開源、二進制中間格式。你可以把它看作並行編程語言中的 LLVM,它支持多種語言編譯成它自己,也支持把自己翻譯到其它語言。

在上面的着色器代碼中,只創建了一個 main 函數,並使用 @stage(compute) 這個特性(Attribute,WGSL 術語)將其標記爲計算着色階段的入口函數。

你可以在着色器代碼中標記多個 @stage(compute),這樣就可以在多個管線中複用一個着色器模塊對象了,只需傳遞不同的 entryPoint 選擇不同的入口函數即可。

但是,@workgroup_size(64) 特性是什麼?

2.4. 並行(Parallelism)

GPU 以延遲爲代價優化了數據吞吐量。想深入這點必須看一下 GPU 的架構,但是作者沒信心講好這塊,所以建議看一看 Fabian Giesen文章

衆所周知,GPU 有非常多個核心構成,可以進行大規模的並行運算。但是,這些核心不像 CPU 並行編程一樣相對獨立運作。首先,GPU 處理核心是分層分組的,不同廠商的 GPU 的設計架構、API 不盡一致。Intel 這裏給了一個不錯的文檔,對他們的架構進行了高級的描述。

在 Intel 的技術中,最小單元被稱作“執行單元(Execution Unit,EU)”,每個 EU 擁有 7 個 SIMT 內核 —— 意思是,它有 7 個以“鎖步”(Lock-step)的方式運行同一個指令的並行計算核。每個內核都有自己的寄存器和調度緩存的指針,儘管執行着相同的操作,但是數據可以是不同的。

所以有時候不推薦在 GPU 上執行 if/else 判斷分支,是因爲 EU 的原因。因爲 EU 遇到分支邏輯的時候,每個內核都要進行 if/else 判斷,這就失去了並行計算的優勢了。

對於循環也是如此。如果某個核心提前完成了計算任務,那它不得不假裝還在運行,等待 EU 內其它核心完成計算。

儘管內核的計算頻率很高,但是從內存中加載數據或者從紋理中採樣像素的時間明顯要更長 —— Fabian 同志說,這起碼要耗費幾百個時鐘週期。這些時間顯然可以拿來算東西。爲了充分利用這些時鐘週期,每個 EU 必須負重前行。

EU 空閒的時候,譬如在等內存的食物過來的時候,它可不會就一直閒下去,它會立馬投入到下一個計算中,只有這下一個計算再次進入等待時,纔會切換回來,切換的過程非常非常短。

GPU 就是以這樣的技術爲代價換來吞吐量的優化的。GPU 通過調度這些任務的切換機制,讓 EU 一直處於忙碌狀態。

image

上圖:Intel 銳炬 Xe 顯卡芯片架構。它被分成 8 個子塊,每個子塊有 8 個 EU;每個 EU 擁有 7 個 SIMT 內核。

不過,根據上圖來看,EU 只是 Intel 顯卡設計架構層級最低的一個,多個 EU 被 Intel 分爲所謂的“子塊(SubSlice)”,子塊中所有的 EU 都可以訪問共有的局部緩存(Shared Local Memory,SLM),大概是 64KB,如果所運行的程序有同步指令,那麼就必須在同一個子塊中運行,因爲這樣才能共享內存。

再往上,子塊就構成了塊(Slice),構成 GPU;對於集成在 CPU 中的 GPU,大約有 170 ~ 700 個內核。對於獨立顯卡,則會有 1500 或以上個內核。

其它廠商也許會用其它的術語,但是架構基本上可以這麼類比理解。

爲了充分利用 GPU 的架構優勢,需要專門寫程序調用,這樣就可以最大限度地壓榨 GPU 的性能。所以,圖形 API 得向外暴露類似的線程模型來調用計算任務。

在 WebGPU API 中,這種線程模型就叫做“工作組(Workgroup)”。

2.5. 工作組(Workgroup)

每個頂點都會被頂點着色器處理一次,每個片元則會被片元着色器處理一次(當然,這是簡單說法,忽略了很多細節)。

而在 GPGPU 中,與頂點、片元類似的概念是需要開發者自己定義的,這個概念叫做 計算項,計算項會被計算着色器處理。

一組計算項就構成了“工作組”,作者稱之爲“工作負載”。工作組中的每個計算項會被同時運行的計算着色器作用。在 WebGPU 中,工作組可以想象成一個三維網格,最小層級的是計算項,計算項構成稍大級別的是工作組,再往上就構成規模更大的工作負載。

image

上圖:這是一個工作負載,其中紅色小立方體由 4³ 個白色小立方體構成,白色小立方是計算項,而紅色小立方體則由這 64 個白色小立方構成,即工作組。

基於上述概念,就可以討論 WGSL 中的 @workgroup_size(x, y, z) 特性了,它的作用很簡單,就是告訴 GPU 這個計算着色器作用的工作組有多大。用上面的圖來說,其實就是紅色小立方的大小。x*y*z 是每個工作組的計算項個數,如果不設某個維度的值,那默認是 1,因此,@workgroup_size(64) 等同於 @workgroup_size(64, 1, 1).

當然,實際 EU 的架構當然不會是這個 3D 網格里面的某個單元。使用這個圖來描述計算項的目的是凸顯出一種局部性質,即假設相鄰的工作組大概率會訪問緩存中相似的區域,所以順次運行相鄰的工作組(紅色小立方)時,命中緩存中已有的數據的機率會更高一些,而無需在再跑去顯存要數據,節省了非常多時間週期。

然而,大多數硬件依舊是順序執行工作組的,所以設置 @workgroup_size(64)@workgroup_size(8, 8) 的兩個不同的着色器實際上差異並不是很大。所以,這個設計上略顯冗餘。

工作組並不是無限維度的,它受設備對象的限制條件約束,打印 device.limits 可以獲取相關的信息:

console.log(device.limits)

/*
{
  // ...
  maxComputeInvocationsPerWorkgroup: 256,
  maxComputeWorkgroupSizeX: 256,
  maxComputeWorkgroupSizeY: 256,
  maxComputeWorkgroupSizeZ: 64,
  maxComputeWorkgroupsPerDimension: 65535,
  // ...
}
*/

可以看到,每個維度上都有最大限制,而且累乘的積也有最大限制。

提示:避免申請每個維度最大限制數量的線程。雖然 GPU 由操作系統底層調度,但如果你的 WebGPU 程序霸佔了 GPU 太久的話,系統有可能會卡死。

那麼,合適的工作組大小建議是多少呢?這需要具體問題具體分析,取決於工作組各個維度有什麼指代含義。作者認爲這答案很含糊,所以他引用了 Corentin 的話:“用 64 作爲工作組的大小(各個維度累乘後),除非你十分清楚你需要調用 GPU 幹什麼事情。”

64 像是個比較穩妥的線程數,在大多數 GPU 上跑得還可以,而且能讓 EU 儘可能跑滿。

2.6. 指令(Command)

到目前爲止,已經寫好了着色器並設置好了管線,剩下的就是要調用 GPU 來執行。由於 GPU 可以是有自己內存的獨立顯卡,所以可以通過所謂的“指令緩衝”或者“指令隊列”來控制它。

指令隊列,是一塊內存(顯示內存),編碼了 GPU 待執行的指令。編碼與 GPU 本身緊密相關,由顯卡驅動負責創建。WebGPU 暴露了一個“CommandEncoder”API 來對接這個術語。

const commandEncoder = device.createCommandEncoder()
const passEncoder = commandEncoder.beginComputePass()
passEncoder.setPipeline(pipeline)
passEncoder.dispatch(1)
passEncoder.end()
const commands = commandEncoder.finish()
device.queue.submit([commands])

commandEncoder 對象有很多方法,可以讓你把某一塊顯存複製到另一塊,或者操作紋理對應的顯存。它還可以創建 PassEncoder(通道編碼器),它可以配置管線並調度編碼指令。

在上述例子中,展示的是計算管線,所以創建的是計算通道編碼器。調用 setPipeline() 設置管線,然後調用 dispatch() 方法告訴 GPU 在每個維度要創建多少個工作組,以備進行計算。

換句話說,計算着色器的調用次數等於每個維度的大小與該維度調用次數的累積。

例如,一個工作組的三個維度大小是 2, 4, 1,在三個維度上要運行 4, 2, 2 次,那麼計算着色器一共要運行 2×4 + 4×2 + 1×2 = 18 次。

順便說一下,通道編碼器是 WebGPU 的抽象概念,它就是文章最開始時作者抱怨 WebGL 全局狀態機的良好替代品。運行 GPU 管線所需的所有數據、狀態都要經過通道編碼器來傳遞。

抽象:指令緩衝也只不過是顯卡驅動或者操作系統的鉤子,它能讓程序調用 GPU 時不會相互干擾,確保相互獨立。指令推入指令隊列的過程,其實就是把程序的狀態保存下來以便待會要用的時候再取出,因爲硬件執行的速度非常快,看起來就是各做各的,沒有受到其它程序的干擾。

跑起代碼,因爲 workgroup_size 特性顯式指定了 64 個工作組,且在這個維度上調用了 1 次,所以最終生成了 64 個線程,雖然這個管線啥事兒都沒做(因爲沒寫代碼),但是至少起作用了,是不是很酷炫?

隨後,我們搞點數據來讓它起作用。

3. 數據交換

如文章開頭所言,作者沒打算直接用 WebGPU 做圖形繪製,而是打算拿它來做物理模擬,並用 Canvas2D 來簡單的可視化。雖然叫是叫物理模擬,實際上就是生成一堆圓幾何,讓它們在平面範圍內隨機運動並模擬他們之間相互碰撞的過程。

爲此,要把一些模擬參數和初始狀態傳遞到 GPU 中,然後跑計算管線,最後讀取結果。

這可以說是 WebGPU 最頭皮發麻的的一部分,因爲有一堆的數據術語和操作要學。不過作者認爲恰好是這些數據概念和數據的行爲模式造就了 WebGPU,使它成爲了高性能的且與設備無關的 API.

3.1. 綁定組的佈局(GPUBindGroupLayout)

爲了與 GPU 進行數據交換,需要一個叫綁定組的佈局對象(類型是 GPUBindGroupLayout)來擴充管線的定義。

首先要說說綁定組(類型是 GPUBindGroup),它是某種管線在 GPU 執行時各個資源的幾何,資源即 Buffer、Texture、Sampler 三種。

而先於綁定組定義的綁定組佈局對象,則記錄了這些資源的數據類型、用途等元數據,使得 GPU 可以提前知道“噢,這麼回事,提前告訴我我可以跑得更快”。

下列創建一個綁定組佈局,簡單起見,只設置一個存儲型(type: "storage")的緩衝資源:

const bindGroupLayout = device.createBindGroupLayout({
  entries: [{
    binding: 1,
    visibility: GPUShaderStage.COMPUTE,
    buffer: {
      type: "storage",
    }
  }]
})

// 緊接着,傳遞給管線
const pipeline = device.createComputePipeline({
  layout: device.createPipelineLayout({
    bindGroupLayouts: [bindGroupLayout]
  }),
  compute: {
    module,
    entryPoint: 'main'
  }
})

binding 這裏設爲了 1,可以自由設定(當然得按順序),它的作用是在 WGSL 代碼中與相同 binding 值的 buffer 變量綁定在一起。

@group(0) @binding(1)
var<storage, write> output: array<f32>;

type 字段是 "storage",即說明這個 Buffer 的類型是存儲型。它還可以設置爲其它的選項,其中 "read-only-storage" 即“只讀存儲型”,即着色器只能讀,但是不能寫這個 Buffer,只讀型緩衝可以優化一些讀寫同步的問題;而 "uniform" 則說明 Buffer 類型是統一數據(Uniform),作用和存儲型差不多(在着色器中值都一樣)。

至此,綁定組佈局對象創建完畢,然後就可以創建綁定組了,這裏就不寫出來了;一旦創建好了對應的綁定組和存儲型 Buffer,那麼 GPU 就可以開始讀取數據了。

但是,在此之前,還有一個問題要討論:暫存緩衝區。

3.2. 暫存緩衝區(Staging Buffer)

這個小節的內容略長,請耐心閱讀。

作者再次強調:GPU 以延遲爲代價,高度優化了數據 IO 性能。GPU 需要相當快的速度向內核提供數據。在 Fabian 他 2011 年的博客中做了一些計算,得出的結論是 GPU 需要維持 3.3 GB/s 的速度才能運行 1280×720 分辨率的紋理的採樣計算。

爲了滿足現在的圖形需求,GPU 還要再快。只有 GPU 的內核與緩衝存儲器高度集成才能實現,這意味着也就難以把這些存儲區交由 CPU 來讀寫。

我們都知道 GPU 有自己的內存,叫顯存,CPU 和 GPU 都可以訪問它,它與 GPU 的集成度不高,一般在電路板的旁邊,它的速度就沒那麼快了。

暫存緩衝區(Staging buffers),是介於顯存和 GPU 之間的緩存,它可以映射到 CPU 端進行讀寫。爲了讀取 GPU 中的數據,要先把數據從 GPU 內的高速緩存先複製到暫存緩衝區,然後把暫存緩衝區映射到 CPU,這樣才能讀取回主內存。對於數據傳遞至 GPU 的過程則類似。

回到代碼中,創建一個可寫的 Buffer,並添加到綁定組,以便計算着色器可以寫入它;同時還創建一個大小一樣的 Buffer 以作爲暫存。創建這些 Buffer 的時候,要用位掩碼來告知其用途(usage),GPU 會根據參數申請、創建這些緩衝區,如果不符合 WebGPU 規則,則拋出錯誤:

const BUFFER_SIZE = 1000
const output = device.createBuffer({
  size: BUFFER_SIZE,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
})
const stagingBuffer = device.createBuffer({
  size: BUFFER_SIZE,
  usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
})

const bindGroup = device.createBindGroup({
  layout: bindGroupLayout,
  entries: [{
    binding: 1,
    resource: {
      buffer: output
    }
  }]
})

注意,createBuffer() 返回的是 GPUBuffer 對象,不是 ArrayBuffer,創建完 Buffer 後還不能馬上寫入或者讀取。爲了實現讀寫 Buffer,需要有單獨的 API 調用,而且 Buffer 必須有 GPUBufferUsage.MAP_READGPUBufferUsage.MAP_WRITE 的用途才能讀或寫。

TypeScript 提示:在各開發環境還未加入 WebGPU API 時,想要獲得 TypeScript 類型提示,還需要安裝 Chrome WebGPU 團隊維護的 @webgpu/types 包到你的項目中。

到目前爲止,不僅有綁定組的佈局對象,還有綁定組本身,現在需要修改通道編碼器部分的代碼以使用這個綁定組,隨後還要把 Buffer 中計算好的數據再讀取回 JavaScript:

const commandEncoder = device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(pipeline)
passEncoder.setBindGroup(0, bindGroup)
passEncoder.dispatch(1)
passEncoder.dispatch(Math.ceil(BUFFER_SIZE / 64))
passEncoder.end()
commandEncoder.copyBufferToBuffer(
  output,
  0, // 從哪裏開始讀取
  stagingBuffer,
  0, // 從哪裏開始寫
  BUFFER_SIZE
)
const commands = commandEncoder.finish()
device.queue.submit([commands])

await stagingBuffer.mapAsync(
  GPUMapMode.READ,
  0, // 從哪裏開始讀,偏移量
  BUFFER_SIZE // 讀多長
 )
const copyArrayBuffer = stagingBuffer.getMappedRange(0, BUFFER_SIZE)
const data = copyArrayBuffer.slice()
stagingBuffer.unmap()
console.log(new Float32Array(data))

稍前的代碼中,管線對象藉助管線佈局添加了綁定組的局對象,所以如果在通道編碼的時候不設置綁定組就會引起調用(dispatch)失敗。

在計算通道 end() 後,指令編碼器緊接着觸發一個緩衝拷貝方法調用,將數據從 output 緩衝複製到 stagingBuffer 緩衝,最後才提交指令編碼的指令緩衝到隊列上。

GPU 會沿着隊列來執行,沒法推測什麼時候會完成計算。但是,可以異步地提交 stagingBuffer 緩衝的映射請求;當 mapAsync 被 resolve 時,stagingBuffer 映射成功,但是 JavaScript 仍未讀取,此時再調用 stagingBuffer.getMappedRange() 方法,就能獲取對應所需的數據塊了,返回一個 ArrayBuffer 給 JavaScript,這個返回的緩衝數組對象就是顯存的映射,這意味着如果 stagingBuffer 的狀態是未映射時,返回的 ArrayBuffer 也隨之沒有了,所以用 slice() 方法來拷貝一份。

顯然,可以在控制檯看到輸出效果:

image

上圖:湊合,但是說明了一個問題,那就是從 GPU 顯存中把這堆 0 給拿下來了

或許,製造點 0 之外的數據會更有說服力。在進行高級計算之前,先搞點人工數據到 Buffer 中,以證明計算管線確實按預期在運行:

@group(0) @binding(1)
var<storage, write> output: array<f32>;

@stage(compute) @workgroup_size(64)
fn main(

  @builtin(global_invocation_id)
  global_id : vec3<u32>,

  @builtin(local_invocation_id)
  local_id : vec3<u32>,

) {
  output[global_id.x] =
    f32(global_id.x) * 1000. + f32(local_id.x);
}

前兩行聲明瞭一個名爲 output 的模塊範圍的變量,它是一個 f32 元素類型的數組。它的兩個特性聲明瞭來源,@group(0) 表示從第一個(索引爲 0)綁定組中獲取第 1 個綁定資源。output 數組是動態長度的,會自動反射對應 Buffer 的長度。

WGSL 變量:與 Rust 不同,let 聲明的變量是不可變的,如果希望變量可變,使用 var 聲明

接下來看 main 函數。它的函數簽名有兩個參數 global_idlocal_id,當然這兩個變量的名稱隨你設定,它們的值取決於對應的內置變量 global_invocation_idlocal_invocation_id,分別指的是 工作負載 中此着色器調用時的全局 x/y/z 座標,以及 工作組 中此着色器調用時的局部 x/y/z 座標。

image

上圖:三個計算項,a、b、c,用綠色字母標註。

這張圖中使用的工作組大小是 @workgroup_size(4, 4, 4),使用圖中的座標軸順序,那麼對於圖中的 a、b、c 計算項:

  • a:local_id = (x=0, y=0, z=0)global_id = (x=0, y=0, z=0)
  • b:local_id = (x=0, y=0, z=0)global_id = (x=4, y=0, z=0)
  • c:local_id = (x=1, y=1, z=0)global_id = (x=5, y=5, z=0)

而對於我們的例子來說,工作組的大小被設爲 @workgroup_size(64, 1, 1),所以 local_id.x 的取值範圍是 0 ~ 63. 爲了能檢查 local_idglobal_id,作者把這兩個值進行編碼,合成一個數字;注意,WGSL 類型是嚴格的,local_idglobal_id 都是 vec3<u32>,因此要顯式地轉換爲 f32 類型才能寫入 output 緩衝區。

image

上圖:GPU 寫入的實際值,注意 local_id 是 63 爲循環的終點,而 global_id 則仍舊在繼續編碼

上圖證明了計算着色器確實向緩衝區輸出了值,但是很容易發現這些數字看似是沒什麼順序的,因爲這是故意留給 GPU 去做的。

3.3. 過度調度

你可能會注意到,計算通道編碼器的調度方法調度次數 Math.ceil(BUFFER_SIZE / 64) * 64 這個值,算出來就是 1024

passEncoder.dispatch(Math.ceil(BUFFER_SIZE / 64))

這直接導致着色器代碼中 global_id.x 的取值能取到 1024,大於 Buffer 的長度 1000.

不過還好,WGSL 是有保護超出數組索引範圍的機制的,即一旦發生對數組索引越界的寫入,那麼總是會寫入最後一個元素。這樣雖然可以避免內存訪問錯誤,但是仍有可能會生成一些無效數據。譬如,你把 JavaScript 端返回的 Float32Array 的最後 3 個元素打印出來,它們是 247055248056608032;如何避免因數組索引越界而可能發生的無效數據問題呢?可以用衛語句提前返回:

fn main( /* ... */ ) {
  if (global_id.x >= arrayLength(&output)) {
    return;
  }
  
  output[global_id.x] = f32(global_id.x) * 100. + f32(local_id.x)
}

若讀者感興趣,可以運行這個例子看效果。

3.4. 麻煩的結構體(內存地址對齊問題)

還記得目標嗎?是在 2D 的 Canvas 中移動一些圓,並讓他們激情地碰撞。

所以,每個圓都要有一個半徑參數和一個座標參數,以及一個速度矢量。可以繼續用 array<f32> 來表示上述數據,例如第一個數字是 x 座標,第二個數字是 y 座標,以此類推。

然而,這看起來有點蠢,WGSL 是允許自定義結構體的,把多條數據關聯在一個結構內。

注意:如果你知道什麼是內存對齊,你可以跳過本小節;如果你不知道,作者也沒打算仔細解釋,他打算直接展示爲什麼要這麼做。

因此,定義一個結構體 Ball,表示 2D 中的圓,並使用 array<Ball> 表示一系列的 2D 圓球。

使用結構體,就不得不討論內存對齊問題。

struct Ball {
  radius: f32;
  position: vec2<f32>;
  velocity: vec2<f32>;
}

@group(0) @binding(1)
var<storage, write> output: array<Ball>;

@stage(compute) @workgroup_size(64)
fn main(
  @builtin(global_invocation_id) global_id: vec3<u32>,
  @builtin(local_invocation_id) local_id: vec3<u32>,
) {
  let num_balls = arrayLength(&output);
  if (global_id.x >= num_balls) {
    return;
  }
  
  output[global_id.x].radius = 999.;
  output[global_id.x].position = vec2<f32>(global_id.xy);
  output[global_id.x].velocity = vec2<f32>(local_id.xy);
}

你可以運行這個代碼,打開控制檯可以看到:

image

上圖:因爲內存對齊的原因,這個 TypedArray 有明顯的數據填充現象

着色器代碼首先把數據 999.0 寫入到結構體的第一個字段 radius 中,以便於觀察兩個結構的分隔界限;但是,這個打印的 Float32Array 中,兩個 999 數字之間,實際上跨越了 6 個數字,譬如上圖中 0~5 位數字是 999, 0, 0, 0, 0, 0,緊隨其後的 6~11 位數字是 999, 0, 1, 0, 1, 0,這就意味着每個結構體都佔據了 6 個數字,但是 Ball 結構體明明只需要 5 個數字即可存儲:radiusposition.xposition.yvelocity.xvelocity.y. 很明顯,每個 radius 後面都塞多了一個 0,這是爲什麼呢?

原因就是內存對齊。每一種 WGSL 中的數據類型都要嚴格執行對齊要求

若一個數據數據類型的對齊尺度是 N(字節),則意味着這個類型的數據值只能存儲在 N 的倍數的內存地址上。舉個例子,f32 的對齊尺度是 4(即 N = 4),vec2<f32> 的對齊尺度是 8(即 N = 8).

假設 Ball 結構的內存地址是從 0 開始的,那麼 radius 的存儲地址可以是 0,因爲 0 是 4 的倍數;緊接着,下個字段 positionvec2<f32> 類型的,對齊尺度是 8,問題就出現了 —— 它的前一個字段 radius 空閒地址是第 4 個字節,並非 position 對齊尺度 8 的倍數,爲了對齊,編譯器在 radius 後面添加了 4 個字節,也就是從第 8 個字節開始才記錄 position 字段的值。這也就說明了控制檯中看到 999 之後的數字爲什麼總是 0 的原因了。

現在,知道結構體在內存中是如何分佈字節數據的了,可以在 JavaScript 中進行下一步操作了。

3.5. 輸入輸出

我們已經從 GPU 中讀取到數據了,現在要在 JavaScript 中解碼它,也就是生成所有 2D 圓的初始狀態,然後再次提交給 GPU 運行計算着色器,讓它“動起來”。初始化很簡單:

let inputBalls = new Float32Array(new ArrayBuffer(BUFFER_SIZE))
for (let i = 0; i < NUM_BALLS; i++) {
  inputBalls[i * 6 + 0] = randomBetween(2, 10) // 半徑
  inputBalls[i * 6 + 1] = 0 // 填充用
  inputBalls[i * 6 + 2] = randomBetween(0, ctx.canvas.width) // x座標
  inputBalls[i * 6 + 3] = randomBetween(0, ctx.canvas.height) // y座標
  inputBalls[i * 6 + 4] = randomBetween(-100, 100) // x 方向速度分量
  inputBalls[i * 6 + 5] = randomBetween(-100, 100) // y 方向速度分量
}

小技巧:如果你以後的程序用到了更復雜的數據結構,使用 JavaScript 拼湊這些字節碼會非常麻煩,你可以用 Google 的 buffer-backed-object 庫去創建複雜的二進制數據(類似序列化)。

還記得如何把 Buffer 傳遞給着色器嗎?不記得的回去看看上文。只需要調整一下計算管線的綁定組佈局即可接收新的 Buffer:

const bindGroupLayout = device.createBindGroupLayout({
  entries: [
    {
      binding: 0,
      visibility: GPUShaderStage.COMPUTE,
      buffer: {
        type: 'read-only-storage'
      }
    },
    {
      binding: 1,
      visibility: GPUShaderStage.COMPUTE,
      buffer: {
        type: 'storage'
      }
    }
  ]
})

然後創建一個新的綁定組來傳遞初始化後的 2D 圓球數據:

const input = device.createBuffer({
  size: BUFFER_SIZE,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
})

const bindGroup = device.createBindGroup({
  layout: bindGroupLayout,
  entries: [
    {
      binding: 0,
      resource: {
        buffer: input // 輸入初始化數據
      }
    },
    {
      binding: 1,
      resource: {
        buffer: output
      }
    }
  ]
})

就像讀取數據一樣,從技術角度來看,爲了輸入初始化的 2D 圓球數據,要創建一個可映射的暫存緩衝區 input,作爲着色器讀取數據的容器。

WebGPU 提供了一個簡單的 API 便於我們把數據寫進 input 緩衝區:

device.queue.writeBuffer(input, 0, inputBalls)

就是這麼簡單,並不需要指令編碼器 —— 也就是說不需要藉助指令緩衝,writeBuffer() 是作用在隊列上的。

device.queue 對象還提供了一些方便操作紋理的 API.

現在,在着色器代碼中要用新的變量來與這個新的 input 緩衝資源綁定:

// ... Ball 結構體定義 ...

@group(0) @binding(0)
var<storage, read> input: array<Ball>;

// ... output Buffer 的定義

let TIME_STEP: f32 = 0.016;

@stage(compute) @workgroup_size(64)
fn main(
  @builtin(global_invocation_id)
  global_id: vec3<u32>
) {
  let num_balls = arrayLength(&output);
  if (global_id.x >= num_balls) {
    return;
  }
  
  // 更新位置
  output[global_id.x].position = 
    input[global_id.x].position +
    input[global_id.x].velocity * TIME_STEP;
}

希望大部分着色器代碼你能看得懂。

最後要做的,只是把 output 緩衝再次讀取回 JavaScript,寫一些 Canvas2D 的可視化代碼把 Ball 的運動效果展示出來(需要用到 requestAnimationFrame),你可以看示例效果:demo

4. 性能

3.5 小節最後演示的代碼只是能讓 Ball 運動起來,還沒有特別複雜的計算。在進行性能觀測之前,要在着色器中加一些適當的物理計算。

作者就不打算解釋物理計算了,寫到這裏,博客已經很長了,但是他簡單的說明了物理效果的核心原理:每個 Ball 都與其它的 Ball 進行碰撞檢測計算。

如果你十分想知道,可以看看最終的演示代碼:final-demo,在 WGSL 代碼中你還可以找到物理計算的資料連接。

作者並未優化物理碰撞算法,也沒有優化 WebGPU 代碼,即使是這樣,在他的 MacBook Air(M1處理器)上表現得也很不錯。

當超過 2500 個 Ball 時,幀數才掉到 60 幀以下,然而使用 Chrome 開發者工具去觀測性能信息時,掉幀並不是 WebGPU 的問題,而是 Canvas2D 的繪製性能不足 —— 使用 WebGL 或 WebGPU 繪圖就不會出現這個問題了。

image

上圖:即使是 14000 個 Ball,WebGPU 在 M1 處理器的 MBA 筆記本上也才用了 16 毫秒的單幀計算時間

作者關閉了 Canvas2D 繪圖,加入 performance.measure() 方法來查看 16毫秒之內究竟可以模擬多少個 Ball 的物理計算。

這性能表現還是沒有優化過的,已經讓作者爲之陶醉。

5. 穩定性與可用性

WebGPU 已經開發了蠻久了,作者認爲制定規範的人希望 API 是穩定的。

話是這麼說沒錯,但是 WebGPU API 目前只能跑在 Chrome 類瀏覽器和 FireFox 瀏覽器上,對 Safari 保持樂觀態度 —— 雖然寫本文時,Safari TP(技術預覽)還沒什麼東西能看。

在穩定性表現上,即使是寫文章的這段時間裏,也是有變化的。

例如,WGSL 着色器代碼的特性語法,從雙方括號改爲 @ 符號:

[[stage(compute), workgroup_size(64)]]
↓
@stage(compute) @workgroup_size(64)

對通道編碼器結束的方法,Firefox 瀏覽器仍然是 endPass(),而 Chrome 類瀏覽器已經改爲最新的 end().

規範中還有一些內容也並不是完全實現在所有瀏覽器上的,用於移動設備的 API 以及部分着色器常量就是如此。

基本上,WebGPU 進入 stable 階段後,不排除會發生很多重大變化。

總結

“在 Web 上能直接使用 GPU”這種現代的 API 看起來很好玩。在經歷過最初的陡峭學習曲線後,作者認爲真的可以使用 JavaScript 調用 GPU 進行大規模並行運算了。

wgpu 是使用 Rust 實現的 WebGPU,你可以在瀏覽器之外使用 Rust 語言調用 WebGPU 規範的 API;wgpu 還支持編譯到 WebAssembly,你甚至可以使用 Rust 的 wgpu 編寫 wasm,然後再放到瀏覽器運行高性能的代碼。

還有個有趣的東西:Deno 藉助 wgpu,內置了 WebGPU 的支持。

如果你有啥問題,你可以去 WebGPU Matrix 頻道(國內可能訪問不太通暢)提問,那裏有一些 WebGPU 的用戶、瀏覽器工程師和制定規範的人。

感謝 Brandon Jones 校對本文,感謝 WebGPU Matrix 頻道解惑。

也感謝原作者分享這篇長文。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章