【譯】在Web中使用GPU Compute

原文連接：Get started with GPU Compute on the Web

背景

正如咱們所知，圖形處理單元（GPU）是計算機中的一個電子子系統，最初專門用於處理圖形。然而，在過去10年中，它已經發展成爲一種更靈活的架構，利用GPU的獨特架構，開發人員能夠實現多種類型的算法，而不只僅是渲染3D圖形。這些功能稱爲GPU Compute，使用GPU做爲通用科學計算的協處理器稱爲general-purpose GPU（GPGPU）編程。

GPU Compute爲最近的機器學習熱潮作出了重大貢獻，例如卷積神經網絡和其餘模型能夠利用此功能在GPU上更高效地運行。因爲當前的Web平臺缺少GPU Compute功能，W3C的「GPU for the Web」社區組正在設計一個API來暴露出GPU API，以便在大多數設備上使用，此API稱爲WebGPU。

WebGPU是一種底層API，例如WebGL。正如咱們所看到的，它很是強大且很是詳盡。但不要緊，咱們關注的是性能。

在本文中，我將重點關注WebGPU的GPU Compute部分，說實話，我會講的淺顯易懂，這樣你們就能夠玩出花樣。我會在即將發表的文章中介紹並深刻探討WebGPU渲染（畫布，紋理等）。

WebGPU目前在macOS的Chrome 78中可經過experimental flag開啓使用。您能夠在chrome://flags/ #enable-unsafe-webgpu中啓用它。API會常常變更，目前不太安全。因爲目前尚未爲WebGPU API實現GPU沙盒，所以能夠讀取其餘進程的GPU數據！所以請勿在啓用WebGPU下模式下瀏覽網頁。

訪問GPU

在WebGPU中訪問GPU很容易。調用navigator.gpu.requestAdapter()會返回一個JavaScript Promise，通過解析異步返回一個GPU adapter（GPU適配器）。咱們能夠將此適配器視爲顯卡。它能夠是集成顯卡（與CPU在同一芯片上）或獨立顯卡（一般是性能更高但功耗更大的PCIe卡）。

一旦得到GPU適配器後，就能夠調用adapter.requestDevice()並返回一個將使用GPU設備進行解析的promise，使用它來進行一些GPU計算。

const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();
複製代碼

這兩個函數都容許傳入所須要的適配器類型（電源首選項）和設備（擴展，限制）的參數。爲簡單起見，咱們將使用本文中的默認選項。

寫buffer

讓咱們看看如何使用JavaScript將數據寫入GPU的內存。因爲現代Web瀏覽器中使用的沙盒模型，這個過程並不簡單。

下面的示例顯示瞭如何將四個字節寫入可從GPU訪問的buffer memory。它調用device.createBufferMappedAsync()來獲得緩衝區的大小並表示用來幹嗎。即便這個API調用不須要使用標誌GPUBufferUsage.MAP_WRITE，也會代表咱們要寫入此buffer。生成的promise將返回GPU buffer對象及其關聯的原始二進制數據array buffer。

若是你已經玩過ArrayBuffer，那麼對寫字節一定很熟悉;使用TypedArray並將值複製進去。

// Get a GPU buffer and an arrayBuffer for writing.
// Upon success the GPU buffer is put in the mapped state.
const [gpuBuffer, arrayBuffer] = await device.createBufferMappedAsync({
  size: 4,
  usage: GPUBufferUsage.MAP_WRITE
});

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

此時，GPU buffer被映射，這意味着它由CPU擁有，而且能夠從JavaScript讀取/寫入。所以GPU能夠訪問它，它必須是未映射的，這就像調用gpuBuffer.unmap()同樣簡單。

映射或未映射的概念就是來防止GPU和CPU同時訪問內存的競爭條件。

讀buffer

如今讓咱們看看如何將GPU緩衝區複製到另外一個GPU緩衝區並將其讀回。

因爲咱們在第一個GPU緩衝區中寫入而且咱們想將其複製到第二個GPU緩衝區，所以須要新的使用標誌GPUBufferUsage.COPY_SRC。使用同步device.createBuffer（）以未映射狀態建立第二個GPU緩衝區。它的用法標誌是GPUBufferUsage.COPY_DST |GPUBufferUsage.MAP_READ，由於它將用做第一個GPU緩衝區的目標，並在執行GPU複製命令後在JavaScript中讀取。

// Get a GPU buffer and an arrayBuffer for writing.
// Upon success the GPU buffer is returned in the mapped state.
const [gpuWriteBuffer, arrayBuffer] = await device.createBufferMappedAsync({
  size: 4,
  usage: GPUBufferUsage.MAP_WRITE | GPUBufferUsage.COPY_SRC
});

// 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.getQueue（）。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.getQueue().submit([copyCommands]);
複製代碼

此時，已發送GPU隊列命令，但不必定執行。要讀取第二個GPU緩衝區，請調用gpuReadBuffer.mapReadAsync（）。一旦執行了全部排隊的GPU命令，它將返回一個承諾，該承諾將使用包含與第一個GPU緩衝區相同的值的ArrayBuffer來解析。

// Read buffer.
const copyArrayBuffer = await gpuReadBuffer.mapReadAsync();
console.log(new Uint8Array(copyArrayBuffer));
複製代碼

能夠戳示例

簡而言之，這是關於緩衝存儲器操做須要記住的內容：

• 必須取消映射GPU緩衝區才能在設備隊列提交中使用。

• 映射後，可使用JavaScript讀取和寫入GPU緩衝區。

• 調用mapReadAsync（），mapWriteAsync（），createBufferMappedAsync（）和createBufferMapped（）時，將映射GPU緩衝區。

着色編程

在GPU上運行的僅執行計算（而且不繪製三角形）的程序稱爲計算着色器。它們由數百個GPU內核（小於CPU內核）並行執行，這些內核一塊兒運行以處理數據。它們的輸入和輸出是WebGPU中的緩衝區。

爲了說明在WebGPU中使用計算着色器，咱們將使用矩陣乘法，這是下面所示的機器學習中的經常使用算法。

簡而言之，咱們要作的是：

1. 建立三個GPU緩衝區（兩個用於矩陣乘法，一個用於結果矩陣）
2. 描述計算着色器的輸入和輸出
3. 編譯計算着色器代碼
4. 設置計算管道
5. 將編碼的命令批量提交給GPU
6. 讀取結果矩陣GPU緩衝區

GPU buffer建立

爲簡單起見，矩陣將表示爲浮點數列表。第一個元素是行數，第二個元素是列數，其他元素是矩陣的實際數。

三個GPU緩衝區是存儲緩衝區，由於咱們須要在計算着色器中存儲和檢索數據。這解釋了爲何GPU緩衝區使用標誌包含GPUBufferUsage.STORAGE。結果矩陣使用標誌也有GPUBufferUsage.COPY_SRC，由於一旦全部GPU隊列命令都被執行，它將被複制到另外一個緩衝區以便讀取。

const adapter = await navigator.gpu.requestAdapter();
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, arrayBufferFirstMatrix] = await device.createBufferMappedAsync({
  size: firstMatrix.byteLength,
  usage: GPUBufferUsage.STORAGE,
});
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, arrayBufferSecondMatrix] = await device.createBufferMappedAsync({
  size: secondMatrix.byteLength,
  usage: GPUBufferUsage.STORAGE,
});
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({
  bindings: [
    {
      binding: 0,
      visibility: GPUShaderStage.COMPUTE,
      type: "storage-buffer"
    },
    {
      binding: 1,
      visibility: GPUShaderStage.COMPUTE,
      type: "storage-buffer"
    },
    {
      binding: 2,
      visibility: GPUShaderStage.COMPUTE,
      type: "storage-buffer"
    }
  ]
});

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

計算着色器代碼

用於乘法矩陣的計算着色器代碼用GLSL編寫，GLSL是WebGL中使用的高級着色語言，其具備基於C編程語言的語法。在不詳細說明的狀況下，您應該在下面找到標有關鍵字緩衝區的三個存儲緩衝區。該程序將使用firstMatrix和secondMatrix做爲輸入，並使用resultMatrix做爲其輸出。

請注意，每一個存儲緩衝區都使用一個綁定限定符，該限定符對應於綁定組佈局中定義的相同索引和上面聲明的綁定組。

const computeShaderCode = `#version 450 layout(std430, set = 0, binding = 0) readonly buffer FirstMatrix { vec2 size; float numbers[]; } firstMatrix; layout(std430, set = 0, binding = 1) readonly buffer SecondMatrix { vec2 size; float numbers[]; } secondMatrix; layout(std430, set = 0, binding = 2) buffer ResultMatrix { vec2 size; float numbers[]; } resultMatrix; void main() { resultMatrix.size = vec2(firstMatrix.size.x, secondMatrix.size.y); ivec2 resultCell = ivec2(gl_GlobalInvocationID.x, gl_GlobalInvocationID.y); float result = 0.0; for (int i = 0; i < firstMatrix.size.y; i++) { int a = i + resultCell.x * int(firstMatrix.size.y); int b = resultCell.y + i * int(secondMatrix.size.y); result += firstMatrix.numbers[a] * secondMatrix.numbers[b]; } int index = resultCell.y + resultCell.x * int(secondMatrix.size.y); resultMatrix.numbers[index] = result; } `;
複製代碼

管道設置

Chrome中的WebGPU目前使用字節碼而不是原始GLSL代碼。這意味着咱們必須在運行計算着色器以前編譯computeShaderCode。幸運的是，@ webgpu / glslang包容許咱們以Chrome中的WebGPU接受的格式編譯computeShaderCode。此字節碼格式基於SPIR-V的安全子集。

請注意，「Web上的GPU」W3C社區組在撰寫WebGPU的着色語言時仍未決定。

import glslangModule from 'https://unpkg.com/@webgpu/glslang/web/glslang.js';
複製代碼

計算管道是實際描述咱們將要執行的計算操做的對象。經過調用device.createComputePipeline（）建立它。它有兩個參數：咱們以前建立的綁定組佈局，以及定義計算着色器（主GLSL函數）的入口點的計算階段和使用glslang.compileGLSL（）編譯的實際計算着色器模塊。

const glslang = await glslangModule();

const computePipeline = device.createComputePipeline({
  layout: device.createPipelineLayout({
    bindGroupLayouts: [bindGroupLayout]
  }),
  computeStage: {
    module: device.createShaderModule({
      code: glslang.compileGLSL(computeShaderCode, "compute")
    }),
    entryPoint: "main"
  }
});
複製代碼

命令提交

在使用咱們的三個GPU緩衝區和帶有綁定組佈局的計算管道實例化綁定組以後，是時候使用它們了。

讓咱們用commandEncoder.beginComputePass（）啓動一個可編程的計算傳遞編碼器。咱們將使用它來編碼將執行矩陣乘法的GPU命令。使用passEncoder.setPindline（computePipeline）設置其管道，使用passEncoder.setBindGroup（0，bindGroup）在索引0處設置其綁定組。索引0對應於GLSL代碼中的set = 0限定符。

如今，咱們來談談這個計算着色器將如何在GPU上運行。咱們的目標是逐步爲結果矩陣的每一個單元並行執行該程序。例如，對於大小爲2乘4的結果矩陣，咱們調用passEncoder.dispatch（2,4）來編碼執行命令。第一個參數「x」是第一個維度，第二個參數「y」是第二個維度，最新的一個「z」是默認爲1的第三個維度，由於咱們在這裏不須要它。在GPU計算世界中，對一組數據執行內核函數的命令編碼稱爲調度。

在咱們的代碼中，「x」和「y」將分別是第一矩陣的行數和第二矩陣的列數。有了它，咱們如今可使用passEncoder.dispatch（firstMatrix [0]，secondMatrix [1]）調度計算調用。

如上圖所示，每一個着色器均可以訪問惟一的gl_GlobalInvocationID對象，該對象將用於瞭解要計算的結果矩陣單元格。

const commandEncoder = device.createCommandEncoder();

const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(computePipeline);
passEncoder.setBindGroup(0, bindGroup);
passEncoder.dispatch(firstMatrix[0] /* x */, secondMatrix[1] /* y */);
passEncoder.endPass();
複製代碼

要結束計算傳遞編碼器，請調用passEncoder.endPass（）。而後，建立一個GPU緩衝區以用做目標，以使用copyBufferToBuffer複製結果矩陣緩衝區。最後，使用copyEncoder.finish（）完成編碼命令，並經過使用GPU命令調用device.getQueue（）。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.getQueue().submit([gpuCommands]);
複製代碼

讀取結果矩陣

讀取結果矩陣就像調用gpuReadBuffer.mapReadAsync（）並記錄生成的promise返回的ArrayBuffer同樣簡單。

在咱們的代碼中，DevTools JavaScript控制檯中記錄的結果是「2,2,50,60,114,140」。

// Read buffer.
const arrayBuffer = await gpuReadBuffer.mapReadAsync();
console.log(new Float32Array(arrayBuffer));
複製代碼

示例請戳這裏

性能調查

那麼GPU上運行矩陣乘法與在CPU上運行矩陣乘法相好比何呢？爲了找到答案，我編寫了剛剛爲CPU描述的程序。正如您在下圖中所看到的，當矩陣的大小大於256乘256時，使用GPU的所有功能彷佛是一個明顯的選擇。

這篇文章只是我探索WebGPU之旅的開始。很快就會有更多文章介紹GPU Compute中更深刻的潛力以及渲染（畫布，紋理，採樣器）在WebGPU中的工做方式。