如何玩轉 WebGL 并行計算

作者 | 滄東

來源 | 阿裡技術公衆号

如今在 Web 端使用 WebGL 進行高性能計算已有不少實踐，例如在端智能領域中的 tensorflow.js，再比如可視化領域中的 Stardust.js。在本文中，我們将介紹以下内容：

• 使用 GPU 進行通用計算（GPGPU）的曆史
• 目前在 Web 端使用圖形 API 實作 GPGPU 的技術原理，以及前端開發者可能遇到的難點
• 相關業界實踐，包括布局計算、動畫插值等
• 局限性與未來展望

一 什麼是 GPGPU

由于硬體結構不同，GPU 與 CPU 擅長執行不同類型的計算任務。CPU 通過複雜的 Cache 設計實作低延遲，包含複雜的控制邏輯（分支預測），ALU 隻占一小部分。而 GPU 為高吞吐量而生，包含大量 ALU。是以在單指令流多資料流（SIMD）場景下，GPU 的運算速度遠超 CPU，并且這種差距還在不斷拉大。

而一些現代 GPU 上甚至有專門負責張量計算、光線追蹤的硬體（Tensor/RT Core），例如 Nvidia 的圖靈架構。這使得在處理這些計算複雜度極高的任務時能獲得更大的性能提升。

這裡就需要引出一個概念，用 GPU 進行除渲染外的通用計算：General-Purpose computation on Graphics Processing Units，即 GPGPU。

自 2002 年提出以來，在實時加解密、圖檔壓縮、随機數生成等計算領域都能看到它的身影，GPU Gems/Pro 上也有專門的章節介紹。經由 Nvidia 提出的 CUDA(Compute Unified Device Architecture) 這一統一計算架構，開發者可以使用 C、Java、Python 等語言編寫自己的并行計算任務代碼。

那麼在 Web 端我們應該如何使用 GPU 的計算能力呢？

二 用 WebGL 實作并行計算的原理

在現代化的圖形 API（Vulkan/Metal/Direct3D）中提供了 Compute Shader 供開發者編寫計算邏輯。考慮到 WebGPU 仍在開發中，目前在 Web 端能使用的圖形渲染 API 隻有 WebGL1/2，它們都不支援 Compute Shader（WebGL 2.0 Compute 已廢棄），是以隻能“曲線救國”。在本文的最後一節我們将展望未來的技術手段。

我們先忽略具體的 API 用法，從 CPU 和 GPU 的角度看兩者在并行計算過程中是如何協作的，前者也常被稱作 host，後者為 device。第一步為資料初始化，需要從 CPU 記憶體中拷貝資料到 GPU 記憶體中，在 WebGL 中會通過紋理綁定完成。第二步 CPU 需要準備送出給 GPU 的指令和資料，完成計算程式的編譯，在 WebGL 中通過調用一系列 API 實作。在第三步中将計算邏輯配置設定給 GPU 各個核心執行，是以這段邏輯也叫做“核函數”。最後把計算結果從 GPU 記憶體中拷貝回 CPU 記憶體，在 WebGL1 中通過讀取紋理中像素值完成。

下面我們從 GPU 程式設計模型和執行模型入手，順便引出線程和線程組的概念，這也是 GPU 可資料并行的關鍵。下圖展示了網格與線程組的層次關系，并不局限于 DirectCompute。

• 通過 dispatch(x, y, z) 配置設定一個 3 維的線程網格（Grid），其中的線程共享全局記憶體空間；
• 網格中包含了許多線程組（Work Group、Thread Group、Thread Block、本地工作組不同叫法），每一個線程組中又包含了許多線程，線程組也是 3 維的，一般在 Shader 中通過 numthreads(x, y, z) 指定。它們可以通過共享記憶體或同步原語進行通信；
• Shader 程式最終會運作在每一個線程上。對于每一個線程，可以擷取自己線上程組中的 3 維坐标，也可以擷取線程組在整個線程網格中的 3 維坐标，以此映射到不同的資料上，實作資料并行的效果；

再回到硬體視角，線程對應 GPU 中的 CUDA 核心，線程組對應 SM（Streaming Multiprocessor），網格就是 GPU。

1 WebGL1 紋理映射

下圖來自「GPGPU 程式設計技術 - 從 GLSL、CUDA 到 OpenCL」，這也是經典的 GPGPU 計算流程。

通常來說圖形渲染 API 最終的輸出目标就是螢幕，顯示渲染結果。但是在 GPGPU 場景中我們隻是希望在 CPU 側讀取最終的計算結果。是以會使用到渲染 API 提供的離屏渲染功能，即渲染到紋理，其中的關鍵技術就是使用幀緩存對象（Framebuffer Object/FBO）作為渲染對象。紋理用來存儲輸入參數和計算結果，是以在建立時我們通常需要開啟浮點數擴充 OES_texture_float，該擴充在 WebGL2 中已經内置。

并行計算發生在光栅化階段，我們将計算邏輯（核函數）寫在 Fragment Shader 中，Vertex Shader 僅負責映射紋理坐标，是以 Geometry 可以使用一個 Quad（4個頂點）或者全屏三角形（3個頂點）。對于每一個像素點來說，它的工作并無變化，平時執行的渲染邏輯此時成了一種計算過程，像素值也成了計算結果。

但這種方式存在一個明顯的限制，對于所有線程，紋理緩存要麼是隻讀的，要麼就是隻寫的，沒法實作一個線程在讀紋理，另一個在寫紋理。本質上是由 GPU 的硬體設計決定的，如果想要實作多個線程同時對同一個紋理進行讀/寫操作，需要設計複雜的同步機制避免讀寫沖突，勢必會影響到線程并行執行的效率。是以在經典 GPGPU 的實作中，通常我們會準備兩個紋理，一個用來儲存輸入資料，一個用來儲存輸出資料。

除此之外，該方法并不支援線程間同步和共享記憶體這些特性，是以一些并行算法無法實作，例如 Bellman-Ford 單源最短路徑算法。

上圖中也提到了乒乓技術，很多算法需要連續運作多次，例如 G6 中使用的布局算法需要疊代多次達到穩定狀态。上一次疊代中輸出的計算結果，需要作為下一次疊代的輸入。在實際實作中，我們會配置設定兩張紋理緩存，每次疊代後對輸入和輸出紋理進行交換，實作類似乒乓的效果。

值得注意的是，由于 readPixels（在 CPU 側讀取紋理中的資料）非常慢，除了擷取最終結果，過程中應當盡可能減少對它的調用，盡可能讓資料留在 GPU 中。

這裡我們不再展開 WebGL1 API 的實際用法，詳細使用方式可以參考相關教程。

2 WebGL2 Transform Feedback

首先不得不提到已廢棄的 WebGL 2.0 Compute（底層為 OpenGL ES 3.1），在草案中能看到例如用于線程間同步的 memoryBarrier 和 shared memory 這些進階特性，但最終工作組還是轉向了 WebGPU。

WebGL2 中提供了另一種在 Vertex Shader 中進行并行計算的手段，即 Transform Feedback，它會跳過光栅化管線是以也不需要 Fragment Shader 參與（實際實作中提供一個空 Shader 即可）。

該方案和 WebGL1 的紋理映射方法有以下不同點：

• 不需要 Fragment Shader 參與，是以可以通過全局變量開啟 gl.enable(gl.RASTERIZER_DISCARD)；
• 計算邏輯寫在 Vertex Shader 中，不再需要晦澀的紋理映射，可以直接使用 Buffer 讀寫資料；
• 讀取結果時可以直接使用 getBufferSubData。不過不變的是，該方法依然很慢；

雖然相比 WebGL1 已經有了不小進步，但依舊缺失 Compute Shader 中的一些重要特性。

同樣，這裡我們也不展開 WebGL2 API 的實際用法，詳細使用方式可以參考相關教程。

三 實作中的難點

即使掌握了以上原理，前端開發者在具體實踐中還是會遇到很大困難。除了圖形 API 和 Shader 本身的學習成本，前端對于 GPU 程式設計模型本身也是比較陌生的。

我們遇到的第一個問題是一個算法是否可并行。有些計算任務非常耗時複雜，但并不能交給 GPU 來做，例如代碼編譯，是以可并行和複雜度并沒有直接關系。關于是否可并行的判斷并無嚴格标準，更多來自經驗以及業界已有的實踐（例如後文會提到的圖布局/分析算法），通常遇到一個 "for every X do Y" 這樣的任務就可以考慮是否能進行資料并行。例如下圖展示了一種單源最短路徑算法，不難發現裡面有周遊每一個節點，針對每一條邊的“松弛”操作，此時我們就可以考慮并行化，讓一個線程處理一個節點。

當我們想把一個已有的可并行算法遷移到 GPU 中時，面臨的第一個問題就是資料結構的設計。GPU 記憶體是線性的，也不存在類似對象這樣結構，是以在遷移算法時不可避免的需要重新設計，如果再考慮到對 GPU 記憶體友好，設計難度會進一步加大。在下面應用示例「關于圖布局/分析算法」一節中将看到關于圖的線形表示。

下一個問題是無論 WebGL1 還是 WebGL2，都缺失了 Compute Shader 中的一些重要特性，是以一些在 CUDA 中已經實作的算法也無法直接移植。關于這個問題在本文最後一節中有詳細的說明。

我們已經反複提到了共享記憶體和同步，這裡舉一個 Reduce 求和的例子幫助讀者了解它們的含義。下圖展示配置設定 16 個線程處理一個長度為 16 的數組，最終由 0 号線程将最終結果輸出到共享記憶體的第一個元素中。該過程可分解為以下步驟：

1. 各個線程從全局記憶體中将資料裝載到共享記憶體内。
2. 進行同步( barrier )，確定對于線程組内的所有線程，共享記憶體資料都是最新的。
3. 在共享記憶體中進行累加，每個線程完成後都需要進行同步。
4. 最後所有線程計算完成後，在第一個線程中把共享記憶體中第一個元素寫入全局輸出記憶體中。

試想如果沒有共享記憶體和同步機制，最終的結果顯然不會是正确的，有點類似并發程式設計中的 mutex，如果沒有讀寫鎖會得到意想不到的混亂結果。

最後，GPU 程式設計中的優化空間很大程度依賴開發者對硬體本身的了解，還是以上面 Reduce 求和為例，在 DirectCompute Optimizations and Best Practices 中能找到基于該版本 5 個以上的優化版本。

另外，GPU 在執行 Shader 時無法中斷，這也帶來了代碼難以調試的問題，很多渲染引擎也同樣面臨這樣的問題，Unity 有 RenderDoc 這樣的工具，WebGL 暫無。

四 應用示例介紹

下面我們着重介紹一些 GPGPU 在可視化領域的應用，它們分别來自圖算法、高性能動畫以及海量資料并行處理場景。

既然是通用計算，我們必然無法覆寫所有領域的計算場景，我們嘗試分析以下計算任務的設計實作思路，希望能給讀者一些啟發，當遇到特定場景的可并行算法時，可以嘗試使用 GPU 加速這個過程。

1 圖算法

布局和分析是圖場景中常見的兩類算法。CUDA 有 nvGRAPH 這樣的高性能圖分析算法庫，包含類似最短路徑、PageRank 等，支援多達 20 億條邊的規模。

在實作具體算法前，我們首先需要思考一個問題，即如何用線性結構表示一個圖。最直覺的資料結構是鄰接矩陣，如下圖所示。如果我們有 6 個節點，就可以用一個 6 x 6 的矩陣表示，有連接配接關系的就在對應元素上 + 1。下圖來自維基百科對于鄰接矩陣的展示。

但這樣的資料結構存在一個明顯的問題，過于稀疏導緻空間浪費，尤其當節點數增多時。鄰接表是更好的選擇，該線性結構分成節點和邊兩部分，充分考慮 GPU 記憶體的順序讀，盡可能壓縮（例如每一個 Edge 的 rgba 分量都存儲了臨接節點的 index）。以斥力計算（G6 的實作）為例，需要周遊除自身外的全部其他節點，這全部都是順序讀操作。同樣的，在計算吸引力時，周遊一個節點的所有邊也都是順序讀。随機讀隻會出現在擷取端點坐标時才會出現。

這裡不展開具體算法實作，遷移 G6 已有布局算法的過程詳見。最終效果依不同算法實作差距很大，效果最好的是 Fruchterman 布局，節點數過千後 GPU 版本有百倍以上的提升，但 GForce 布局在少量節點的情況下甚至不如 CPU 版本。

2 SandDance

SandDance 提供了多元資料在多種布局下流暢切換的效果，它擴充了 Vega 規範，在 2D 場景中增加了深度資訊，同時使用 Deck.gl 做渲染。具體到布局切換使用的技術，Luma.gl（Deck.gl 的底層渲染引擎） 提供了基于 WebGL2 Transform Feedback 的進階封裝，用于在 GPU 中完成動畫和資料變換的插值。對比傳統的在 CPU 中做插值動畫性能要高很多。在通用渲染引擎中，該技術也常用于粒子特效的實作。

3 P4: Portable Parallel Processing Pipelines

P4 緻力于海量資料的處理和渲染，在運作時生成資料聚合和渲染的 Shader 代碼，前者有點類似 tfjs 中的一些 op。值得一提的是通過 WebGL 的 Blending 操作實作了一些 Reduce 操作（例如最大最小值、計數、求和、平均值）。例如在實作 Reduce 求和時使用到的 blendEquation 為 gl.ADD。

五 目前局限性與未來展望

我們可以看出 WebGL 受限于底層 API 能力，在很多計算相關的特性上有不同程度的缺失，導緻很多可并行算法無法實作。另一方面，可視化領域又缺失有不少适合的場景，我們迫切需要下一代能力更強的 Web API。

WebGPU 作為 WebGL 的繼任者，底層依賴各個作業系統上更現代化的圖形 API，提供了更低級的接口，這意味着開發者對 GPU 有更多直接控制以及更少的驅動資源消耗，渲染計算一視同仁。目前已經可以在 Chrome/Edge/Safari 的預覽版本中使用它。

WebGPU 在 Shader 語言的選擇上抛棄了 WebGL 使用的 GLSL，轉向新的 WGSL。在我們關心的計算相關特性上，它提供了 storage/workgroupBarrier 同步方法。有了這些特性，一些算法就可以移植到 Web 端了，例如單源最短路徑等其他圖算法的 CUDA 開源實作，筆者嘗試用 WGSL 實作它。

目前一個更成熟的實踐是，Apache TVM 社群加入了 WebAssembly 和 WebGPU 後端支援。在 MacOS 上可以獲得和直接本地運作 native metal 幾乎一樣的效率。

總之在可預見的未來，這無疑是 Web 端 GPGPU 的最佳選擇。

相關連結：

• https://webglfundamentals.org/webgl/lessons/webgl-gpgpu.html
• https://webgl2fundamentals.org/webgl/lessons/webgl-gpgpu.html
• https://www.lewuathe.com/illustration-of-distributed-bellman-ford-algorithm.html
• http://on-demand.gputechconf.com/gtc/2010/presentations/S12312-DirectCompute-Pre-Conference-Tutorial.pdf
• https://ieeexplore.ieee.org/abstract/document/8468065
• https://tvm.apache.org/2020/05/14/compiling-machine-learning-to-webassembly-and-webgpu

【公開課】了解 Pod 和容器設計模式

本節課程由阿裡雲與 CNCF 共同推出。通過課程學習，您将了解到 Kubernetes 體系中最重要的一個基礎知識點，也就是 Pod 與容器設計模式。阿裡巴巴進階技術專家、CNCF 官方大使張磊将通過真實案例為您闡述“為什麼我們需要 Pod ”、 Pod 的實作機制，以及 Kubernetes 非常提倡的一個概念，叫做容器設計模式。

點選這裡

，即可看到課程～