Halaman ini diterjemahkan oleh Cloud Translation API.
Switch to English

MLIR CodeGen untuk XLA

XLA beroperasi pada HloInstruction dan melakukan banyak optimasi pada representasi ini, berbagi banyak hal ini di antara perangkat yang ditargetkan. Seperti beberapa titik jadwal linier dihitung dan buffer memori ditetapkan untuk setiap nilai secara statis. Codegen khusus perangkat beroperasi dengan melintasi urutan ini dan memanggil "emitor" untuk menghasilkan representasi yang cocok untuk perangkat (misalnya fungsi LLVM tunggal per perhitungan XLA pada CPU, atau urutan "thunks" yang mengenkapsulasi operasi GPU dan kemungkinan menghasilkan PTX saat menargetkan GPU).

Sebagai langkah pementasan, saat ini kami sedang dalam proses mencegat proses tepat setelah XLA menyelesaikan fase penugasan-penyangga dan memancarkan modul MLIR dalam dialek lhlo . Dari sana kami melakukan codegen menggunakan komponen MLIR (Linalg, affine, dan terutama dialek GPU) tergantung pada perangkat.

Di bawah ini adalah rencana catatan untuk memigrasikan XLA / GPU secara bertahap dengan menggunakan lhlo sebagai input codegen.

Tugas

Tuan rumah Alat
Masukkan format HloInstruction * (Tugas 1) HloInstruction * (Tugas 1)
Format output xla :: Thunk (Tugas 2) LLVM IR (Tugas 3)
  • Tugas 1 mengubah format input host dan perangkat dari HloInstruction * ke LHLO.
  • Tugas 2 mengubah format output host dari thunks ke "beberapa landasan pendaratan untuk host" (lihat di bawah).
  • Tugas 3 memigrasi output perangkat dari LLVM IR ke beberapa bentuk MLIR. Ini opsional untuk proyek ini, dan lihat bagian "Migrasi Perangkat LLVM IR" untuk detailnya.

Proyek ini memprioritaskan memiliki model runnable ujung ke ujung dengan LHLO-emitor diaktifkan sebanyak mungkin. Ini menyiratkan bahwa daftar urutan tujuan berdasarkan prioritas:

  • Jadikan XLA / GPU dapat dijalankan dengan emulator LHLO, dengan Thunks dan emitter yang ada tidak dimodifikasi.
  • Hilangkan referensi ke HloInstruction * di LHLO, kasus per kasus:
    • Alihkan emitor lawas ke emitor berbasis MLIR (mis. Linalg), atau
    • Secara otomatis menerjemahkan emitor yang ada untuk mengambil representasi MLIR (bermigrasi ke Standar dengan Dialek GPU).

Migrasi Thunks (Tugas 2)

xla :: gpu :: Thunk adalah struktur data yang:

  • Dapat dipanggil ke dari host (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • Membawa berbagai data dalam subkelasnya.
  • Berinteraksi dengan BufferAllocation :: Slice dan StreamExecutor.
  • Meluncurkan kernel
  • Panggilan ke semua pustaka runtime.

Biaya itu termasuk:

  • Merupakan data konfigurasi op-spesifik (mis. Konfigurasi konvolusi).
  • Memigrasi bentuk op dan bentuk operan.
  • Merupakan pohon thunks (sementara, kondisi, dll).

Pekerjaan migrasi tidak tergantung dari migrasi LHLO / emitor. Di bawah sumber daya terbatas, itu diprioritaskan di balik LHLO / migrasi emitor.

Kami memiliki beberapa pilihan tentang cara menurunkan bagian sisi host dari LHLO:

  • TFRT
    • (Pro) pembungkus CUDA dan HIP yang bagus untuk digunakan.
    • (Pro) mudah untuk mengimplementasikan panggilan perpustakaan (cuDNN, cuBLAS, cuFFT, dll), karena ops TFRT diinterpretasikan oleh kode C ++.
    • (Con) sisi host sedang dalam pengembangan dan tidak diuji.
  • Kode CPU yang dimasukkan
    • (Pro) kemampuan rendah yang hebat. Buat beberapa loop dan kondisi dan selesai.
    • (Con) GPUDialect belum memodelkan rantai / aliran / asinkronisasi / alokasi perangkat.
    • (Kon) Dukungan runtime CUDA / HIP minimal (jalur toolkit, versi, pemuatan dinamis, dll).
  • XLA runtime (interpreting) yang ada

Keputusan: mengadopsi TFRT, tetapi juga mendukung kode CPU jitting di TFRT.

Migrasi Perangkat LLVM IR (Tugas 3)

Emitor unsur menghasilkan op target dengan mengisinya elemen demi elemen. Setiap elemen output tergantung pada serangkaian elemen dari operan. Semua elemen dijelaskan dengan menggabungkan buffer dengan indeks dinamis. Cukup untuk menggambarkan hampir semua operasi "matematika", tetapi untuk alasan kinerja, hanya sebagian besar operasi "matematika" yang diterapkan langsung di (Cpu | Gpu) ElementalIrEmitter.

ElementalIrEmitter unik karena:

  • Sebagian besar kode dibagi antara XLA / GPU dan CPU.
  • Ini mewakili sebagian besar ops yang terlihat dalam model, termasuk semua ops elemen-bijaksana.
  • Sebagian besar fusi semata-mata bergantung pada ElementalIrEmitter.
  • Secara struktural sederhana, karena menggambarkan ketergantungan data DAG antara elemen op dan elemen operan.
  • Ini sebagian besar portabel dan tingkat tinggi (misalnya tidak seperti GPU kReduce dan GPU kCopy).
  • Dukungan bentuk dinamis mudah untuk setidaknya operasi elemen-bijaksana.

Sekarang, untuk semua ops, dipancarkan secara elemental atau tidak, ada beberapa rasa dari status akhir dari masing-masing op XLA:

  1. Kode perangkat tetap sebagai LLVM IR.
  2. Refactor emitor lama menjadi seperti LHLO -> MLIR LLVM Dialect:
    • (Biaya) Akan menjadi pekerjaan yang dibuang jika kita ingin bermigrasi ke Standar.
    • (Manfaat) Mudah dan mekanis. Dapat dilakukan dalam waktu singkat.
    • (Manfaat) Tidak lebih menguntungkan dibandingkan dengan (1).
  3. Refactor emitter lama menjadi seperti LHLO -> MLIR GPU + Standard + Loops:
    • (Biaya) Mengangkat emitor yang ada ke Standar menghadirkan beberapa tantangan. Pointer dan GEP perlu dikonversi ke MemRefs dan SubViews. Memastikan kelengkapan amdgpu adalah satu lagi.
    • (Biaya) XLA / GPU sangat bergantung pada metadata LLVM:
      • range untuk indeks blok / utas.
      • align , dereferenceable , invariant.load , alias.scope , noalias untuk memuat / menyimpan.
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable untuk loop berurutan.
    • (Manfaat) Bisa jangka panjang. Lebih portabel.
  4. Refactor emitter lama menjadi LHLO -> Linalg, dan tulis emitor Linalg baru
    • (Biaya) Ini adalah kasus per kasus. Dibandingkan dengan opsi sebelumnya, implementasi baru yang sesuai dengan kinerja XLA perlu melalui benchmark <-> mengoptimalkan alur kerja, yang dapat menjadi biaya yang signifikan untuk beberapa operasi.
    • (Manfaat) tumpukan terpadu; dukungan masyarakat; portabilitas; lebih banyak potensi optimisasi.

Kesimpulan:

  • Jangan pilih (2). (1) atau (3) lebih baik dari (2). (2) harganya lebih mahal dari (1), karena memerlukan banyak refactoring mekanis. Dengan (1) kita masih bisa mencapai tujuan memungkinkan XLA untuk mengambil emiter MLIR. Ini dengan melakukan LHLO -> LLVM IR -> menjalankan emitor perangkat legacy.
  • Operasi ElementalIrEmitter berlaku untuk (4), tetapi tidak secara bertahap. Tidak ada cara untuk melakukannya op by op, karena semua ops yang dipancarkan secara elemen terhubung ke grafik yang sama. Pekerjaan ini juga dapat berfungsi sebagai titik penyatuan beberapa kekuatan yang sedang berjalan (xla / service / mlir_gpu, generator kernel, Linalg).
  • Semua operasi lainnya berlaku untuk (1). Sebagai sasaran peregangan, mereka mungkin dimigrasikan ke (3) atau (4).

Prioritas

Sementara ketiga tugas yang disebutkan di atas dapat diparalelkan, di bawah sumber daya terbatas mereka harus diserialisasi. Prioritisasi berfokus pada hasil yang terlihat untuk menyelesaikan setiap tugas.

Prioritasnya adalah: Task1 (LHLO untuk emitor legacy)> Task 2 (Thunks)> Task 3 (Emitor MLIR).

Pada akhir Tugas 1, pengguna XLA dapat menghasilkan LHLO (mis. Generator kernel) dan menjalankannya. Format kompilasi tidak akan menjadi serialable MLIR.

Pada akhir Tugas 2, LHLO menurunkan MLIR yang layak dan bersambung. Ini memungkinkan kompilasi offline.

Pada akhir Tugas 3, semua penghasil XLA berbasis MLIR dalam implementasinya.

Desain yang rinci

Langkah 1: (Tugas 1) Selesaikan LHLO dan Buat Legacy Emitter Ambil LHLO

Langkah ini membuat semua pemancar XLA / GPU yang ada berinteraksi dengan operasi MLIR. Langkah ini murni refactoring dan NFC.

Langkah ini sebagian besar mekanis, tetapi ada baiknya memperhatikan perbedaan berikut antara HloComputation yang tidak diuji dan LHLO:

  • Setiap HloInstruction memiliki akses langsung ke operandnya (DAG aliran data). Sebaliknya, setiap op LHLO hanya memiliki akses ke buffer operandnya (bipartit antara ops dan buffer). LHLO ops harus melalui rantai use-def untuk mengakses op operan mereka.
  • Emitor peninggalan yang belum ditelusuri secara empiris hampir tidak pernah mengakses operan mereka. Satu-satunya pengecualian adalah kReduce.
  • Emitter legacy yang belum diarsip mengakses BufferAssignment hanya untuk mendapatkan irisan, bukan untuk mengakses struktur data aux seperti dataflow_analysis () atau alias_analysis (). llvm_ir membangun alias_analysis () sendiri berdasarkan informasi slice.

Kesimpulannya adalah bahwa LHLO harus pas dengan tanpa kerumitan besar.

Langkah 2: Dukungan Profiling (Opsional)

Langkah ini hanya diperlukan jika kita mulai membuang beberapa logika XLA Thunk (lihat langkah selanjutnya).

Sebelum benar-benar menyalakan pemancar berbasis MLIR, kita perlu membuat profil untuk pemancar berbasis MLIR.

Saat ini XLA melakukan profil sendiri dengan memanggil timer StreamExecutor. Timer di bawah kap menyisipkan dua peristiwa sebelum dan sesudah peluncuran kernel, dan mengukur waktu sinkronisasi antara kedua peristiwa ini.

Ada sekitar tiga pendekatan untuk mendukung pembuatan profil di MLIR:

  • Jalankan profiler ujung ke ujung
  • Tambahkan op profil untuk setiap op di LHLO, menggunakan profiler yang diinjeksi.

Pendekatan "ujung ke ujung" transparan untuk MLIR, tetapi menderita masalah yang sama yang membuat XLA tidak menggunakannya di tempat pertama: panggilan perpustakaan yang dikumpulkan oleh profiler (nvprof / ...) tidak dapat dengan mudah berhubungan dengan HLO ops. Sebagai contoh, cuDNN meluncurkan beberapa kernel untuk setiap HLO, dan sulit untuk mengatakan kernel mana yang sesuai dengan HLO mana.

Pendekatan "profiler yang diinjeksi" membutuhkan:

  • LHLO untuk mengambil profiler sebagai parameter.
  • memasukkan profile.start / profile.end sebelum dan sesudah setiap op.
  • pass dari profil yang lebih rendah. {start, end} ke implementasi C ++.

Penentuan profil yang tepat tidak dapat dengan mudah dilakukan untuk operasi yang dihasilkan MLIR, karena:

  • MLIR tidak memiliki timer, atau tergantung pada TFRT / StreamExecutor.
  • MLIR tidak mudah memanggil fungsi C dengan parameter yang rumit.

Langkah 3: (Tugas 2) Memigrasi Thunks

Sebagai catatan, ada sekitar tiga macam thunks:

  • KernelThunk, yang meluncurkan kernel.
  • Control flow thunks, yang memiliki logika kontrol aliran host (kondisional, sementara, untuk, urutan) dan meluncurkan kernel body.
  • Pustaka perpustakaan: cuDNN, cuBLAS, cuFFT, NCCL, dll.

Rencananya adalah:

  • Jadikan Thunks (de) serializable.
  • Bantu meningkatkan TFRT ke kondisi di mana ia dapat mendukung semantik ini.
  • Saat keadaan membaik, migrasikan individual thunks secara bertahap.

Item tindakan ini hanya dipesan sebagian. Urutan eksekusi aktual / paralelisme teknik harus dievaluasi saat berjalan.

Langkah 4: (Tugas 3) Migrasi ElementalIrEmitter

Setelah profiling siap, kita dapat menyelesaikan dan menyetel semua emitor berbasis ElementalIrEmitter di MLIR. Kemudian kami menyalakannya secara default, dengan asumsi bahwa semua emitor berbasis MLIR ini menggunakan aliran tunggal.

Perhatikan bahwa itu juga menguntungkan untuk melakukan migrasi ElementalIrEmitter XLA / CPU, karena mereka berbagi sebagian besar kode.

Dengan semua pembandingan dan perburuan kinerja selesai (TODO: define parity kinerja), kami mengaktifkan emitor elemen berbasis MLIR baru, dan menghapus ElementalIrEmitter warisan.

Langkah ini juga menyediakan transisi fusi mudah (op bersarang) untuk migrasi nanti.

Langkah 5: Dukungan Multi-Stream atau Drop

Kami tidak dapat menghapus beberapa emitor sampai kami mendukungnya di MLIR, atau kami membatalkan fitur. Ini adalah jumlah pekerjaan yang relatif besar di MLIR dan sedikit keuntungan untuk XLA. Kami harus menyelidiki pengguna saat ini dari pengguna XLA / GPU multi-stream, dan mencoba untuk menghapus fitur ini jika masuk akal.

Langkah 6: (Tugas 3) Ops Perangkat Bermigrasi

Langkah ini memigrasikan semua ops yang tidak diuji, lalu kami dapat menghapus semua emitor yang tidak diuji.

Ini membutuhkan penulisan ulang / refactor untuk kCopy dan kReduce. kReduce sudah dikerjakan untuk banyak, sehingga jumlah aktual pekerjaan yang perlu dilakukan masih harus dilihat.