Introdução à computação de GPU na Web

Nesta postagem, exploramos a API WebGPU experimental por meio de exemplos e ajuda você a começar a realizar cálculos com paralelismo de dados usando a GPU.

François Beaufort
François Beaufort

Contexto

Como você já deve saber, a Unidade de Processamento Gráfico (GPU) é um subsistema eletrônico de um computador originalmente especializado para processar gráficos. No entanto, nos últimos 10 anos, ela evoluiu para uma arquitetura mais flexível, permitindo que os desenvolvedores implementem muitos tipos de algoritmos, não apenas gráficos 3D, e aproveitem a arquitetura exclusiva da GPU. Esses recursos são chamados de computação em GPU, e o uso de uma GPU como coprocessador para computação científica de uso geral é chamado de programação da GPU de uso geral (GPGPU, na sigla em inglês).

A computação com GPU contribuiu significativamente para o crescimento recente do machine learning, já que as redes neurais de convolução e outros modelos podem aproveitar a arquitetura para serem executados de forma mais eficiente nas GPUs. Como a atual plataforma da Web não tem recursos de computação de GPU, o grupo da comunidade "GPU para a Web" do W3C está criando uma API para expor as APIs de GPU modernas que estão disponíveis na maioria dos dispositivos atuais. Essa API é chamada de WebGPU.

A WebGPU é uma API de baixo nível, como a WebGL. Ele é muito poderoso e detalhado, como você verá. Mas não tem problema. Queremos apenas o desempenho.

Neste artigo, vou me concentrar na parte de computação de GPU da WebGPU e, para ser honesto, estou apenas falando do básico, para que você possa começar a jogar por conta própria. Vou me aprofundar na renderização da WebGPU (tela, textura etc.) nos próximos artigos.

Acessar a GPU

É fácil acessar a GPU na WebGPU. Chamar navigator.gpu.requestAdapter() retorna uma promessa de JavaScript que será resolvida de forma assíncrona com um adaptador de GPU. Pense nesse adaptador como uma placa de vídeo. Ela pode ser integrada (no mesmo chip que a CPU) ou discreta (geralmente uma placa PCIe com melhor desempenho, mas que consome mais energia).

Quando tiver o adaptador de GPU, chame adapter.requestDevice() para receber uma promessa que será resolvida com um dispositivo de GPU que você vai usar para fazer alguns cálculos de GPU.

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

As duas funções têm opções que permitem ser específico sobre o tipo de adaptador (preferência de energia) e o dispositivo (extensões, limites) que você quer. Para simplificar, vamos usar as opções padrão neste artigo.

Gravar memória de buffer

Vamos aprender a usar JavaScript para gravar dados na memória da GPU. Esse processo não é simples por causa do modelo de sandbox usado em navegadores da Web modernos.

O exemplo abaixo mostra como gravar quatro bytes na memória em buffer acessível pela GPU. Ele chama device.createBuffer(), que assume o tamanho e o uso do buffer. Mesmo que a flag de uso GPUBufferUsage.MAP_WRITE não seja necessária para essa chamada específica, vamos ser explícitos para gravar nesse buffer. Isso resulta em um objeto de buffer da GPU mapeado na criação graças a mappedAtCreation definido como verdadeiro. Em seguida, o buffer de dados binários brutos associado pode ser recuperado chamando o método de buffer da GPU getMappedRange().

A gravação de bytes é familiar se você já jogou com ArrayBuffer. Use um TypedArray e copie os valores nele.

// 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]);

Nesse ponto, o buffer da GPU é mapeado, o que significa que ele é de propriedade da CPU e pode ser acessado em leitura/gravação a partir do JavaScript. Para que a GPU possa acessá-lo, ela precisa ser desmapeada, o que é tão simples quanto chamar gpuBuffer.unmap().

O conceito de mapeado/não mapeado é necessário para evitar disputas em que a GPU e a CPU acessam a memória ao mesmo tempo.

Ler memória de buffer

Agora vamos aprender a copiar um buffer GPU em outro e fazer a leitura dele de volta.

Como estamos gravando no primeiro buffer de GPU e queremos copiá-lo para um segundo buffer de GPU, uma nova sinalização de uso GPUBufferUsage.COPY_SRC é necessária. O segundo buffer de GPU é criado em um estado não mapeado desta vez com device.createBuffer(). A flag de uso dela é GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ, porque será usada como destino do primeiro buffer da GPU e lida em JavaScript depois que os comandos de cópia da GPU forem executados.

// 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
});

Como a GPU é um coprocessador independente, todos os comandos dela são executados de forma assíncrona. É por isso que há uma lista de comandos da GPU criada e enviada em lotes quando necessário. Na WebGPU, o codificador de comandos da GPU retornado por device.createCommandEncoder() é o objeto JavaScript que cria um lote de comandos "armazenados em buffer" que serão enviados à GPU em algum momento. Por outro lado, os métodos em GPUBuffer não estão em buffer, o que significa que são executados atomicamente no momento em que são chamados.

Depois de ter o codificador de comandos da GPU, chame copyEncoder.copyBufferToBuffer(), como mostrado abaixo, para adicionar esse comando à fila de comandos para execução posterior. Por fim, conclua a codificação dos comandos chamando copyEncoder.finish() e envie-os para a fila de comandos do dispositivo da GPU. A fila é responsável por processar envios feitos via device.queue.submit() com os comandos da GPU como argumentos. Isso executará atomicamente todos os comandos armazenados na matriz em ordem.

// 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]);

Os comandos da fila da GPU já foram enviados, mas não necessariamente executados. Para ler o segundo buffer de GPU, chame gpuReadBuffer.mapAsync() com GPUMapMode.READ. Ela retorna uma promessa que será resolvida quando o buffer da GPU for mapeado. Em seguida, receba o intervalo mapeado com gpuReadBuffer.getMappedRange() que contém os mesmos valores do primeiro buffer de GPU depois que todos os comandos da GPU na fila forem executados.

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

Teste esta amostra.

Resumindo, veja o que você precisa lembrar sobre as operações de memória de buffer:

  • Os buffers de GPU precisam ser desmapeados para serem usados no envio da fila do dispositivo.
  • Quando mapeados, os buffers de GPU podem ser lidos e gravados em JavaScript.
  • Os buffers de GPU são mapeados quando mapAsync() e createBuffer() com mappedAtCreation definido como verdadeiro são chamados.

Programação de sombreador

Os programas em execução na GPU que só realizam cálculos (e não desenham triângulos) são chamados de sombreadores de computação. Eles são executados em paralelo por centenas de núcleos de GPU (menores que os núcleos de CPU) que operam em conjunto para processar dados. A entrada e a saída são buffers na WebGPU.

Para ilustrar o uso de sombreadores de computação na WebGPU, vamos jogar com a multiplicação de matrizes, um algoritmo comum em machine learning ilustrado abaixo.

Diagrama de multiplicação de matrizes
Diagrama de multiplicação de matrizes

Resumindo, aqui está o que vamos fazer:

  1. Crie três buffers de GPU: dois para as matrizes multiplicarem e um para a matriz de resultado.
  2. Descrever a entrada e a saída do sombreador de computação
  3. Compilar o código do sombreador de computação
  4. Configurar um pipeline de computação
  5. Enviar em lote os comandos codificados para a GPU
  6. Ler o buffer de GPU da matriz de resultado

Criação de buffers de GPU

Para simplificar, as matrizes serão representadas como uma lista de números de ponto flutuante. O primeiro elemento é o número de linhas, o segundo é o número de colunas e o restante são os números reais da matriz.

Representação simples de uma matriz em JavaScript e seu equivalente na notação matemática
Representação simples de uma matriz em JavaScript e o equivalente na notação matemática

Os três buffers de GPU são de armazenamento, porque precisamos armazenar e extrair dados no sombreador de computação. Isso explica por que as flags de uso do buffer da GPU incluem GPUBufferUsage.STORAGE para todas elas. A flag de uso da matriz resultante também tem GPUBufferUsage.COPY_SRC porque será copiada em outro buffer para leitura depois que todos os comandos da fila da GPU forem executados.

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
});

Vincular layout do grupo e vincular grupo

Os conceitos de layout e grupo de vinculação são específicos para a WebGPU. Um layout de grupo de vinculação define a interface de entrada/saída esperada por um sombreador, enquanto um grupo de vinculação representa os dados reais de entrada/saída de um sombreador.

No exemplo abaixo, o layout do grupo de vinculação espera dois buffers de armazenamento somente leitura nas vinculações de entrada numeradas 0, 1 e um buffer de armazenamento em 2 para o sombreador de computação. Por outro lado, o grupo de vinculação, definido para esse layout, associa buffers de GPU às entradas: gpuBufferFirstMatrix à vinculação 0, gpuBufferSecondMatrix à vinculação 1 e resultMatrixBuffer à vinculação 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
      }
    }
  ]
});

Calcular código do sombreador

O código do sombreador de computação para multiplicar matrizes é escrito em WGSL, a linguagem WebGPU Shader, que pode ser traduzido para SPIR-V. Sem entrar em detalhes, confira abaixo os três buffers de armazenamento identificados com var<storage>. O programa vai usar firstMatrix e secondMatrix como entradas e resultMatrix como saída.

Cada buffer de armazenamento tem uma decoração binding usada que corresponde ao mesmo índice definido nos layouts de grupos de vinculação e nos grupos de vinculação declarados acima.

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;
    }
  `
});

Configuração do pipeline

O pipeline de computação é o objeto que descreve a operação de computação que vamos realizar. Crie-o chamando device.createComputePipeline(). Ela usa dois argumentos: o layout do grupo de vinculação criado anteriormente e um estágio de computação que define o ponto de entrada do sombreador de computação (a função WGSL main) e o módulo real do sombreador de computação criado com device.createShaderModule().

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

Envio de comandos

Depois de instanciar um grupo de vinculação com nossos três buffers de GPU e um pipeline de computação com um layout de grupo de vinculação, é hora de usá-los.

Vamos iniciar um codificador de passagem de computação programável com commandEncoder.beginComputePass(). Ele será usado para codificar comandos da GPU que farão a multiplicação de matrizes. Defina o pipeline com passEncoder.setPipeline(computePipeline) e o grupo de vinculação no índice 0 com passEncoder.setBindGroup(0, bindGroup). O índice 0 corresponde à decoração group(0) no código WGSL.

Agora vamos conferir como o sombreador de computação será executado na GPU. Nosso objetivo é executar esse programa em paralelo para cada célula da matriz de resultado, passo a passo. Em uma matriz de resultados de tamanho 16 por 32, por exemplo, para codificar o comando de execução, em um @workgroup_size(8, 8), chamaríamos passEncoder.dispatchWorkgroups(2, 4) ou passEncoder.dispatchWorkgroups(16 / 8, 32 / 8). O primeiro argumento "x" é a primeira dimensão, o segundo "y" é a segunda dimensão e o último "z" é a terceira dimensão, que tem como padrão 1, porque não precisamos disso aqui. No mundo da computação da GPU, a codificação de um comando para executar uma função do kernel em um conjunto de dados é chamada de despacho.

Execução em paralelo para cada célula de matriz de resultado
Execução em paralelo para cada célula da matriz de resultados

O tamanho da grade do grupo de trabalho do sombreador de computação é (8, 8) no código WGSL. Por isso, "x" e "y", que são respectivamente o número de linhas da primeira matriz e o número de colunas da segunda matriz, será dividido por 8. Com isso, agora podemos enviar uma chamada de computação com passEncoder.dispatchWorkgroups(firstMatrix[0] / 8, secondMatrix[1] / 8). O número de grades do grupo de trabalho a serem executadas são os argumentos dispatchWorkgroups().

Como mostrado no desenho acima, cada sombreador terá acesso a um objeto builtin(global_invocation_id) exclusivo que será usado para saber qual célula da matriz de resultados será computada.

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();

Para encerrar o codificador de passagem de computação, chame passEncoder.end(). Em seguida, crie um buffer de GPU a ser usado como destino para copiar o buffer da matriz de resultado com copyBufferToBuffer. Por fim, conclua a codificação dos comandos com copyEncoder.finish() e envie-os para a fila do dispositivo da GPU chamando device.queue.submit() com os comandos da 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]);

Ler matriz de resultados

Ler a matriz de resultado é tão fácil quanto chamar gpuReadBuffer.mapAsync() com GPUMapMode.READ e aguardar a resolução da promessa de retorno, o que indica que o buffer da GPU agora está mapeado. Nesse ponto, é possível receber o intervalo mapeado com gpuReadBuffer.getMappedRange().

Resultado da multiplicação de matrizes
Resultado da multiplicação de matrizes

No nosso código, o resultado registrado no Console JavaScript do DevTools é "2, 2, 50, 60, 114, 140".

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

Parabéns! Você conseguiu. Teste o exemplo.

Um último truque

Uma maneira de facilitar a leitura do código é usar o prático método getBindGroupLayout do pipeline de computação para inferir o layout do grupo de vinculação do módulo do sombreador. Esse truque elimina a necessidade de criar um layout personalizado de grupo de vinculação e especificar um layout de pipeline no pipeline de computação, como mostrado abaixo.

Uma ilustração de getBindGroupLayout para o exemplo anterior está disponível.

 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: [

Descobertas de desempenho

Como a execução da multiplicação de matrizes em uma GPU se compara à execução em uma CPU? Para descobrir, criei o programa que acabamos de descrever para uma CPU. Como mostrado no gráfico abaixo, usar a capacidade total da GPU parece uma escolha óbvia quando o tamanho das matrizes é maior que 256 por 256.

Comparativo de GPU x CPU
Comparativo de mercado de GPU x CPU

Este artigo é apenas o começo da minha jornada explorando a WebGPU. Em breve, você vai receber mais artigos com mais detalhes sobre a computação com GPU e sobre como a renderização (tela, textura, amostragem) funciona na WebGPU.