Arsitektur XLA

Mengapa kami membangun XLA?

Kami memiliki beberapa tujuan agar XLA dapat bekerja dengan TensorFlow:

  • Meningkatkan kecepatan eksekusi. Kompilasi subgrafik untuk mengurangi waktu eksekusi Ops yang berumur pendek untuk menghilangkan overhead dari runtime TensorFlow, menggabungkan operasi pipeline untuk mengurangi overhead memori, dan mengkhususkan diri pada bentuk tensor yang diketahui untuk memungkinkan propagasi konstan yang lebih agresif.

  • Meningkatkan penggunaan memori. Menganalisis dan menjadwalkan penggunaan memori, pada prinsipnya menghilangkan banyak buffer penyimpanan perantara.

  • Kurangi ketergantungan pada Ops kustom. Hilangkan kebutuhan akan banyak Ops kustom dengan meningkatkan kinerja Ops tingkat rendah yang digabungkan secara otomatis agar sesuai dengan kinerja Ops kustom yang digabungkan dengan tangan.

  • Kurangi jejak seluler. Hilangkan waktu proses TensorFlow dengan mengompilasi subgraf sebelumnya dan memancarkan pasangan file objek/header yang dapat ditautkan langsung ke aplikasi lain. Hasilnya dapat mengurangi jejak untuk inferensi seluler beberapa kali lipat.

  • Meningkatkan portabilitas. Buatlah relatif mudah untuk menulis backend baru untuk perangkat keras baru, di mana sebagian besar program TensorFlow akan berjalan tanpa dimodifikasi pada perangkat keras tersebut. Ini kontras dengan pendekatan yang mengkhususkan Ops monolitik individual untuk perangkat keras baru, yang mengharuskan program TensorFlow ditulis ulang untuk menggunakan Ops tersebut.

Bagaimana cara kerja XLA?

Bahasa input ke XLA disebut "HLO IR", atau hanya HLO (Operasi Tingkat Tinggi). Semantik HLO dijelaskan pada Operasi Semantik halaman. Hal ini paling nyaman untuk memikirkan HLO sebagai IR compiler .

XLA mengambil grafik ("perhitungan") yang didefinisikan dalam HLO dan mengkompilasinya menjadi instruksi mesin untuk berbagai arsitektur. XLA adalah modular dalam arti bahwa itu adalah mudah untuk slot yang dalam backend alternatif untuk menargetkan beberapa arsitektur HW baru . Backend CPU untuk x64 dan ARM64 serta backend GPU NVIDIA ada di pohon sumber TensorFlow.

Diagram berikut menunjukkan proses kompilasi di XLA:

XLA dilengkapi dengan beberapa optimasi dan melewati analisis yang adalah target-independen, seperti CSE , target operasi independen fusion, dan buffer analisis untuk mengalokasikan memori runtime untuk perhitungan.

Setelah langkah target-independen, XLA mengirimkan perhitungan HLO ke backend. Backend dapat melakukan pengoptimalan tingkat HLO lebih lanjut, kali ini dengan mempertimbangkan informasi dan kebutuhan spesifik target. Misalnya, backend GPU XLA dapat melakukan penggabungan operasi yang bermanfaat secara khusus untuk model pemrograman GPU dan menentukan cara mempartisi komputasi ke dalam aliran. Pada tahap ini, backend juga dapat mencocokkan pola operasi atau kombinasi tertentu dengan panggilan pustaka yang dioptimalkan.

Langkah selanjutnya adalah pembuatan kode target-spesifik. CPU dan GPU backends disertakan dengan XLA menggunakan LLVM untuk tingkat rendah IR, optimasi, dan kode-generasi. Backend ini memancarkan LLVM IR yang diperlukan untuk mewakili komputasi XLA HLO secara efisien, dan kemudian memanggil LLVM untuk memancarkan kode asli dari LLVM IR ini.

Backend GPU saat ini mendukung GPU NVIDIA melalui backend LLVM NVPTX; backend CPU mendukung beberapa ISA CPU.