W chwili pisania tego tekstu StableHLO jest gotowe do zastąpienia MHLO/HLO jako interfejs kompilatora. Może być generowany przez TensorFlow, JAX i PyTorch, może być używany przez XLA i kilka wtyczek PJRT innych firm. Zawiera wszystkie funkcje publiczne udostępniane przez MHLO/HLO, a także dodatkowe funkcje.
Ten dokument opisuje kolejne kroki w ramach projektu StableHLO, kategoryzując bieżące prace widoczne w systemie śledzenia problemów i porządkując je według zaplanowanych wyników.
Aktualne kamienie milowe
Nasze obecne kamienie milowe są zgodne z 2 głównymi trendami:
- Maksymalizacja korzyści z StableHLO dla całej społeczności OpenXLA.
- Ujednolicenie środowiska programisty dla wszystkich członków OpenXLA.
- Wycofanie MHLO: w IV kwartale 2024 r. zaczęliśmy wewnętrznie wycofywać MHLO, przenosząc przydatne przejścia, w tym kanonizację i wzorce folderów, do StableHLO. Gdy proces migracji okaże się prosty, udostępnimy dokument RFC z harmonogramem migracji zewnętrznych na stabilną wersję HLO. Prawdopodobnie nastąpi to w I kwartale 2025 r. Planujemy dać zespołom wystarczająco dużo czasu i pomocy, aby mogły przejść na StableHLO w I kwartale 2025 r.
- Przenieś optymalizacje niezależne od sprzętu do StableHLO: zgodnie z trendem opisanym powyżej chcemy, aby StableHLO było najlepszym miejscem do konsolidacji uproszczeń grafu niezależnych od sprzętu, tak aby wszystkie wtyczki PJRT, w tym te, które są konwertowane z StableHLO na kompilator IR niebędący XLA, mogły korzystać z maksimum korzyści. Częścią tego celu jest konsolidacja wzorów używanych w Google AI Edge, projekcie JAX-Enzyme i innych projektach w repozytorium StableHLO. Część tego procesu konsolidacji już się rozpoczęła, ale większość prac zostanie wykonana w pierwszym kwartale 2025 r.
- Komponentyzacja OpenXLA: zaczęliśmy tworzyć specjalne komponenty w openxla/xla dla HLO, które przypominają konfigurację repozytorium StableHLO (ref), a także zaczęliśmy przenosić wszystkie backendy OpenXLA do wtyczek PJRT. Dodatkowo inwestujemy w rozwiązanie widocznych problemów z UX, które wykryliśmy w tych wtyczkach PJRT. Chodzi m.in. o dokładne informowanie o wersji StableHLO w wtyczkach StableHLO, aby nowe wtyczki mogły od razu korzystać z nowych funkcji (ref).
- Zwiększenie wydajności złożonych operacji: w III kwartale 2024 r. dodaliśmy do HLO złożone operacje, umożliwiając pełne wsparcie dla abstrakcji w kompilatorze. W IV kwartale 2024 r. dodaliśmy do XLA inlinera do tworzenia złożonych elementów i dodaliśmy przejścia w HLO/StableHLO do tworzenia nieznanych złożonych elementów z ich rozkładów. Obecnie rozważamy dodanie interfejsów JAX do generowania kompozycji z ramy (interfejsy PyTorch już istnieją) oraz dodanie dokumentacji Colab na temat prawidłowego korzystania z kompozycji. Planujemy to na IV kwartał 2024 r.
Dotychczasowe kamienie milowe
W pierwszej połowie 2024 r. wydaliśmy StableHLO w wersji 1.0, w której wdrożyliśmy ulepszenia o wysokiej priorytecie, m.in. ukończyliśmy specyfikację publicznego przesunięcia, ulepszyliśmy zgodność z wersjami nowszymi i starszymi, aby obsługiwać wdrażanie na urządzeniu z testowaniem jednostek zgodności, rozszerzyliśmy możliwości rozszerzalności za pomocą operacji złożonych, wdrożyliśmy wsparcie dla specyfikacji dynamizmu, udostępniliśmy pełny zestaw danych testowych z ocenionymi wynikami złotego standardu i wiele więcej.