Al momento della stesura di questo articolo, StableHLO è pronto a sostituire MHLO/HLO come interfaccia del compilatore. Può essere prodotto da TensorFlow, JAX e PyTorch, può essere utilizzato da XLA e da diversi plug-in PJRT di terze parti e include tutte le funzionalità pubbliche fornite da MHLO/HLO, oltre a funzionalità aggiuntive.
Questo documento descrive i passaggi successivi per il progetto StableHLO, classifica il lavoro in corso registrato nel tracker dei problemi e organizza questo lavoro in risultati pianificati.
Traguardi attuali
I nostri attuali traguardi seguono due tendenze principali:
- Massimizzare il vantaggio di StableHLO per l'intera community OpenXLA.
- Unificare l'esperienza degli sviluppatori per tutti i membri di OpenXLA.
- Ritiro di MHLO: nel quarto trimestre del 2024 abbiamo iniziato a esplorare il ritiro interno di MHLO, eseguendo la migrazione di passaggi utili, tra cui la canonizzazione e i pattern di cartelle, a StableHLO. Una volta che il processo di migrazione sarà dimostrato essere semplice internamente, prevediamo di condividere un RFC con le tempistiche per le migrazioni esterne a StableHLO. Questo dovrebbe accadere nel primo trimestre del 2025 e prevediamo di dare ai team tutto il tempo e l'assistenza necessari per eseguire la migrazione a StableHLO nel primo semestre del 2025.
- Esegui la migrazione delle ottimizzazioni indipendenti dall'hardware a StableHLO: in linea con la tendenza sopra indicata, vogliamo che StableHLO sia il luogo migliore per consolidare le semplificazioni dei grafici indipendenti dall'hardware, in modo che tutti i plug-in PJRT, inclusi quelli che eseguono la conversione da StableHLO a un IR del compilatore non XLA, possano trarre il massimo vantaggio. Parte di questo obiettivo prevede il consolidamento dei pattern utilizzati in Google AI Edge, nel progetto JAX-Enzyme e in altri progetti nel repository StableHLO. Parte di questo consolidamento è già stata avviata, ma il flusso di lavoro verrà in gran parte ripreso e completato nel primo trimestre del 2025.
- Componentizzazione di OpenXLA: abbiamo iniziato a creare componenti dedicati in openxla/xla per HLO che assomigliano alla configurazione del repo StableHLO (ref), nonché iniziato a spostare tutti i backend OpenXLA dietro i plug-in PJRT. Inoltre, stiamo investendo per risolvere i problemi di UX più importanti rilevati in questi plug-in PJRT, ad esempio la comunicazione precisa della versione StableHLO nei plug-in StableHLO, in modo che le nuove funzionalità possano essere utilizzate immediatamente dai nuovi plug-in (ref).
- Rendere i compositi operativi dall'inizio alla fine: nel terzo trimestre del 2024 abbiamo aggiunto i compositi a HLO, attivando il supporto completo dello stack del compilatore per le astrazioni. Nel quarto trimestre del 2024 abbiamo insegnato all'inliner XLA i compositi e abbiamo aggiunto passaggi in HLO/StableHLO per l'inserimento in linea dei compositi sconosciuti con le relative decomposizioni. Al momento stiamo valutando l'aggiunta di API JAX dedicate per la generazione di compositi dal framework (le API PyTorch esistono già), nonché l'aggiunta della documentazione di Colab su come utilizzare correttamente i compositi, da completare nel quarto trimestre del 2024.
Traguardi passati
Nel primo semestre del 2024 abbiamo rilasciato la versione 1.0 di StableHLO, che ha implementato miglioramenti di alta priorità, tra cui il completamento delle specifiche dell'insieme di operazioni pubblico, l'interprete di riferimento, l'estensione della compatibilità sia in avanti che all'indietro per supportare il deployment on-device con test di unità di compatibilità, l'estensibilità tramite operazioni composite, il supporto del dinamismo specificato, il rilascio di una suite completa di dati di test con risultati di riferimento valutati e altro ancora.