Mengembangkan backend baru untuk XLA

Panduan ini ditujukan bagi engineer sistem yang ingin XLA menghasilkan output program yang menargetkan hardware mereka secara efisien. Panduan ini tidak berisi langkah demi langkah dan mengasumsikan pengetahuan tentang LLVM, Bazel, dan XLA.

XLA menyediakan antarmuka abstrak yang dapat diterapkan oleh arsitektur atau akselerator baru untuk membuat backend guna menjalankan output program ML dari 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 penerapan akan termasuk dalam salah satu skenario berikut:

  1. Arsitektur CPU yang sudah ada yang belum didukung secara resmi oleh XLA, dengan atau tanpa backend LLVM yang ada.
  2. Hardware yang tidak mirip CPU dengan backend LLVM yang ada.
  3. Hardware yang tidak mirip 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 CPUXLA yang ada. XLA memudahkan untuk menargetkan CPU yang berbeda dengan 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 di-build dengan XLA. Dalam mode JIT, backend CPU XLA memancarkan kode untuk CPU {i>host<i}. Untuk kompilasi sebelumnya, 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 harus dapat digunakan kembali.

Skenario 2: Hardware yang tidak mirip 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 ini sudah mengeluarkan IR LLVM. Bergantung pada sifat hardware, ada kemungkinan banyak aspek pembuatan IR LLVM harus diubah, tetapi banyak kode dapat dibagikan dengan backend yang ada.

Contoh yang bagus untuk diikuti adalah backend GPU dari XLA. Backend GPU menargetkan ISA non-CPU seperti 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 logika emisi IR LLVM, tetapi bagian lainnya bersifat unik.

Skenario 3: Hardware yang tidak mirip CPU tanpa backend LLVM yang ada

Jika LLVM tidak dapat digunakan, opsi terbaiknya adalah mengimplementasikan backend baru bagi 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 implementasi StreamExecutor yang sudah ada untuk mengetahui 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 menyusun data literal XLA dari handle memori perangkat tertentu. Dengan kata lain, ini membantu merangkum transfer data dari host ke perangkat dan kembali.