Mengembangkan backend baru untuk XLA

Panduan ini ditujukan untuk engineer sistem yang ingin XLA menghasilkan output program yang menargetkan hardware mereka secara efisien. Panduan ini tidak menjelaskan langkah demi langkah dan mengasumsikan Anda telah memiliki pengetahuan tentang LLVM, Bazel, dan XLA.

XLA menyediakan antarmuka abstrak yang dapat diterapkan oleh arsitektur atau akselerator baru untuk membuat backend guna menjalankan program ML yang di-output oleh XLA. Penargetan ulang XLA seharusnya jauh lebih sederhana dan lebih skalabel daripada menerapkan setiap operasi yang ada dari framework frontend seperti PyTorch atau TensorFlow untuk hardware baru.

Sebagian besar implementasi akan termasuk dalam salah satu skenario berikut:

  1. Arsitektur CPU yang sudah ada belum didukung secara resmi oleh XLA, dengan atau tanpa backend LLVM yang sudah ada.
  2. Hardware seperti non-CPU dengan backend LLVM yang sudah ada.
  3. Hardware seperti non-CPU tanpa backend LLVM yang ada.

Skenario 1: Arsitektur CPU yang ada belum secara resmi didukung oleh XLA

Dalam skenario ini, mulailah dengan melihat backend CPU XLA yang ada. XLA memudahkan penargetan CPU yang berbeda menggunakan LLVM, karena perbedaan utama antara backend XLA untuk CPU adalah kode yang dihasilkan oleh LLVM.

Jika vendor hardware memiliki backend LLVM untuk hardware mereka, Anda dapat dengan mudah menautkan backend dengan LLVM yang dibuat dengan XLA. Dalam mode JIT, backend CPU XLA memancarkan kode untuk CPU host. Untuk kompilasi di awal, xla::AotCompilationOptions dapat menyediakan triple LLVM untuk mengonfigurasi arsitektur target.

Jika backend LLVM tidak ada, tetapi ada jenis generator kode lain, sebagian besar backend CPU yang ada seharusnya dapat digunakan kembali.

Skenario 2: Hardware yang tidak menyerupai CPU dengan backend LLVM yang ada

Anda dapat membuat model implementasi xla::Compiler baru pada class xla::CPUCompiler dan xla::GPUCompiler yang ada, karena class tersebut sudah mengeluarkan IR LLVM. Bergantung pada sifat hardware, ada kemungkinan bahwa banyak aspek pembuatan IR LLVM harus diubah, tetapi banyak kode dapat dibagikan dengan backend yang ada.

Contoh bagus untuk diikuti adalah backend GPU dari XLA. Backend GPU menargetkan ISA non-CPU sehingga beberapa aspek pembuatan kodenya unik untuk domain GPU. Jenis hardware lain, misalnya DSP seperti Hexagon (yang memiliki backend LLVM upstream), dapat menggunakan kembali bagian-bagian logika emisi IR LLVM, tetapi bagian lainnya akan bersifat unik.

Skenario 3: Hardware seperti non-CPU tanpa backend LLVM yang ada

Jika LLVM tidak dapat digunakan, opsi terbaiknya adalah mengimplementasikan backend baru untuk XLA bagi hardware yang diinginkan. Opsi ini memerlukan upaya paling besar. Class yang perlu diimplementasikan adalah sebagai berikut:

  • StreamExecutor: Untuk banyak perangkat, tidak semua metode StreamExecutor diperlukan. Lihat penerapan StreamExecutor yang sudah ada untuk detailnya.
  • xla::Compiler: Class ini mengenkapsulasi kompilasi komputasi HLO ke dalam xla::Executable.
  • xla::Executable: Class ini digunakan untuk meluncurkan komputasi yang dikompilasi di platform.
  • xla::TransferManager: Class ini memungkinkan backend menyediakan mekanisme khusus platform untuk membuat data literal XLA dari handle memori perangkat tertentu. Dengan kata lain, data ini membantu mengenkapsulasi transfer data dari host ke perangkat, lalu kembali.