Roadmap StableHLO

Pada saat penulisan, StableHLO siap menggantikan MHLO/HLO sebagai antarmuka compiler. Model ini dapat dihasilkan oleh TensorFlow, JAX, dan PyTorch, dapat digunakan oleh XLA dan beberapa plugin PJRT pihak ketiga, serta memiliki semua fitur publik yang disediakan oleh MHLO/HLO serta fungsi tambahan.

Dokumen ini menjelaskan langkah-langkah berikutnya untuk project StableHLO, yang mengategorikan pekerjaan yang sedang berlangsung yang tercermin dalam pelacak masalah dan mengatur pekerjaan ini ke dalam hasil yang direncanakan.

Pencapaian Saat Ini

Pencapaian kami saat ini mengikuti dua tren utama:

  1. Maksimalkan manfaat StableHLO untuk seluruh komunitas OpenXLA.
  2. Menyatukan pengalaman developer untuk semua pelanggan OpenXLA.
  • Penghentian MHLO: Pada Kuartal 4'24, kami telah mulai mempelajari penghentian MHLO secara internal, yang memigrasikan kartu yang berguna, termasuk kanonisasi dan pola folder, ke StableHLO. Setelah proses migrasi terbukti sederhana secara internal, kami berencana membagikan RFC dengan linimasa untuk migrasi eksternal ke StableHLO. Hal ini kemungkinan akan terjadi pada K1'25, dan kami berencana untuk memberikan waktu dan dukungan yang cukup kepada tim untuk bermigrasi ke StableHLO pada K1'25.
  • Memigrasikan pengoptimalan independen hardware ke StableHLO: Dengan mengikuti tren di atas, kami ingin StableHLO menjadi tempat terbaik untuk menggabungkan penyederhanaan grafik independen hardware, sehingga semua plugin PJRT termasuk yang dikonversi dari StableHLO ke IR compiler non-XLA dapat melihat manfaat maksimal. Bagian dari sasaran ini melibatkan penggabungan pola yang digunakan di Google AI Edge, project JAX-Enzyme, dan project lainnya di repositori StableHLO. Beberapa konsolidasi ini telah dimulai, tetapi alur kerja sebagian besar akan dimulai dan selesai pada K1'25.
  • Pembentukan Komponen OpenXLA: Kami telah mulai membuat komponen khusus di openxla/xla untuk HLO yang menyerupai penyiapan repo StableHLO (ref), serta mulai memindahkan semua backend OpenXLA di belakang plugin PJRT. Kami juga berinvestasi dalam memperbaiki masalah UX yang mencolok yang kami temukan di plugin PJRT ini, termasuk hal-hal seperti memiliki komunikasi versi StableHLO yang akurat di plugin StableHLO, sehingga fitur baru dapat langsung digunakan oleh plugin baru (ref).
  • Membuat komposit berfungsi secara menyeluruh: Pada Kuartal 3'24, kami menambahkan komposit ke HLO, yang memungkinkan dukungan stack compiler penuh untuk abstraksi. Pada Kuartal 4'24, kami mengajarkan penginline XLA tentang komposit dan menambahkan kartu di HLO/StableHLO untuk melakukan penginlinean komposit yang tidak diketahui dengan dekomposisinya. Sekarang kami sedang menyelidiki penambahan JAX API khusus untuk membuat komposit dari framework (PyTorch API sudah ada), serta menambahkan dokumentasi Colab tentang cara menggunakan komposit dengan benar, yang akan selesai pada Kuartal 4'24.

Pencapaian Sebelumnya

Pada K1 2024, kami merilis StableHLO v1.0 yang menerapkan peningkatan prioritas tinggi, termasuk menyelesaikan spesifikasi opset publik, penafsir referensi, memperluas kompatibilitas maju dan mundur untuk mendukung deployment di perangkat dengan pengujian unit kompatibilitas, ekstensi melalui operasi komposit, dukungan dinamisme spesifikasi, merilis rangkaian data pengujian lengkap dengan hasil emas yang dievaluasi, dan lainnya.