執筆時点では、StableHLO はコンパイラ インターフェースとして MHLO/HLO に代わる準備ができています。TensorFlow、JAX、PyTorch で生成でき、XLA と複数のサードパーティ PJRT プラグインで使用できます。また、MHLO/HLO が提供するすべての公開機能と追加機能を備えています。
このドキュメントでは、StableHLO プロジェクトの次のステップについて説明します。問題トラッカーに反映されている進行中の作業を分類し、この作業を計画された成果物に整理します。
現在のマイルストーン
現在のマイルストーンは、主に次の 2 つのトレンドに沿っています。
- OpenXLA コミュニティ全体で StableHLO のメリットを最大限に活用する。
- すべての OpenXLA メンバーのデベロッパー エクスペリエンスを統一します。
- MHLO の非推奨: 2024 年第 4 四半期に、MHLO の社内非推奨の検討を開始し、正規化やフォルダパターンなどの有用なパスを StableHLO に移行しました。移行プロセスが社内で簡単であることが確認できたら、StableHLO への外部移行のタイムラインを含む RFC を共有する予定です。これは 2025 年第 1 四半期に行われる可能性があり、2025 年前半に StableHLO に移行するための十分な時間とサポートをチームに提供することを計画しています。
- ハードウェアに依存しない最適化を StableHLO に移行する: 上記の傾向に沿って、StableHLO をハードウェアに依存しないグラフの簡素化を統合する最適な場所にしたいと考えています。これにより、StableHLO から XLA 以外のコンパイラ IR に変換するプラグインを含むすべての PJRT プラグインが最大限のメリットを得ることができます。この目標の一部として、Google AI Edge、JAX-Enzyme プロジェクト、その他のプロジェクトで使用されているパターンを StableHLO リポジトリに統合します。この統合の一部はすでに始まっていますが、ワークストリームの大部分は 2025 年第 1 四半期に開始され、完了する予定です。
- OpenXLA のコンポーネント化: StableHLO リポジトリの設定（参照）に似た HLO 用の openxla/xla に専用のコンポーネントの作成を開始しました。また、すべての OpenXLA バックエンドを PJRT プラグインの背後に移動し始めました。また、これらの PJRT プラグインで見つかった目立つ UX の問題の修正にも取り組んでいます。たとえば、StableHLO プラグインで StableHLO バージョンの正確な通信を行うなど、新しい機能を新しいプラグインですぐに使用できるようにしています（参照）。
- コンポジットをエンドツーエンドで機能させる: 2024 年第 3 四半期に、HLO にコンポジットを追加し、抽象化の完全なコンパイラ スタック サポートを可能にしました。2024 年第 4 四半期に、XLA インライナーに複合体について学習させ、未知の複合体を分解してインライン化するためのパスを HLO/StableHLO に追加しました。現在、フレームワークからコンポジットを生成するための専用の JAX API の追加（PyTorch API はすでに存在）と、コンポジットを適切に使用する方法に関する Colab ドキュメントの追加を検討しています。これは 2024 年第 4 四半期に完了する予定です。
過去のマイルストーン
2024 年第 1 四半期に、StableHLO v1.0 をリリースしました。これには、公開 opset 仕様の完成、リファレンス インタープリタ、上位互換性と下位互換性を拡張して、互換性単体テストによるオンデバイス デプロイをサポートする、コンポジット オペレーションによる拡張性、仕様化されたダイナミズム サポート、評価済みのゴールド結果を含む完全なテストデータ スイートのリリースなど、優先度の高い改善が実装されています。