Начните работу с вычислениями на графическом процессоре в Интернете

В этом посте на примерах рассматривается экспериментальный API WebGPU и помогает приступить к выполнению параллельных вычислений с использованием графического процессора.

Франсуа Бофор
François Beaufort
GitHub

Фон

Как вы, возможно, уже знаете, графический процессор (GPU) — это электронная подсистема внутри компьютера, которая изначально специализировалась на обработке графики. Однако за последние 10 лет он развился в сторону более гибкой архитектуры, позволяющей разработчикам реализовывать множество типов алгоритмов, а не только рендеринг 3D-графики, используя при этом преимущества уникальной архитектуры графического процессора. Эти возможности называются вычислениями на графическом процессоре, а использование графического процессора в качестве сопроцессора для научных вычислений общего назначения называется программированием на графическом процессоре общего назначения (GPGPU).

Вычисления на графическом процессоре внесли значительный вклад в недавний бум машинного обучения, поскольку сверточные нейронные сети и другие модели могут использовать преимущества этой архитектуры для более эффективной работы на графических процессорах. Поскольку в текущей веб-платформе отсутствуют возможности вычислений на графическом процессоре, группа сообщества W3C «GPU for the Web» разрабатывает API для предоставления современных API-интерфейсов графического процессора, доступных на большинстве современных устройств. Этот API называется WebGPU .

WebGPU — это низкоуровневый API, такой как WebGL. Как вы увидите, это очень мощный и многословный инструмент. Но это нормально. Нам нужна производительность.

В этой статье я сосредоточусь на вычислительной части WebGPU, связанной с вычислениями на графическом процессоре, и, честно говоря, я лишь прикоснулся к поверхности, чтобы вы могли начать играть самостоятельно. Я буду копаться глубже и освещать рендеринг WebGPU (холст, текстуру и т. д.) в следующих статьях.

Доступ к графическому процессору

Доступ к графическому процессору в WebGPU прост. Вызов navigator.gpu.requestAdapter() возвращает обещание JavaScript, которое будет асинхронно разрешаться с помощью адаптера графического процессора. Думайте об этом адаптере как о видеокарте. Он может быть либо интегрированным (на том же чипе, что и ЦП), либо дискретным (обычно это карта PCIe, которая более производительна, но потребляет больше энергии).

Когда у вас есть адаптер графического процессора, вызовите adapter.requestDevice() чтобы получить обещание, которое будет разрешено с помощью устройства графического процессора, которое вы будете использовать для выполнения некоторых вычислений на графическом процессоре.

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

Обе функции принимают параметры, которые позволяют вам указать тип адаптера (настройки мощности) и устройства (расширения, ограничения), которые вам нужны. Для простоты в этой статье мы будем использовать параметры по умолчанию.

Запись буферной памяти

Давайте посмотрим, как использовать JavaScript для записи данных в память графического процессора. Этот процесс непрост из-за модели песочницы, используемой в современных веб-браузерах.

В примере ниже показано, как записать четыре байта в буферную память, доступную из графического процессора. Он вызывает device.createBuffer() , который принимает размер буфера и его использование. Несмотря на то, что флаг использования GPUBufferUsage.MAP_WRITE не требуется для этого конкретного вызова, давайте четко укажем, что мы хотим производить запись в этот буфер. В результате объект буфера графического процессора отображается при создании благодаря значению mappedAtCreation установленному в true. Затем связанный буфер необработанных двоичных данных можно получить, вызвав метод 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]);

На этом этапе буфер графического процессора отображается, то есть он принадлежит процессору и доступен для чтения/записи из JavaScript. Чтобы графический процессор мог получить к нему доступ, его необходимо отменить, что так же просто, как вызов gpuBuffer.unmap() .

Концепция сопоставления/несопоставления необходима для предотвращения состояний гонки, когда графический процессор и процессор одновременно обращаются к памяти.

Чтение буферной памяти

Теперь давайте посмотрим, как скопировать буфер графического процессора в другой буфер графического процессора и прочитать его обратно.

Поскольку мы записываем данные в первый буфер графического процессора и хотим скопировать их во второй буфер графического процессора, требуется новый флаг использования GPUBufferUsage.COPY_SRC . На этот раз второй буфер графического процессора создается в несопоставленном состоянии с помощью device.createBuffer() . Его флаг использования — GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ , поскольку он будет использоваться в качестве места назначения первого буфера графического процессора и считываться в 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
});

Поскольку графический процессор является независимым сопроцессором, все команды графического процессора выполняются асинхронно. Вот почему существует список команд графического процессора, который создается и при необходимости отправляется пакетно. В WebGPU кодировщик команд графического процессора, возвращаемый функцией device.createCommandEncoder() представляет собой объект JavaScript, который создает пакет «буферизованных» команд, которые в какой-то момент будут отправлены в графический процессор. С другой стороны, методы GPUBuffer «небуферизованы», то есть они выполняются атомарно в момент вызова.

Получив кодировщик команд графического процессора, вызовите copyEncoder.copyBufferToBuffer() как показано ниже, чтобы добавить эту команду в очередь команд для последующего выполнения. Наконец, завершите команды кодирования, вызвав метод copyEncoder.finish() , и отправьте их в очередь команд устройства графического процессора. Очередь отвечает за обработку отправок, выполненных через device.queue.submit() с командами графического процессора в качестве аргументов. Это приведет к атомарному выполнению всех команд, хранящихся в массиве, по порядку.

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

На этом этапе команды очереди графического процессора были отправлены, но не обязательно выполнены. Чтобы прочитать второй буфер графического процессора, вызовите gpuReadBuffer.mapAsync() с GPUMapMode.READ . Он возвращает обещание, которое будет выполнено при отображении буфера графического процессора. Затем получите сопоставленный диапазон с помощью gpuReadBuffer.getMappedRange() , который содержит те же значения, что и первый буфер графического процессора, после выполнения всех команд графического процессора в очереди.

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

Вы можете попробовать этот образец .

Короче говоря, вот что вам нужно помнить относительно операций с буферной памятью:

  • Буферы графического процессора должны быть отключены для использования при отправке очереди устройств.
  • При сопоставлении буферы графического процессора можно читать и записывать с помощью JavaScript.
  • Буферы графического процессора сопоставляются при вызове mapAsync() и createBuffer() с mappedAtCreation установленным в true.

Программирование шейдеров

Программы, работающие на графическом процессоре и выполняющие только вычисления (но не рисующие треугольники), называются вычислительными шейдерами. Они выполняются параллельно сотнями ядер графического процессора (которые меньше ядер ЦП), которые работают вместе для обработки данных. Их ввод и вывод — это буферы в WebGPU.

Чтобы проиллюстрировать использование вычислительных шейдеров в WebGPU, мы поиграем с умножением матриц — распространенным алгоритмом машинного обучения, показанным ниже.

Диаграмма умножения матрицы
Диаграмма умножения матрицы

Вкратце, вот что мы собираемся сделать:

  1. Создайте три буфера графического процессора (два для матриц для умножения и один для матрицы результатов).
  2. Описать ввод и вывод вычислительного шейдера.
  3. Скомпилируйте код вычислительного шейдера
  4. Настройка вычислительного конвейера
  5. Отправьте в пакетном режиме закодированные команды на графический процессор.
  6. Чтение буфера графического процессора матрицы результатов

Создание буферов графического процессора

Для простоты матрицы будут представлены в виде списка чисел с плавающей запятой. Первый элемент — это количество строк, второй элемент — количество столбцов, а остальное — фактические числа матрицы.

Простое представление матрицы в JavaScript и его эквивалент в математической записи.
Простое представление матрицы в JavaScript и его эквивалент в математической записи.

Три буфера графического процессора являются буферами хранения, поскольку нам необходимо хранить и извлекать данные в вычислительном шейдере. Это объясняет, почему флаги использования буфера графического процессора для всех из них включают GPUBufferUsage.STORAGE . Флаг использования матрицы результатов также имеет GPUBufferUsage.COPY_SRC , поскольку он будет скопирован в другой буфер для чтения после выполнения всех команд очереди графического процессора.

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 для вычислительного шейдера. С другой стороны, группа привязок, определенная для этого макета группы привязок, связывает буферы графического процессора с записями: 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"
  }
});

Подача команд

После создания экземпляра группы привязки с нашими тремя буферами графического процессора и вычислительного конвейера с макетом группы привязки пришло время их использовать.

Давайте запустим программируемый кодер вычислительных проходов с помощью commandEncoder.beginComputePass() . Мы будем использовать это для кодирования команд графического процессора, которые будут выполнять умножение матрицы. Установите его конвейер с помощью passEncoder.setPipeline(computePipeline) и его группу привязки с индексом 0 с помощью passEncoder.setBindGroup(0, bindGroup) . Индекс 0 соответствует оформлению group(0) в коде WGSL.

Теперь давайте поговорим о том, как этот вычислительный шейдер будет работать на графическом процессоре. Наша цель — выполнить эту программу параллельно для каждой ячейки матрицы результатов, шаг за шагом. Например, для матрицы результатов размером 16 на 32, чтобы закодировать команду выполнения в @workgroup_size(8, 8) , мы должны вызвать passEncoder.dispatchWorkgroups(2, 4) или passEncoder.dispatchWorkgroups(16 / 8, 32 / 8) . Первый аргумент «x» — это первое измерение, второй «y» — второе измерение, а последний «z» — это третье измерение, которое по умолчанию равно 1, поскольку оно нам здесь не нужно. В мире вычислений на графических процессорах кодирование команды для выполнения функции ядра над набором данных называется диспетчеризацией.

Параллельное выполнение для каждой ячейки матрицы результатов
Параллельное выполнение для каждой ячейки матрицы результатов

Размер сетки рабочей группы для нашего вычислительного шейдера равен (8, 8) в нашем коде WGSL. По этой причине «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() . Затем создайте буфер графического процессора, который будет использоваться в качестве места назначения для копирования буфера матрицы результатов с помощью copyBufferToBuffer . Наконец, завершите команды кодирования с помощью copyEncoder.finish() и отправьте их в очередь устройства графического процессора, вызвав метод device.queue.submit() с командами графического процессора. 

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

Чтение матрицы результатов

Считать результирующую матрицу так же просто, как вызвать gpuReadBuffer.mapAsync() с GPUMapMode.READ и дождаться разрешения возвращаемого обещания, которое указывает на то, что буфер графического процессора теперь сопоставлен. На этом этапе можно получить сопоставленный диапазон с помощью 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: [

Результаты производительности

Так чем же умножение матриц на графическом процессоре отличается от его выполнения на процессоре? Чтобы выяснить это, я написал только что описанную программу для центрального процессора. И, как вы можете видеть на графике ниже, использование полной мощности графического процессора кажется очевидным выбором, когда размер матриц превышает 256 на 256.

Тест графического процессора и процессора
Тест графического процессора и процессора

Эта статья была лишь началом моего пути изучения WebGPU . В ближайшее время ожидайте новых статей, в которых будет более подробно рассмотрены вычисления на графическом процессоре и то, как рендеринг (холст, текстура, сэмплер) работает в WebGPU.