Дорожная карта StableHLO

На момент написания StableHLO готов заменить MHLO/HLO в качестве интерфейса компилятора. Его можно создавать с помощью TensorFlow, JAX и PyTorch, его можно использовать с помощью XLA и нескольких плагинов 3p PJRT, и он обладает всеми общедоступными функциями, предоставляемыми MHLO/HLO, а также дополнительными функциями.

В этом документе описываются следующие шаги проекта StableHLO, классифицируя текущую работу, отраженную в системе отслеживания проблем, и распределяя эту работу по запланированным результатам.

Текущие вехи

Наши текущие достижения следуют двум основным тенденциям:

  1. Получите максимальную выгоду от StableHLO для всего сообщества OpenXLA.
  2. Унифицируйте опыт разработчиков для всех участников OpenXLA.
  • Устаревание MHLO : в четвертом квартале 24 года мы начали изучать внутреннюю устаревание MHLO, перенося полезные этапы, включая канонизацию и шаблоны папок, в StableHLO. Как только процесс миграции окажется тривиальным внутри, мы планируем опубликовать RFC с указанием сроков внешней миграции на StableHLO. Скорее всего, это произойдет в первом квартале 25 года, и мы планируем предоставить командам достаточно времени и поддержки для перехода на StableHLO в первом полугодии 25 года.
  • Перенесите аппаратно-независимые оптимизации в StableHLO . Следуя вышеописанной тенденции, мы хотим, чтобы StableHLO был лучшим местом для консолидации аппаратно-независимых упрощений графов, чтобы все плагины PJRT, включая те, которые конвертируют из StableHLO в IR компилятора, отличного от XLA, могли получить максимальную выгоду. Часть этой цели включает в себя объединение шаблонов, используемых в Google AI Edge , проекте JAX-Enzyme и других проектах, в репозитории StableHLO. Часть этой консолидации уже началась, но рабочий процесс в основном активизируется и завершится в первом квартале 2025 года.
  • Компонентизация OpenXLA : мы начали создавать специальные компоненты в openxla/xla для HLO, которые напоминают настройку репозитория StableHLO ( ссылка ), а также начали перемещать все серверные части OpenXLA за плагины PJRT. Мы также вкладываем средства в исправление заметных проблем с пользовательским интерфейсом, которые мы обнаруживаем в этих плагинах PJRT, включая такие вещи, как обеспечение точной связи версии StableHLO в плагинах StableHLO, чтобы новые функции могли быть немедленно использованы новыми плагинами ( ссылка ).
  • Заставьте композиты работать e2e : в третьем квартале 2024 года мы добавили композиты в HLO, обеспечив полную поддержку абстракций стеком компилятора. В четвертом квартале 2024 года мы научили инлайнер XLA композитам и добавили проходы в HLO/StableHLO для встраивания неизвестных композитов с их разложением. Сейчас мы изучаем возможность добавления выделенных API-интерфейсов JAX для создания композитов из платформы (API PyTorch уже существуют), а также добавления документации Colab о том, как правильно использовать композиты, которая будет завершена в четвертом квартале 2024 года.

Прошлые вехи

В первом полугодии 2024 года мы выпустили StableHLO v1.0 , в котором реализованы высокоприоритетные улучшения, в том числе завершение общедоступной спецификации opset, эталонного интерпретатора, расширение прямой и обратной совместимости для поддержки развертывания на устройстве с модульным тестированием совместимости, расширяемость с помощью составных операций, спецификация. поддержка динамизма, выпущен полный набор тестовых данных с оцененными золотыми результатами и многое другое.