開始在網路上使用 GPU Compute

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

François Beaufort
François Beaufort
GitHub

背景

您可能已經知道，圖形處理器 (GPU) 是電腦中的電子子系統，原本專門用於處理圖形。不過，在過去 10 年間，它已演變成更具彈性的架構，可讓開發人員實作多種算法，而非僅是算繪 3D 圖形，同時還能充分利用 GPU 的獨特架構。這些功能稱為 GPU 運算，而將 GPU 用於通用科學運算的副處理器，則稱為通用 GPU (GPGPU) 程式設計。

近期機器學習蓬勃發展，GPU 運算功不可沒，因為卷積神經網路和其他模型可利用架構在 GPU 上更有效率地執行。由於目前的網路平台缺乏 GPU 運算功能，W3C 的「GPU for the Web」社群群組正在設計 API，以便公開目前大多數裝置可用的新型 GPU API。這個 API 稱為 WebGPU。

WebGPU 是低階 API，類似於 WebGL。您會發現，這項工具非常強大且冗長。但沒關係。我們要的是效能。

在本文中，我將著重於 WebGPU 的 GPU 運算部分，老實說，我只是略微介紹一下，讓您可以自行開始玩耍。我會在後續文章中深入探討 WebGPU 算繪 (畫布、紋理等)。

存取 GPU

在 WebGPU 中存取 GPU 非常簡單。呼叫 navigator.gpu.requestAdapter() 會傳回 JavaScript Promise，該 Promise 會以非同步方式解析 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 緩衝區已對應，也就是說，它由 CPU 擁有，且可透過 JavaScript 讀取/寫入。因此，您必須取消對應，讓 GPU 能夠存取，這就像呼叫 gpuBuffer.unmap() 一樣簡單。

需要使用已對應/未對應的概念，才能避免 GPU 和 CPU 同時存取記憶體的競爭狀態。

讀取緩衝區記憶體

接下來，我們來看看如何將 GPU 緩衝區複製到另一個 GPU 緩衝區，並讀取該緩衝區。

由於我們是在第一個 GPU 緩衝區中寫入資料，且想將資料複製到第二個 GPU 緩衝區，因此需要使用新的用法標記 GPUBufferUsage.COPY_SRC。這次使用 device.createBuffer() 時，第二個 GPU 緩衝區是在未對應的狀態下建立。其用途旗標為 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 上的「unbuffered」方法，也就是在呼叫時會自動執行。

取得 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()。它會傳回一個 promise，在 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 專屬。繫結群組版面配置會定義著色器預期的輸入/輸出介面，而繫結群組則代表著色器的實際輸入/輸出資料。

在以下範例中，繫結群組版面配置會預期在編號為 0 和 1 的項目繫結中，有兩個唯讀儲存空間緩衝區，以及在 2 中為運算著色器提供的儲存空間緩衝區。另一方面，為這個繫結群組版面配置定義的繫結群組，會將 GPU 緩衝區與項目建立關聯：gpuBufferFirstMatrix 與繫結 0、gpuBufferSecondMatrix 與繫結 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> 標示的三個儲存空間緩衝區。程式會使用 firstMatrix 和 secondMatrix 做為輸入內容，並將 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」是第三個維度，預設值為 1，因為我們在這裡不需要這個維度。在 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() 取得對應的範圍。

矩陣相乘結果
矩陣相乘結果

在我們的程式碼中，DevTools 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 中如何進行算繪 (畫布、紋理、取樣器)。