Comienza a usar GPU Compute en la Web

En esta entrada, se explora la API experimental de WebGPU mediante ejemplos y se te ayuda a comenzar a realizar cálculos de datos paralelos con la GPU.

François Beaufort

Información general

Como ya sabrás, la unidad de procesamiento gráfico (GPU) es un subsistema electrónico dentro de una computadora que se especializó originalmente para el procesamiento de gráficos. Sin embargo, en los últimos 10 años, evolucionó hacia una arquitectura más flexible que les permitía a los desarrolladores implementar muchos tipos de algoritmos, no solo renderizar gráficos 3D, y aprovechar la arquitectura única de la GPU. Estas funciones se denominan procesamiento con GPU, y el uso de una GPU como coprocesador para el procesamiento científico de uso general se denomina programación de GPU de uso general (GPGPU).

El procesamiento con GPU contribuyó de manera significativa al auge reciente del aprendizaje automático, ya que las redes neuronales de convolución y otros modelos pueden aprovechar la arquitectura para ejecutarse de manera más eficiente en las GPU. Dado que la plataforma web actual carece de capacidades de procesamiento con GPU, el grupo de la comunidad "GPU para la Web" de W3C está diseñando una API para exponer las APIs de GPU modernas que están disponibles en la mayoría de los dispositivos actuales. Esta API se llama WebGPU.

WebGPU es una API de bajo nivel, como WebGL. Es muy poderosa y muy detallada, como verás. Pero está bien. Lo que buscamos es el rendimiento.

En este artículo, me enfocaré en la parte de procesamiento con GPU de WebGPU y, para ser honesto, solo estaré al principio para que puedas comenzar a jugar por tu cuenta. Analizaremos en más detalle la renderización de WebGPU (lienzo, textura, etc.) en los próximos artículos.

Nota: Por ahora, WebGPU está disponible en Chrome Canary detrás de una marca experimental. Puedes habilitarla en chrome://flags/#enable-unsafe-webgpu. La API cambia constantemente y actualmente no es segura. Como la zona de pruebas de GPU aún no se implementó para la API de WebGPU, es posible leer datos de GPU para otros procesos. No navegues en la Web con esa opción habilitada.

Accede a la GPU

Acceder a la GPU es fácil en WebGPU. La llamada a navigator.gpu.requestAdapter() muestra una promesa de JavaScript que se resolverá de forma asíncrona con un adaptador de GPU. Piensa en este adaptador como la tarjeta gráfica. Puede estar integrada (en el mismo chip que la CPU) o discreta (por lo general, una tarjeta PCIe que tiene mejor rendimiento, pero usa más energía).

Una vez que tengas el adaptador de GPU, llama a adapter.requestDevice() para obtener una promesa que se resolverá con un dispositivo GPU que usarás para realizar algunos cálculos de GPU.

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

Ambas funciones toman opciones que te permiten especificar el tipo de adaptador (preferencia de alimentación) y el dispositivo (extensiones, límites) que deseas. Para simplificar, usaremos las opciones predeterminadas en este artículo.

Escribir en la memoria del búfer

Veamos cómo usar JavaScript para escribir datos en la memoria de la GPU. Este proceso no es sencillo debido al modelo de zona de pruebas que se usa en los navegadores web modernos.

En el siguiente ejemplo, se muestra cómo escribir cuatro bytes para almacenar en búfer la memoria accesible desde la GPU. Llama a device.createBuffer(), que toma el tamaño del búfer y su uso. Aunque la marca de uso GPUBufferUsage.MAP_WRITE no es necesaria para esta llamada específica, seamos explícitos que queremos escribir en este búfer. Como resultado, se obtiene un objeto de búfer de GPU asignado durante la creación gracias a que mappedAtCreation se establece en verdadero. Luego, se puede recuperar el búfer de datos binarios sin procesar asociado llamando al método del búfer de GPU getMappedRange().

La escritura de bytes resulta familiar si ya jugaste con ArrayBuffer. Usa un TypedArray y copia los valores en él.

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

En este punto, se asigna el búfer de la GPU, lo que significa que es propiedad de la CPU, y se puede acceder a él mediante lectura y escritura desde JavaScript. Para que la GPU pueda acceder a él, no debe estar asignado, lo que es tan simple como llamar a gpuBuffer.unmap().

El concepto de asignado y no asignado es necesario para evitar condiciones de carrera en las que la GPU y la CPU acceden a la memoria al mismo tiempo.

Lectura de memoria del búfer

Ahora, veamos cómo copiar un búfer de GPU a otro y volver a leerlo.

Como se escribe en el primer búfer de GPU y queremos copiarlo en un segundo búfer de GPU, se requiere una nueva marca de uso GPUBufferUsage.COPY_SRC. El segundo búfer de GPU se crea en un estado sin asignar esta vez con device.createBuffer(). Su marca de uso es GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ, ya que se usará como el destino del primer búfer de GPU y se leerá en JavaScript una vez que se ejecuten los comandos de copia de la GPU.

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

Debido a que la GPU es un coprocesador independiente, todos los comandos de la GPU se ejecutan de forma asíncrona. Es por eso que hay una lista de comandos de GPU que se crean y envían por lotes cuando es necesario. En WebGPU, el codificador de comandos de GPU que muestra device.createCommandEncoder() es el objeto JavaScript que compila un lote de comandos "almacenados en búfer" que se enviarán a la GPU en algún momento. Por otro lado, los métodos de GPUBuffer no tienen búfer, lo que significa que se ejecutan de forma atómica en el momento en que se los llama.

Una vez que tengas el codificador de comandos de la GPU, llama a copyEncoder.copyBufferToBuffer() como se muestra a continuación para agregar este comando a la cola de comandos y ejecutarlo más tarde. Por último, llama a copyEncoder.finish() para finalizar los comandos de codificación y envíalos a la cola de comandos del dispositivo GPU. La cola es responsable de controlar las entregas realizadas a través de device.queue.submit() con los comandos de la GPU como argumentos. Esto ejecutará de forma atómica todos los comandos almacenados en el array en orden.

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

En este punto, se enviaron los comandos de la cola de GPU, pero no necesariamente se ejecutaron. Para leer el segundo búfer de GPU, llama a gpuReadBuffer.mapAsync() con GPUMapMode.READ. Muestra una promesa que se resolverá cuando se asigne el búfer de la GPU. Luego, obtén el rango asignado con gpuReadBuffer.getMappedRange(), que contiene los mismos valores que el primer búfer de GPU una vez que se hayan ejecutado todos los comandos de GPU en cola.

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

Puedes probar esta muestra.

En resumen, esto es lo que debes recordar sobre las operaciones de la memoria del búfer:

• Los búferes de GPU no deben estar asignados para usarse en el envío de cola de dispositivos.
• Cuando se asignan, los búferes de GPU se pueden leer y escribir en JavaScript.
• Los búferes de GPU se asignan cuando se llama a mapAsync() y createBuffer() con mappedAtCreation establecido como verdadero.

Programación de sombreadores

Los programas que se ejecutan en la GPU y que solo realizan cálculos (y no dibujan triángulos) se denominan sombreadores de cómputos. Se ejecutan en paralelo cientos de núcleos de GPU (que son más pequeños que los núcleos de CPU) que operan juntos para procesar datos. Su entrada y salida son búferes en WebGPU.

Para ilustrar el uso de sombreadores de cómputos en WebGPU, jugaremos con la multiplicación de matrices, un algoritmo común en aprendizaje automático que se ilustra a continuación.

Diagrama de multiplicación de matrices

En resumen, haremos lo siguiente:

1. Crea tres búferes de GPU (dos para multiplicar las matrices y uno para la matriz de resultados).
2. Describir la entrada y la salida del sombreador de cómputos
3. Compila el código de sombreador de cómputos
4. Configura una canalización de cómputos
5. Envía los comandos codificados por lotes a la GPU
6. Lee el búfer de GPU de la matriz de resultados

Creación de búferes de GPU

Para simplificar, las matrices se representarán como una lista de números de punto flotante. El primer elemento es la cantidad de filas, el segundo es la cantidad de columnas y el resto son los números reales de la matriz.

Representación simple de una matriz en JavaScript y su equivalente en notación matemática

Los tres búferes de GPU son búferes de almacenamiento, ya que necesitamos almacenar y recuperar datos en el sombreador de cómputos. Esto explica por qué las marcas de uso del búfer de GPU incluyen GPUBufferUsage.STORAGE para todas ellas. La marca de uso de la matriz de resultados también tiene GPUBufferUsage.COPY_SRC porque se copiará en otro búfer para leer una vez que se hayan ejecutado todos los comandos de la cola de 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
});

Vincula el diseño del grupo y el grupo de vinculaciones

Los conceptos de diseño de grupos de vinculaciones y grupos de vinculaciones son específicos de WebGPU. Un diseño de grupo de vinculaciones define la interfaz de entrada y salida que espera un sombreador, mientras que un grupo de vinculación representa los datos de entrada y salida reales para un sombreador.

En el siguiente ejemplo, el diseño del grupo de vinculaciones espera dos búferes de almacenamiento de solo lectura en vinculaciones de entradas numeradas 0, 1 y un búfer de almacenamiento en 2 para el sombreador de cómputos. Por otro lado, el grupo de vinculaciones, definido para este diseño de grupo de vinculaciones, asocia búferes de GPU a las entradas: gpuBufferFirstMatrix a la 0 de vinculación, gpuBufferSecondMatrix a la vinculación 1 y resultMatrixBuffer a la vinculación 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
      }
    }
  ]
});

Código del sombreador de cómputos

El código de sombreador de cómputos para multiplicar matrices se escribe en WGSL, el lenguaje de sombreador de WebGPU, que se puede traducir trivialmente a SPIR-V. Sin entrar en detalle, deberías encontrar debajo de los tres búferes de almacenamiento identificados con var<storage>. El programa usará firstMatrix y secondMatrix como entradas, y resultMatrix como su salida.

Ten en cuenta que cada búfer de almacenamiento tiene una decoración binding usada que corresponde al mismo índice definido en los diseños de grupos de vinculaciones y los grupos de vinculaciones declarados anteriormente.

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

Configuración de la canalización

La canalización de cómputos es el objeto que realmente describe la operación de cómputos que vamos a realizar. Llama a device.createComputePipeline() para crearlo. Toma dos argumentos: el diseño del grupo de vinculaciones que creamos antes y una etapa de cómputos que define el punto de entrada de nuestro sombreador de cómputos (la función WGSL de main) y el módulo de sombreador de cómputos real creado con device.createShaderModule().

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

Envío de comandos

Después de crear una instancia de un grupo de vinculaciones con nuestros tres búferes de GPU y una canalización de cómputos con un diseño de grupo de vinculaciones, es hora de usarlos.

Comencemos un codificador de pase de cómputos programable con commandEncoder.beginComputePass(). Lo usaremos para codificar los comandos de la GPU que harán la multiplicación de matrices Configura su canalización con passEncoder.setPipeline(computePipeline) y su grupo de vinculaciones en el índice 0 con passEncoder.setBindGroup(0, bindGroup). El índice 0 corresponde a la decoración group(0) en el código WGSL.

Ahora, hablemos sobre cómo se ejecutará este sombreador de cómputos en la GPU. Nuestro objetivo es ejecutar este programa en paralelo para cada celda de la matriz de resultados, paso a paso. En el caso de una matriz de resultados de 16 por 32, por ejemplo, para codificar el comando de ejecución, en un @workgroup_size(8, 8), llamaremos a passEncoder.dispatchWorkgroups(2, 4) o passEncoder.dispatchWorkgroups(16 / 8, 32 / 8). El primer argumento "x" es la primera dimensión, el segundo, "y" es la segunda, y el último argumento "z" es la tercera dimensión cuya configuración predeterminada es 1, ya que no la necesitamos en este caso. En el mundo de procesamiento de GPU, codificar un comando para ejecutar una función de kernel en un conjunto de datos se denomina despacho.

Ejecución en paralelo para cada celda de la matriz de resultados

El tamaño de la cuadrícula del grupo de trabajo para nuestro sombreador de cómputos es (8, 8) en nuestro código WGSL. Debido a eso, “x” e “y”, que son respectivamente el número de filas de la primera matriz y el número de columnas de la segunda matriz, se dividirán por 8. Con eso, ahora podemos enviar una llamada de procesamiento con passEncoder.dispatchWorkgroups(firstMatrix[0] / 8, secondMatrix[1] / 8). La cantidad de cuadrículas de grupos de trabajo que se ejecutarán son los argumentos dispatchWorkgroups().

Como se ve en el dibujo anterior, cada sombreador tendrá acceso a un objeto builtin(global_invocation_id) único que se usará para saber qué celda de la matriz de resultados se debe calcular.

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 finalizar el codificador de pases de cómputos, llama a passEncoder.end(). Luego, crea un búfer de GPU para usarlo como destino y copiar el búfer de matriz de resultados con copyBufferToBuffer. Por último, finaliza los comandos de codificación con copyEncoder.finish() y envíalos a la cola del dispositivo de la GPU llamando a device.queue.submit() con los comandos de la 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]);

Leer la matriz de resultados

Para leer la matriz de resultados, es tan fácil como llamar a gpuReadBuffer.mapAsync() con GPUMapMode.READ y esperar a que se resuelva la promesa que se muestra, lo que indica que el búfer de GPU ahora está asignado. En este punto, es posible obtener el rango asignado con gpuReadBuffer.getMappedRange().

Resultado de la multiplicación de matrices

En nuestro código, el resultado registrado en la Consola de JavaScript de Herramientas para desarrolladores es “2, 2, 50, 60, 114, 140”.

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

¡Felicitaciones! Lo logró Puedes jugar con la muestra.

Un último truco

Una forma de facilitar la lectura de tu código es usar el práctico método getBindGroupLayout de la canalización de cómputos para inferir el diseño del grupo de vinculaciones desde el módulo del sombreador. Este truco quita la necesidad de crear un diseño de grupo de vinculaciones personalizado y de especificar un diseño de canalización en tu canalización de cómputos, como puedes ver a continuación.

Una ilustración de getBindGroupLayout para la muestra anterior está disponible.

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

Hallazgos sobre el rendimiento

Entonces, ¿cómo se compara la ejecución de la multiplicación de matrices en una GPU con la de ejecutarla en una CPU? Para descubrirlo, escribí el programa que acabamos de describir para una CPU. Y, como puedes ver en el siguiente gráfico, usar toda la potencia de la GPU parece una opción obvia cuando el tamaño de las matrices es superior a 256 por 256.

Comparativas de GPU frente a CPU

Este artículo fue solo el comienzo de mi recorrido por explorar WebGPU. Pronto habrá más artículos en los que se brindarán más detalles sobre el procesamiento con GPU y sobre el funcionamiento de la renderización (lienzo, textura, muestra) en WebGPU.