Команда 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.