XLA (Accelerated Linear Algebra) adalah compiler machine learning (ML) yang mengoptimalkan aljabar linear, sehingga memberikan peningkatan kecepatan eksekusi dan penggunaan memori. Halaman ini memberikan ringkasan singkat tentang tujuan dan arsitektur compiler XLA.
Tujuan
Saat ini, XLA mendukung beberapa frontend framework ML (termasuk PyTorch, TensorFlow, dan JAX) serta merupakan bagian dari project OpenXLA – sebuah ekosistem teknologi compiler open source untuk ML yang dikembangkan secara kolaboratif oleh organisasi hardware dan software ML terkemuka. Sebelum project OpenXLA dibuat, XLA dikembangkan di dalam project TensorFlow, tetapi tujuan dasarnya tetap sama:
Meningkatkan kecepatan eksekusi. Kompilasi subgrafik untuk mengurangi waktu eksekusi operasi berdurasi singkat dan menghilangkan overhead dari runtime, menggabungkan operasi pipeline untuk mengurangi overhead memori, dan mengkhususkan bentuk tensor yang diketahui agar memungkinkan propagasi konstan yang lebih agresif.
Meningkatkan penggunaan memori. Menganalisis dan menjadwalkan penggunaan memori untuk menghilangkan banyak buffer penyimpanan perantara.
Mengurangi ketergantungan pada operasi kustom. Hilangkan kebutuhan akan banyak operasi kustom dengan meningkatkan performa operasi level rendah yang digabungkan secara otomatis agar sesuai dengan performa operasi kustom yang awalnya digabungkan secara manual.
Meningkatkan portabilitas. Menulis backend baru untuk hardware baru menjadi relatif mudah, sehingga sebagian besar model ML dapat berjalan tanpa dimodifikasi pada hardware tersebut. Hal ini berbeda dengan pendekatan operasi monolitik individual yang mengkhususkan diri untuk hardware baru, yang mengharuskan model ditulis ulang untuk memanfaatkan operasi tersebut.
Cara kerjanya
Compiler XLA mengambil grafik model dari framework ML yang ditentukan dalam StableHLO dan mengompilasinya ke dalam petunjuk mesin untuk berbagai arsitektur. StableHLO menentukan set operasi berversi (HLO = operasi tingkat tinggi) yang menyediakan lapisan portabilitas antara framework ML dan compiler.
Secara umum, proses kompilasi yang mengonversi grafik model menjadi file yang dapat dieksekusi yang dioptimalkan untuk target mencakup langkah-langkah berikut:
XLA melakukan beberapa penerusan pengoptimalan dan analisis bawaan pada grafik Stabilo yang tidak bergantung pada target, seperti CSE, penggabungan operasi yang tidak bergantung pada target, dan analisis buffer untuk mengalokasikan memori runtime untuk komputasi. Selama tahap pengoptimalan ini, XLA juga mengonversi dialek StableHLO menjadi dialek HLO internal.
XLA mengirimkan komputasi HLO ke backend untuk pengoptimalan tingkat HLO lebih lanjut, kali ini dengan mempertimbangkan kebutuhan dan informasi khusus target. Misalnya, backend GPU dapat melakukan fusi operasi yang bermanfaat khususnya untuk model pemrograman GPU dan menentukan cara mempartisi komputasi menjadi aliran. Pada tahap ini, backend juga dapat cocok dengan pola operasi atau kombinasi tertentu dari operasi tersebut untuk panggilan library yang dioptimalkan.
Backend kemudian melakukan pembuatan kode khusus target. Backend CPU dan GPU yang disertakan dengan XLA menggunakan LLVM untuk IR tingkat rendah, pengoptimalan, dan pembuatan kode. Backend ini mengeluarkan IR LLVM yang diperlukan untuk merepresentasikan komputasi HLO secara efisien, lalu memanggil LLVM untuk memunculkan kode native dari IR LLVM ini.
Dalam proses ini, compiler XLA bersifat modular karena mudah dimasukkan ke backend alternatif untuk menargetkan beberapa arsitektur HW baru. Backend GPU saat ini mendukung GPU NVIDIA melalui backend LLVM NVPTX. Backend CPU mendukung beberapa ISA CPU.