PyTorch представляет OpenReg — симулятор акселераторов для тестирования аппаратного обеспечения

Команда PyTorch анонсировала выпуск OpenReg — самодостаточного симулятора акселераторов, предназначенного для упрощения интеграции пользовательского аппаратного обеспечения с фреймворком. Новый инструмент использует механизм PrivateUse1 DispatchKey, ставший стандартом сообщества для подключения внешних устройств.

Архитектура и назначение

OpenReg позиционируется как минималистичная референсная реализация, а не полнофункциональный высокопроизводительный бэкенд. Его основная цель — служить эталоном для проверки механизмов интеграции новых аппаратных ускорителей в экосистему PyTorch.

Архитектура OpenReg

Ключевые компоненты системы

Архитектура проекта включает три основных аспекта:

• CPU-симулятор акселератора — позволяет разрабатывать, тестировать и отлаживать пользовательские операции PyTorch без доступа к реальному оборудованию
• Интеграция с CI/CD — автоматизированное тестирование в рамках процессов непрерывной интеграции
• Сопутствующая документация — подробные руководства по интеграции оборудования

Функциональные возможности

OpenReg предоставляет разработчикам комплексный набор инструментов для работы с пользовательскими акселераторами:

• Регистрация операторов — поддержка как встроенных операторов PyTorch, так и пользовательских реализаций
• Механизм fallback — обеспечивает непрерывную работу при недоступности зарегистрированных операторов
• Поддержка runtime — управление событиями, потоками, устройствами и генерацией случайных чисел
• Кросс-платформенность — работа на Linux, macOS и Windows
• Автоматическая mixed precision — оптимизация вычислений с использованием разных типов данных
• Автозагрузка — автоматическая инициализация при импорте torch

Установка и использование

Для работы с OpenReg требуется сборка из исходного кода, которая выполняется простой командой:

cd test/cpp_extensions/open_registration_extension/torch_openreg
python -m pip install --no-build-isolation -e . # для разработки
python -m pip install --no-build-isolation . # для установки

Использование симулятора максимально упрощено и аналогично работе с CUDA-устройствами:

import torch
if not torch.openreg.is_available():
 print("OpenReg backend is not available in this build.")
 exit()
print("OpenReg backend is available!")
device = torch.device("openreg")
x = torch.tensor([[1., 2.], [3., 4.]], device=device)
y = x + 2
print("Result y:n", y)
print(f"Device of y: {y.device}")
z = y.cpu()
print("Result z:n", z)
print(f"Device of z: {z.device}")

Появление OpenReg — логичный шаг в эволюции PyTorch как открытой платформы. Вместо того чтобы каждый производитель оборудования изобретал велосипед для интеграции, сообщество предлагает стандартизированный подход. Особенно ценно, что это не просто документация, а работающий код с тестами — именно так и должны выглядеть современные open-source проекты. Интересно, сколько времени пройдет до появления первых коммерческих реализаций на основе этого симулятора.

Перспективы развития

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

По материалам PyTorch Blog.