開始在網路上使用 GPU Compute

本文透過範例探索實驗性 WebGPU API,並協助您開始使用 GPU 執行資料平行運算。

François Beaufort
François Beaufort

背景

如您所知,圖形處理器 (GPU) 是電腦中的電子子系統,專門用於處理圖形。不過,過去 10 年來,此架構已發展出更靈活的架構,讓開發人員得以導入多種演算法,而不僅僅是轉譯 3D 圖形,還能利用 GPU 的獨特架構。這些功能稱為 GPU Compute,而使用 GPU 做為一般用途科學運算的輔助處理器,則稱為一般用途 GPU (GPGPU) 程式設計。

GPU 運算對近期機器學習的蓬勃發展有重大貢獻,因為卷積類神經網路和其他模型可利用架構來提高 GPU 的執行效率。隨著目前 Web Platform 缺少 GPU 運算功能,W3C 的「網路 GPU」社群群正在設計 API,以公開大多數目前裝置可用的新式 GPU API。這個 API 稱為 WebGPU

WebGPU 是一種低階 API,例如 WebGL。如你所見但沒關係。我們希望的是效能。

在本文中,我會將重點放在 WebGPU 的 GPU 運算部分。首先是這樣的,我只是正在刮開表面,以便您開始玩遊戲。我將深入探討即將發布的文章,並深入探討 WebGPU 轉譯 (畫布、紋理等)。

存取 GPU

在 WebGPU 中,您可以輕鬆存取 GPU。呼叫 navigator.gpu.requestAdapter() 會傳回 JavaScript 承諾,該承諾會以 GPU 轉接程式以非同步方式解析。您可以把這個轉接器當成顯示卡。這種套件可以整合 (在與 CPU 相同的晶片上) 或獨立 (通常是 PCIe 卡,效能較高但耗電量更高)。

取得 GPU 轉接器後,請呼叫 adapter.requestDevice() 以取得會與 GPU 裝置解析的承諾,以便您執行一些 GPU 運算。

const adapter = await navigator.gpu.requestAdapter();
if (!adapter) { return; }
const device = await adapter.requestDevice();

這兩種函式都會提供選項,讓您明確指定所需的轉接程式類型 (電源偏好設定) 和裝置 (擴充功能、限制)。為求簡單起見,本文將使用本文的預設選項。

寫入緩衝區記憶體

讓我們一起瞭解如何使用 JavaScript,將資料寫入 GPU 的記憶體中。由於新式網路瀏覽器採用沙箱機制,此程序並不容易。

以下範例說明如何寫入四個位元組,以便從 GPU 存取的緩衝區記憶體。它會呼叫 device.createBuffer(),取用緩衝區的大小及其使用情況。即使這個特定呼叫不需要使用 GPUBufferUsage.MAP_WRITE,但我們明確地要寫入這個緩衝區。如果 mappedAtCreation 設為 true,就能在建立時產生對應的 GPU 緩衝區物件。接著,您可以呼叫 GPU 緩衝區方法 getMappedRange(),藉此擷取相關聯的原始二進位資料緩衝區。

如果您已使用 ArrayBuffer 進行播放,就會很熟悉寫入位元組;請使用 TypedArray 並將值複製到其中。

// Get a GPU buffer in a mapped state and an arrayBuffer for writing.
const gpuBuffer = device.createBuffer({
  mappedAtCreation: true,
  size: 4,
  usage: GPUBufferUsage.MAP_WRITE
});
const arrayBuffer = gpuBuffer.getMappedRange();

// Write bytes to buffer.
new Uint8Array(arrayBuffer).set([0, 1, 2, 3]);

此時,系統會對應 GPU 緩衝區,這代表 GPU 為 CPU 所有,且可透過 JavaScript 讀取/寫入。為了讓 GPU 能夠存取該功能,您必須進行取消對應,就像呼叫 gpuBuffer.unmap() 一樣簡單。

必須採用「對應/未對應」的概念,才能避免同時 GPU 和 CPU 存取記憶體發生競爭狀況。

讀取緩衝區記憶體

現在來瞭解如何將 GPU 緩衝區複製到其他 GPU 緩衝區並讀回。

由於我們是在第一個 GPU 緩衝區寫入,且想將其複製到第二個 GPU 緩衝區,因此需要新的使用旗標 GPUBufferUsage.COPY_SRC。第二個 GPU 緩衝區會以 device.createBuffer() 未對應的狀態建立。其使用旗標為 GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,因為其會做為第一個 GPU 緩衝區的目的地,並在執行 GPU 複製指令後以 JavaScript 讀取。

// Get a GPU buffer in a mapped state and an arrayBuffer for writing.
const gpuWriteBuffer = device.createBuffer({
  mappedAtCreation: true,
  size: 4,
  usage: GPUBufferUsage.MAP_WRITE | GPUBufferUsage.COPY_SRC
});
const arrayBuffer = gpuWriteBuffer.getMappedRange();

// Write bytes to buffer.
new Uint8Array(arrayBuffer).set([0, 1, 2, 3]);

// Unmap buffer so that it can be used later for copy.
gpuWriteBuffer.unmap();

// Get a GPU buffer for reading in an unmapped state.
const gpuReadBuffer = device.createBuffer({
  size: 4,
  usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
});

由於 GPU 是獨立的輔助處理器,因此所有 GPU 指令都會以非同步方式執行。這就是為什麼系統會視需要建構並分批傳送的 GPU 指令清單。在 WebGPU 中,device.createCommandEncoder() 傳回的 GPU 指令編碼器是 JavaScript 物件,用於建構一批「緩衝處理」指令,並在某個時間點發送給 GPU。另一方面,GPUBuffer 的方法為「無緩衝處理」,代表這些方法在呼叫時,會以不可分割的形式執行。

取得 GPU 指令編碼器後,請呼叫 copyEncoder.copyBufferToBuffer() (如下所示),將這個指令新增至指令佇列,以供日後執行。最後,請呼叫 copyEncoder.finish() 並提交至 GPU 裝置指令佇列,完成編碼指令。佇列負責處理透過 device.queue.submit() 提交並使用 GPU 指令做為引數的提交內容。系統會以不可分割的形式,依序執行陣列中儲存的所有指令。

// Encode commands for copying buffer to buffer.
const copyEncoder = device.createCommandEncoder();
copyEncoder.copyBufferToBuffer(
  gpuWriteBuffer /* source buffer */,
  0 /* source offset */,
  gpuReadBuffer /* destination buffer */,
  0 /* destination offset */,
  4 /* size */
);

// Submit copy commands.
const copyCommands = copyEncoder.finish();
device.queue.submit([copyCommands]);

此時,系統會傳送 GPU 佇列指令,但不一定會執行。如要讀取第二個 GPU 緩衝區,請使用 GPUMapMode.READ 呼叫 gpuReadBuffer.mapAsync()。這個方法會傳回承諾,並在對應 GPU 緩衝區時加以解析。執行所有排入佇列的 GPU 指令後,即可使用 gpuReadBuffer.getMappedRange() 取得包含與第一個 GPU 緩衝區相同的值對應的範圍。

// Read buffer.
await gpuReadBuffer.mapAsync(GPUMapMode.READ);
const copyArrayBuffer = gpuReadBuffer.getMappedRange();
console.log(new Uint8Array(copyArrayBuffer));

你可以試用這個範例

簡而言之,以下是緩衝區記憶體作業的相關注意事項:

  • GPU 緩衝區必須取消對應,才能在裝置佇列提交時使用。
  • 對應之後,即可以 JavaScript 讀取及寫入 GPU 緩衝區。
  • 呼叫 mapAsync()createBuffer()mappedAtCreation 設為 true 時,系統會對應 GPU 緩衝區。

著色器程式設計

在 GPU 上執行的程式,只執行運算 (且不會繪製三角形) 的程式稱為「運算著色器」。且會由數百個 GPU 核心 (比 CPU 核心小) 平行執行,這些核心會共同運作來處理資料。其輸入和輸出是 WebGPU 中的緩衝區。

為說明在 WebGPU 中使用運算著色器的方式,我們將進行矩陣乘法,這是機器學習中常見的演算法,如下圖所示。

矩陣乘法圖表
矩陣乘法圖表

簡單來說,我們會進行下列工作:

  1. 建立三個 GPU 緩衝區 (兩個矩陣相乘,另一個代表結果矩陣)
  2. 說明運算著色器的輸入和輸出內容
  3. 編譯運算著色器程式碼
  4. 設定運算管線
  5. 將編碼指令批次提交至 GPU
  6. 讀取結果矩陣 GPU 緩衝區

建立 GPU 緩衝區

為求簡單起見,系統會使用浮點數清單呈現矩陣。第一個元素是列數,第二個元素是欄數,其餘元素則是矩陣的實際數字。

以簡易方式呈現 JavaScript 中的矩陣
JavaScript 中的矩陣簡單表示法,以及它在數學記號標記中的對等項目

三個 GPU 緩衝區是儲存空間緩衝區,因為我們需要在運算著色器中儲存及擷取資料。這說明瞭為何所有 GPU 緩衝區使用標記都包含 GPUBufferUsage.STORAGE。結果矩陣用量標記也具有 GPUBufferUsage.COPY_SRC,因為所有 GPU 佇列指令都執行完畢後,系統會將其複製到其他緩衝區進行讀取。

const adapter = await navigator.gpu.requestAdapter();
if (!adapter) { return; }
const device = await adapter.requestDevice();


// First Matrix

const firstMatrix = new Float32Array([
  2 /* rows */, 4 /* columns */,
  1, 2, 3, 4,
  5, 6, 7, 8
]);

const gpuBufferFirstMatrix = device.createBuffer({
  mappedAtCreation: true,
  size: firstMatrix.byteLength,
  usage: GPUBufferUsage.STORAGE,
});
const arrayBufferFirstMatrix = gpuBufferFirstMatrix.getMappedRange();
new Float32Array(arrayBufferFirstMatrix).set(firstMatrix);
gpuBufferFirstMatrix.unmap();


// Second Matrix

const secondMatrix = new Float32Array([
  4 /* rows */, 2 /* columns */,
  1, 2,
  3, 4,
  5, 6,
  7, 8
]);

const gpuBufferSecondMatrix = device.createBuffer({
  mappedAtCreation: true,
  size: secondMatrix.byteLength,
  usage: GPUBufferUsage.STORAGE,
});
const arrayBufferSecondMatrix = gpuBufferSecondMatrix.getMappedRange();
new Float32Array(arrayBufferSecondMatrix).set(secondMatrix);
gpuBufferSecondMatrix.unmap();


// Result Matrix

const resultMatrixBufferSize = Float32Array.BYTES_PER_ELEMENT * (2 + firstMatrix[0] * secondMatrix[1]);
const resultMatrixBuffer = device.createBuffer({
  size: resultMatrixBufferSize,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
});

繫結群組版面配置和繫結群組

繫結群組版面配置和繫結群組的概念是 WebGPU 專用的概念。繫結群組版面配置會定義著色器預期的輸入/輸出介面,繫結群組則代表著色器的實際輸入/輸出資料。

在下方範例中,繫結群組版面配置預期在編號項目繫結 01 之間需要兩個唯讀儲存空間緩衝區,並為運算著色器使用 2 的儲存空間緩衝區。另一方面,針對此繫結群組版面配置定義的繫結群組會將 GPU 緩衝區與以下項目建立關聯:gpuBufferFirstMatrix0gpuBufferSecondMatrix 與繫結 1 建立關聯,並將 resultMatrixBuffer 與繫結 2 建立關聯。

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

const bindGroup = device.createBindGroup({
  layout: bindGroupLayout,
  entries: [
    {
      binding: 0,
      resource: {
        buffer: gpuBufferFirstMatrix
      }
    },
    {
      binding: 1,
      resource: {
        buffer: gpuBufferSecondMatrix
      }
    },
    {
      binding: 2,
      resource: {
        buffer: resultMatrixBuffer
      }
    }
  ]
});

運算著色器程式碼

用來擴充矩陣的運算著色器程式碼是以 WGSL 編寫而成,WebGPU 著色器語言可大幅轉譯為 SPIR-V。如果不詳細說明,應該可以在使用 var<storage> 識別的三個儲存體緩衝區下方。程式會使用 firstMatrixsecondMatrix 做為輸入內容,並使用 resultMatrix 做為輸出內容。

請注意,每個儲存空間緩衝區都有 binding 裝飾,而該裝飾會對應於上方宣告的繫結群組版面配置和繫結群組中定義的索引。

const shaderModule = device.createShaderModule({
  code: `
    struct Matrix {
      size : vec2f,
      numbers: array<f32>,
    }

    @group(0) @binding(0) var<storage, read> firstMatrix : Matrix;
    @group(0) @binding(1) var<storage, read> secondMatrix : Matrix;
    @group(0) @binding(2) var<storage, read_write> resultMatrix : Matrix;

    @compute @workgroup_size(8, 8)
    fn main(@builtin(global_invocation_id) global_id : vec3u) {
      // Guard against out-of-bounds work group sizes
      if (global_id.x >= u32(firstMatrix.size.x) || global_id.y >= u32(secondMatrix.size.y)) {
        return;
      }

      resultMatrix.size = vec2(firstMatrix.size.x, secondMatrix.size.y);

      let resultCell = vec2(global_id.x, global_id.y);
      var result = 0.0;
      for (var i = 0u; i < u32(firstMatrix.size.y); i = i + 1u) {
        let a = i + resultCell.x * u32(firstMatrix.size.y);
        let b = resultCell.y + i * u32(secondMatrix.size.y);
        result = result + firstMatrix.numbers[a] * secondMatrix.numbers[b];
      }

      let index = resultCell.y + resultCell.x * u32(secondMatrix.size.y);
      resultMatrix.numbers[index] = result;
    }
  `
});

管道設定

運算管線是實際描述將執行運算作業的物件。呼叫 device.createComputePipeline() 即可建立容器。這個方法會採用兩個引數:先前建立的繫結群組版面配置,以及定義運算著色器進入點 (main WGSL 函式) 的運算階段,以及使用 device.createShaderModule() 建立的實際運算著色器模組。

const computePipeline = device.createComputePipeline({
  layout: device.createPipelineLayout({
    bindGroupLayouts: [bindGroupLayout]
  }),
  compute: {
    module: shaderModule,
    entryPoint: "main"
  }
});

提交指令

使用三個 GPU 緩衝區和具有繫結群組版面配置的運算管道將繫結群組執行個體化後,就可以開始使用這些項目了。

接著使用 commandEncoder.beginComputePass() 來啟動可程式化運算傳遞編碼器。我們會用它來對執行矩陣乘法的 GPU 指令進行編碼使用 passEncoder.setPipeline(computePipeline) 設定管道,並使用 passEncoder.setBindGroup(0, bindGroup) 將其繫結群組設為索引 0。索引 0 對應 WGSL 程式碼中的 group(0) 裝飾。

接著來談談這個運算著色器在 GPU 上的執行方式。我們的目標是針對結果矩陣的每個儲存格,逐步執行此程式。舉例來說,若是大小為 16 x 32 的結果矩陣,要為執行指令編碼,我們會在 @workgroup_size(8, 8) 上呼叫 passEncoder.dispatchWorkgroups(2, 4)passEncoder.dispatchWorkgroups(16 / 8, 32 / 8)。第一個引數「x」是第一個維度,第二個維度「y」是第二個維度,最新的「z」是第三個維度,因為此處用不到。在 GPU 運算世界中,將指令編碼以在一組資料上執行核心函式,稱為「分派」。

每個結果矩陣儲存格平行執行
每個結果矩陣儲存格平行執行

在 WGSL 程式碼中,運算著色器的工作群組格線大小為 (8, 8)。因此,「x」和「y」分別代表第一個矩陣的列數,第二個矩陣的欄數將除以 8。透過這種方式,我們現在可以使用 passEncoder.dispatchWorkgroups(firstMatrix[0] / 8, secondMatrix[1] / 8) 分派運算呼叫。要執行的工作群組格線數量是 dispatchWorkgroups() 引數。

如上圖所示,每個著色器都可以存取不重複的 builtin(global_invocation_id) 物件,用來判斷要計算哪個結果矩陣儲存格。

const commandEncoder = device.createCommandEncoder();

const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(computePipeline);
passEncoder.setBindGroup(0, bindGroup);
const workgroupCountX = Math.ceil(firstMatrix[0] / 8);
const workgroupCountY = Math.ceil(secondMatrix[1] / 8);
passEncoder.dispatchWorkgroups(workgroupCountX, workgroupCountY);
passEncoder.end();

如要結束運算傳遞編碼器,請呼叫 passEncoder.end()。接著,建立 GPU 緩衝區做為目的地,使用 copyBufferToBuffer 複製結果矩陣緩衝區。最後,使用 copyEncoder.finish() 完成編碼指令,並使用 GPU 指令呼叫 device.queue.submit(),將這些指令提交至 GPU 裝置佇列。

// Get a GPU buffer for reading in an unmapped state.
const gpuReadBuffer = device.createBuffer({
  size: resultMatrixBufferSize,
  usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
});

// Encode commands for copying buffer to buffer.
commandEncoder.copyBufferToBuffer(
  resultMatrixBuffer /* source buffer */,
  0 /* source offset */,
  gpuReadBuffer /* destination buffer */,
  0 /* destination offset */,
  resultMatrixBufferSize /* size */
);

// Submit GPU commands.
const gpuCommands = commandEncoder.finish();
device.queue.submit([gpuCommands]);

讀取結果矩陣

讀取結果矩陣就像使用 GPUMapMode.READ 呼叫 gpuReadBuffer.mapAsync(),並等待傳回承諾解決,表示 GPU 緩衝區已對應。此時,您可以使用 gpuReadBuffer.getMappedRange() 取得對應的範圍。

矩陣乘法結果
矩陣乘法結果

在我們程式碼中,登入開發人員工具 JavaScript 控制台的結果為「2, 2, 50, 60, 114, 140」。

// Read buffer.
await gpuReadBuffer.mapAsync(GPUMapMode.READ);
const arrayBuffer = gpuReadBuffer.getMappedRange();
console.log(new Float32Array(arrayBuffer));

恭喜!您成功了。您可以暢玩本範例

最後一個技巧

其中一種讓程式碼更容易閱讀,其中一種方法是使用運算管線的便利 getBindGroupLayout 方法,從著色器模組推斷繫結群組版面配置。這個技巧讓您不必建立自訂繫結群組版面配置,並在運算管道中指定管道版面配置,如下所示。

這裡會顯示先前範例的 getBindGroupLayout 插圖。

 const computePipeline = device.createComputePipeline({
-  layout: device.createPipelineLayout({
-    bindGroupLayouts: [bindGroupLayout]
-  }),
   compute: {
-// Bind group layout and bind group
- const bindGroupLayout = device.createBindGroupLayout({
-   entries: [
-     {
-       binding: 0,
-       visibility: GPUShaderStage.COMPUTE,
-       buffer: {
-         type: "read-only-storage"
-       }
-     },
-     {
-       binding: 1,
-       visibility: GPUShaderStage.COMPUTE,
-       buffer: {
-         type: "read-only-storage"
-       }
-     },
-     {
-       binding: 2,
-       visibility: GPUShaderStage.COMPUTE,
-       buffer: {
-         type: "storage"
-       }
-     }
-   ]
- });
+// Bind group
  const bindGroup = device.createBindGroup({
-  layout: bindGroupLayout,
+  layout: computePipeline.getBindGroupLayout(0 /* index */),
   entries: [

效能發現項目

那麼在 GPU 上執行矩陣乘法與在 CPU 上執行時,有何差異?為了找出答案,我編寫了剛剛描述 CPU 的程式。如下圖所示,當矩陣大小大於 256 x 256 時,使用 GPU 的完整能力似乎是顯而易見的選擇。

GPU 與 CPU 基準
GPU 與 CPU 基準測試

本文只是我開始探索 WebGPU 歷程的第一步。我們很快就會推出更多文章,以便深入探討 GPU 運算,以及 WebGPU 中的轉譯 (畫布、紋理、取樣器) 運作方式。