XLA (Accelerated Linear Algebra) adalah compiler machine learning (ML) yang mengoptimalkan aljabar linear, yang memberikan peningkatan kecepatan eksekusi dan penggunaan memori. Halaman ini menyediakan ringkasan singkat tentang tujuan dan arsitektur compiler XLA.
Tujuan
Saat ini, XLA mendukung beberapa frontend framework ML (termasuk PyTorch, TensorFlow, dan JAX) dan merupakan bagian dari project OpenXLA, yakni 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 yang berumur pendek 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. Menuliskan backend baru untuk hardware baru saja menjadi relatif mudah, sehingga sebagian besar model ML dapat berjalan tanpa dimodifikasi pada hardware tersebut. Hal ini berbeda dengan pendekatan mengkhususkan operasi monolitik individual untuk hardware baru, yang mengharuskan model ditulis ulang untuk menggunakan 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 kumpulan 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 target mencakup langkah-langkah berikut:
XLA melakukan beberapa penerusan pengoptimalan dan analisis bawaan pada grafik StableHLO yang tidak bergantung pada target, seperti CSE, fusi operasi target independen, 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 informasi dan kebutuhan khusus target. Misalnya, backend GPU dapat melakukan fusi operasi yang bermanfaat khusus untuk model pemrograman GPU dan menentukan cara mempartisi komputasi menjadi stream. Pada tahap ini, backend juga dapat mencocokkan pola operasi atau kombinasi tertentu dengan 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 memancarkan IR LLVM yang diperlukan untuk merepresentasikan komputasi HLO dengan cara yang efisien, lalu memanggil LLVM untuk memunculkan kode native dari IR LLVM ini.
Dalam proses ini, compiler XLA bersifat modular dalam arti 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.