Au moment de la rédaction de cet article, StableHLO est prêt à remplacer MHLO/HLO en tant qu'interface de compilation. Il peut être produit par TensorFlow, JAX et PyTorch, il peut être consommé par XLA et plusieurs plug-ins PJRT tiers, et il dispose de toutes les fonctionnalités publiques fournies par MHLO/HLO, ainsi que de fonctionnalités supplémentaires.
Ce document décrit les prochaines étapes du projet StableHLO, en catégorisant le travail en cours reflété dans le suivi des problèmes et en organisant ce travail en livrables planifiés.
Étapes actuelles
Nos jalons actuels suivent deux tendances principales:
- Maximiser les avantages de StableHLO pour l'ensemble de la communauté OpenXLA.
- Unifier l'expérience développeur pour tous les membres d'OpenXLA
- Abandon de MHLO: au 4e trimestre 2024, nous avons commencé à explorer l'abandon interne de MHLO, en migrant les passes utiles, y compris la canonisation et les modèles de dossiers, vers StableHLO. Une fois que le processus de migration aura été prouvé trivial en interne, nous prévoyons de partager une RFC avec des délais pour les migrations externes vers StableHLO. Cela se produira probablement au premier trimestre 2025. Nous prévoyons de laisser suffisamment de temps et d'assistance aux équipes pour qu'elles migrent vers StableHLO au premier semestre 2025.
- Migrer les optimisations indépendantes du matériel vers StableHLO: conformément à la tendance ci-dessus, nous souhaitons que StableHLO soit le meilleur endroit pour consolider les simplifications de graphiques indépendantes du matériel, afin que tous les plug-ins PJRT, y compris ceux qui convertissent de StableHLO vers un IR de compilateur non XLA, puissent bénéficier au maximum. Pour ce faire, nous devons consolider les modèles utilisés dans Google AI Edge, le projet JAX-Enzyme et d'autres projets, tous situés dans le dépôt StableHLO. Une partie de cette consolidation a déjà commencé, mais le flux de travail sera largement repris et terminé au premier trimestre 2025.
- Componentisation OpenXLA: nous avons commencé à créer des composants dédiés dans openxla/xla pour HLO, qui ressemblent à la configuration du dépôt StableHLO (référence), et avons commencé à déplacer tous les backends OpenXLA derrière les plug-ins PJRT. Nous investissons également dans la résolution des problèmes d'expérience utilisateur importants que nous découvrons dans ces plug-ins PJRT, par exemple en assurant une communication précise sur la version StableHLO dans les plug-ins StableHLO afin que les nouvelles fonctionnalités puissent être utilisées immédiatement par les nouveaux plug-ins (référence).
- Fonctionnement des composites de bout en bout: au troisième trimestre 2024, nous avons ajouté des composites à HLO, ce qui permet de prendre en charge entièrement la pile du compilateur pour les abstractions. Au quatrième trimestre 2024, nous avons appris au compilateur XLA à connaître les composites et ajouté des passes dans HLO/StableHLO pour l'intégration de composites inconnus avec leurs décompositions. Nous étudions actuellement l'ajout d'API JAX dédiées pour générer des composites à partir du framework (les API PyTorch existent déjà), ainsi que l'ajout de documentation Colab sur l'utilisation correcte des composites, qui sera finalisé au quatrième trimestre 2024.
Anciens jalons
Au cours du premier semestre 2024, nous avons publié StableHLO v1.0, qui a implémenté des améliorations de haute priorité, y compris la finalisation de la spécification publique de l'ensemble d'opérations, l'interprète de référence, l'extension de la compatibilité ascendante et descendante pour prendre en charge le déploiement sur l'appareil avec des tests unitaires de compatibilité, l'extensibilité via des opérations composites, la prise en charge du dynamisme spécifié, la publication d'une suite de données de test complète avec des résultats d'or évalués, et plus encore.