WebGPU бэкенд: отличия от WebGL, compute shaders, статус поддержки

ONNX Runtime Web предоставляет возможность запуска нейросетевых моделей прямо в браузере с использованием различных графических бэкендов, среди которых WebGL и WebGPU. WebGPU — современный графический и вычислительный API, разработанный как преемник WebGL, с целью дать веб-приложениям доступ к возможностям GPU на уровне, близком к нативным приложениям. Основные отличия WebGPU от WebGL заключаются в архитектуре, парадигме вычислений и производительности.

  • Подход к вычислениям: WebGL изначально ориентирован на рендеринг графики и использует фрагментные шейдеры для выполнения операций, что накладывает ограничения на гибкость и эффективность вычислений. WebGPU создан с акцентом на общие вычисления (general-purpose computing), предоставляя полноценные compute shaders, что позволяет эффективно реализовать линейную алгебру, свертки и матричные операции, критичные для нейросетей.
  • Контроль над ресурсами GPU: WebGPU предоставляет более прямой доступ к буферам, текстурам и состояниям GPU, снижая накладные расходы и повышая предсказуемость выполнения операций. WebGL скрывает многие детали за графическими абстракциями, что иногда вызывает неоптимальное использование памяти и вычислительных ресурсов.
  • Асинхронность и синхронизация: WebGPU поддерживает явное управление очередями команд и синхронизацией, позволяя реализовать параллельное выполнение нескольких вычислительных потоков. WebGL ограничен синхронизацией через рендер-пайпы, что затрудняет эффективное выполнение сложных вычислительных графов.

Compute shaders в ONNX Runtime Web

Compute shaders — ключевой элемент WebGPU бэкенда, позволяющий запускать вычислительные задачи на GPU без привязки к графическому конвейеру. В контексте ONNX Runtime Web это используется для операций линейной алгебры: матричных умножений, сверток, нормализаций, операций активации и других тензорных вычислений.

  • Преимущества использования compute shaders:

    • Возможность писать низкоуровневый код вычислений, оптимизированный под конкретные размеры тензоров.
    • Повышенная производительность на современных GPU, особенно при больших батчах данных.
    • Возможность многопоточности и распределения вычислений по ядрам GPU без вмешательства CPU.
  • Особенности интеграции в ONNX Runtime Web:

    • Compute shaders компилируются в WebGPU pipeline и управляются через GPUCommandEncoder.
    • ONNX Runtime автоматически распределяет операции между CPU и GPU в зависимости от доступного оборудования и типа операции.
    • Некоторые операции имеют оптимизированные шейдеры для WebGPU, что снижает время выполнения по сравнению с WebGL бэкендом.

Статус поддержки WebGPU в браузерах

Поддержка WebGPU всё ещё развивается, но ключевые браузеры уже реализовали базовые возможности API:

  • Google Chrome: стабильная поддержка WebGPU начиная с версии 113 и выше, доступ через флаг в ранних версиях.
  • Microsoft Edge: наследует поддержку Chrome, WebGPU доступен по умолчанию.
  • Mozilla Firefox: экспериментальная поддержка через флаг dom.webgpu.enabled.
  • Safari: частичная поддержка WebGPU через WebKit, требуется включение через настройки разработчика.

Некоторые ограничения текущей реализации включают:

  • Ограничения на размеры буферов и текстур, различающиеся между платформами.
  • Разная производительность шейдеров на интегрированных GPU (Intel, Apple M-серия) по сравнению с дискретными.
  • Возможные несовпадения в поведении при асинхронных вычислениях и необходимости ручного управления синхронизацией.

Использование WebGPU в ONNX Runtime Web

Для включения WebGPU бэкенда используется инициализация ort.InferenceSession с параметром executionProviders: ['webgpu']. Это автоматически включает поддержку compute shaders для всех операций, где реализованы GPU-шейдеры.

import * as ort from 'onnxruntime-web';

const session = await ort.InferenceSession.create('model.onnx', {
  executionProviders: ['webgpu'],
});

Особенности конфигурации:

  • Буферизация данных: тензоры передаются в GPU через GPUBuffer, что снижает накладные расходы при повторных вычислениях.
  • Пакетная обработка: рекомендуется использовать большие батчи для максимального использования возможностей GPU.
  • Fallback на CPU: операции, не поддерживаемые в compute shaders, автоматически выполняются на CPU, обеспечивая корректность работы модели.

WebGPU бэкенд позволяет ONNX Runtime Web достичь более высокой производительности, особенно для моделей с интенсивными тензорными вычислениями, такими как сверточные нейронные сети и трансформеры. Он обеспечивает современный интерфейс для вычислений, приближенный к возможностям десктопного GPU, сохраняя при этом кросс-платформенность веб-приложений.