Dari HLO ke Thunks

Dokumen ini menguraikan perjalanan modul Pengoptimal Tingkat Tinggi (HLO) XLA dari keadaan awalnya hingga menjadi file yang dapat dieksekusi akhir. Terkadang kami akan menghilangkan "modul" dan menyebutnya hanya sebagai "HLO".

Diagram HLO ke thunk

HLO pra-pengoptimalan

Kita mulai dengan modul HLO pra-pengoptimalan. HLO pra-pengoptimalan tidak berisi operasi (ops) yang dianggap internal untuk XLA, seperti fusion atau bitcast. Operasi tidak memiliki tata letak pada tahap ini, atau jika ada, tata letak tersebut akan diabaikan. HLO pra-pengoptimalan biasanya dihasilkan oleh framework tingkat yang lebih tinggi seperti TensorFlow dan JAX. Saat menggunakan flag XLA -xla_dump_to, HLO pra-pengoptimalan akan di-dump ke file dengan akhiran nama file “before_optimizations.txt”.

Mengoptimalkan Modul HLO

Pipeline XLA:GPU mengubah HLO pra-pengoptimalan menjadi HLO yang dioptimalkan dengan menjalankan serangkaian proses. Pass dapat dikelompokkan secara semantik dan dijalankan dalam urutan berikut:

Hal ini mencakup penerusan seperti Shardy Partitioner atau penerusan untuk sharding SPMD.

Penerusan pengoptimalan

Hal ini dapat mencakup pengesahan legalisasi dan penyederhanaan.

Penerusan pengoptimalan kolektif

Mirip dengan Penerusan pengoptimalan, tetapi berfokus pada operasi kolektif.

Penerusan tugas tata letak

Setiap operasi HLO diberi tata letak yang merupakan bagian dari bentuk instruksi. Tata letak mengontrol cara tensor ditata secara fisik dalam memori.

Contoh bentuk dengan tata letak:

f32[10,20,30]{2,0,1}

Setelah jenis elemen, ada dimensi logis bentuk, diikuti dengan permutasi tata letak dalam urutan kecil ke besar. Dalam contoh ini, dimensi paling kecil adalah 30, dimensi kedua paling kecil adalah 10, dan dimensi terbesar adalah 20.

Tujuan penetapan tata letak adalah untuk meminimalkan jumlah transposisi fisik yang diperlukan menggunakan strategi greedy. Dimulai dengan batasan tata letak tertentu (misalnya, library cuDNN/cuBLAS mengharapkan dimensi berurutan) dan menyebarkan tata letak “ke bawah” lalu “ke atas” grafik HLO. Di akhir propagasi tata letak, beberapa petunjuk mungkin memiliki tata letak yang bertentangan, satu dipropagasi dari operand, satu dipropagasi dari pengguna. Untuk mengatasi konflik ini, instruksi HLO copy dimasukkan yang mengubah tata letak dari tata letak operand ke tata letak instruksi.

Penerusan normalisasi tata letak

Mengingat agak sulit untuk mengetahui bentuk fisik, normalisasi tata letak mencoba menulis ulang bentuk sehingga menggunakan tata letak default {rank-1, rank-2, …, 0}. Pada contoh di atas, bentuk yang dinormalisasi adalah f32[20,10,30]{2,1,0}. Operasi penyalinan yang mengubah tata letak ditulis ulang sebagai kombinasi transpose dan bitcast. Mengingat saat ini kami tidak dapat menormalisasi semua operasi, masih ada beberapa operasi yang mungkin memiliki tata letak non-default, terutama gather dan dot. Pada batas antara operasi yang dinormalisasi dan operasi yang tidak dinormalisasi akan ada operasi bitcast yang merepresentasikan transposisi, yaitu transposisi dengan tata letak yang ditetapkan sehingga secara fisik menjadi no-op.

Normalisasi tata letak juga membuat beberapa transposisi implisit menjadi eksplisit, yang penting karena codegen dapat menangani transposisi eksplisit dengan emitter khusus. Misalnya, pembentukan ulang secara teknis diizinkan untuk memiliki tata letak fisik yang berbeda antara operand dan hasil (misalnya, karena peringkat yang berbeda). Pass ReshapeDecomposer yang berjalan sebagai bagian dari pass normalisasi tata letak mengubah pembentukan ulang menjadi urutan transpose, pembentukan ulang bitcast, dan transpose.

Mengirimkan pengoptimalan penetapan tata letak pasca-

Pass yang paling penting di sini adalah fusi Triton (fusi GEMM + fusi Softmax/Layernorm) atau penulisan ulang ke panggilan library. Penyetelan otomatis juga berjalan di langkah ini, di mana XLA memilih antara berbagai emitter, memilih algoritma terbaik untuk konvolusi atau titik, menemukan pengelompokan terbaik untuk penggabungan yang ditangani oleh emitter Triton, dll.

Kartu gabungan

Dua lintasan utama adalah fusi PriorityFusion dan Multi-Output.

Di PriorityFusion, kita membentuk gabungan yang dipandu oleh model biaya. Saat menggabungkan, kami akan mengizinkan operasi duplikat dengan beberapa pengguna jika operasi dapat digabungkan ke semua pengguna. Kami juga akan mengizinkan perluasan fusi Softmax Triton yang ada jika memungkinkan.

Penggabungan Multi-Output adalah proses terpisah yang memungkinkan penggabungan operasi/penggabungan yang berbagi operand. Operasi ini juga dapat menggabungkan operand/penggabungan operand ke pengguna tanpa duplikasi dengan menambahkan output tambahan, sehingga pengguna lain dari operasi yang akan digabungkan dapat dialihkan ke output ini. Pass ini harus berhati-hati agar tidak memperkenalkan siklus ke dalam grafik HLO.

Setelah fusi Multi-Output, penghapusan subekspresi umum (pass HloCSE) berjalan, yang berpotensi menggabungkan kembali operasi yang sebelumnya diduplikasi jika operasi tersebut berakhir dalam fusi yang sama.

Beberapa proses pasca-penggabungan

Beberapa lintasan terkait kolektif (seperti mengubahnya menjadi asinkron, atau menerapkan urutan relatif kolektif tertentu).

Terakhir, kita menjalankan CopyInsertion tempat salinan ditambahkan untuk memastikan bahwa operasi di tempat tidak menimpa data yang masih diperlukan di tempat lain.

Di akhir pengoptimalan, HLO yang dioptimalkan akan di-dump jika menggunakan tanda -xla_dump_to ke file yang memiliki akhiran nama file "after_optimizations.txt". Jika Anda ingin mengekspor HLO setelah melewati proses perantara yang benar-benar mengubah HloModule, Anda dapat menggunakan flag -xla_dump_hlo_pass_re=.* (atau ekspresi reguler tertentu untuk membatasi proses tertentu).

Penjadwalan

Modul HLO tanpa jadwal masih memiliki beberapa derajat kebebasan dalam urutan pemrosesan operasi. Setiap pengurutan topologi yang memperhatikan hubungan operand/hasil dan dependensi kontrol valid. Penjadwalan menentukan pesanan spesifik yang akan digunakan. Masalah utama pada tahap ini adalah konsumsi memori maksimum yang bergantung pada masa aktif tensor. Pada langkah awal, kami mencoba berbagai algoritma penjadwal dan memilih jadwal yang akan meminimalkan penggunaan memori puncak. Perhatikan bahwa pada tahap ini kita belum menggunakan buffer fisik (yang akan terjadi di "Penetapan Buffer") dan mensimulasikan penggunaan memori.

Kemudian, LatencyHidingSchedulerpass berjalan dan mencoba memaksimalkan tumpang-tindih komputasi-komunikasi. Namun, hal itu dapat meningkatkan penggunaan memori lagi.

Terakhir, jika konsumsi memori puncak lebih tinggi daripada jumlah memori yang tersedia, kita akan menjalankan HloRematerialization. Penerusan ini berupaya mengurangi penggunaan memori dengan mengorbankan performa, karena misalnya beberapa penggabungan mungkin dibagi dan beberapa operasi mungkin diduplikasi agar memiliki masa aktif buffer yang lebih pendek. Jika rematerialisasi terjadi, sebaiknya selidiki cara mengurangi persyaratan memori di sisi model (misalnya, menggunakan ukuran batch yang lebih kecil).

Penetapan Buffer

Tepat sebelum menurunkan ke LLVM IR, kita menjalankan proses penetapan buffer yang akan menetapkan slice buffer ke setiap instruksi dalam grafik HLO. Penetapan buffer berjalan dalam beberapa langkah:

  1. HloDataflowAnalysis menetapkan HloValues (pada dasarnya buffer logis) ke instruksi. Untuk operasi di tempat, HloValue operand dapat digunakan kembali. Operasi dapat menentukan lebih dari satu HloValue (misalnya, dengan bentuk hasil tuple).

  2. HloAliasAnalysis mencoba menggabungkan buffer untuk operasi anti-aliasing, dan menghitung pemetaan dari HloValue ke HloBuffer.

  3. BufferAssignment menghitung pemetaan HloBuffers ke slice buffer di dalam buffer besar sedemikian rupa sehingga slice buffer yang sama tidak digunakan untuk HloBuffers yang berbeda dengan masa aktif yang tumpang-tindih. Untuk operasi yang mungkin memiliki alias, tidak masalah jika ada sedikit tumpang-tindih (waktu berakhir satu HloBuffer mungkin bertepatan dengan waktu mulai HloBuffer lainnya). Saat menggunakan tanda -xla_dump_to, beberapa informasi tentang penetapan buffer di-dump ke file dengan akhiran nama "after_optimizations-buffer-assignment.txt".

Thunk

Setelah grafik HLO dioptimalkan dan dijadwalkan, grafik tersebut akan diturunkan menjadi urutan thunk linear untuk backend tertentu (CPU atau GPU).

Di XLA, Thunk adalah abstraksi unit kerja mandiri yang dijalankan oleh runtime. Hal ini dapat berupa peluncuran kernel yang dikompilasi, operasi tertentu, panggilan library, konstruksi alur kontrol, komunikasi kolektif, dan sebagainya. Urutan Thunk merepresentasikan seluruh yang dapat dieksekusi untuk backend tertentu.

Emisi Thunk

Proses konversi komputasi HLO terjadwal menjadi urutan thunk disebut "emisi thunk". Hal ini ditangani oleh class emitter khusus di setiap backend.

Untuk Backend GPU, hal ini ditangani oleh IrEmitterUnnested. EmitHloComputation melakukan iterasi melalui daftar Instruksi HLO terjadwal dalam komputasi dan mengirim ke metode Emit... khusus (misalnya, EmitFusion, EmitConvolutionThunk, EmitWhile). Setiap metode ini membuat objek Thunk yang sesuai dan menambahkannya ke urutan thunk.

Untuk Backend CPU, ThunkEmitter melakukan peran ini dan disusun dengan cara yang serupa. ThunkSequence akhir disematkan dalam CpuExecutable.

Perhatikan bahwa setiap instruksi dalam komputasi entri modul HLO mungkin tidak sesuai dengan (kTuple, kConstant, ..), satu, atau beberapa (misalnya instruksi pengurutan) thunk dalam urutan thunk akhir.

Buffer Perintah: Mengoptimalkan Eksekusi di GPU

Hardware GPU modern memungkinkan perekaman urutan operasi GPU (peluncuran kernel, penyalinan memori, dll.) sekali, lalu memutar ulang urutan tersebut beberapa kali dengan overhead CPU minimal. Ini adalah pengoptimalan performa yang penting, terutama untuk workload dengan banyak kernel kecil yang diluncurkan dengan cepat. XLA menggunakan Buffer Perintah sebagai abstraksi dari Grafik CUDA atau Grafik HIP. Antarmuka inti ditentukan dalam GpuCommandBuffer.

Buffer perintah ditampilkan dalam urutan thunk oleh CommandBufferThunk.

Emitter tidak menghasilkan thunk ini langsung dari petunjuk HLO. Sebagai gantinya, hal ini dilakukan oleh CommandBufferConversionPass yang berjalan di ThunkSequence itu sendiri.

Pass mengidentifikasi sub-urutan berdekatan dari thunk yang kompatibel (misalnya, serangkaian KernelThunk dan GemmThunk). Kemudian, mengganti sub-urutan yang ditemukan dengan satu CommandBufferThunk. Thunk baru merangkum logika thunk asli sebagai daftar objek CommandBufferCmd ringan. Saat CommandBufferThunk dieksekusi untuk pertama kalinya di aliran GPU tertentu, CommandBufferThunk akan "merekam" urutan perintahnya ke dalam buffer perintah hardware. Pada semua eksekusi berikutnya, perintah ini hanya mengeluarkan satu perintah ke GPU untuk "memutar ulang" urutan yang direkam. Tindakan ini menghindari overhead CPU saat meluncurkan setiap kernel secara terpisah.

Dapat dieksekusi

Produk akhir pipeline kompilasi XLA adalah Dapat Dieksekusi mandiri dan khusus platform. Objek ini merangkum semua informasi yang diperlukan untuk menjalankan program yang dikompilasi di perangkat target, seperti CPU atau GPU. Ini adalah jembatan antara compiler dan runtime. Runtime modern seperti PJRT menggunakan abstraksi tingkat yang sedikit lebih tinggi (lihat PjRtExecutable), tetapi pada akhirnya akan membungkus executable khusus backend.

Executable berisi beberapa informasi penting yang dihasilkan selama kompilasi. Meskipun konten persisnya bervariasi menurut backend, umumnya mencakup:

  • Kode yang Dikompilasi: Ini adalah kode mesin tingkat rendah yang akan berjalan di perangkat. Untuk CPU, ini biasanya berupa satu atau beberapa file objek. Untuk GPU, ini adalah kode perangkat yang dikompilasi dalam format PTX atau HSACO, yang dimuat ke GPU saat runtime.

  • Rencana Eksekusi (ThunkSequence): Inti dari logika runtime. Ini adalah urutan linear objek Thunk. Setiap thunk mewakili satu unit kerja, seperti meluncurkan kernel, memanggil fungsi library (misalnya, cuBLAS), atau menangani alur kontrol. Runtime menjalankan program dengan melakukan iterasi melalui urutan ini.

  • Tata Letak Memori (BufferAssignment): Metadata penting ini, yang dihasilkan oleh BufferAssigner, menjelaskan tata letak memori lengkap untuk komputasi. Hal ini menentukan ukuran setiap buffer dan cara memori dialokasikan dan digunakan kembali untuk parameter, output, dan nilai sementara. Runtime menggunakan ini untuk mengalokasikan memori perangkat dan meneruskan pointer yang benar ke setiap thunk.

  • (opsional) Modul HLO: Untuk proses debug dan pembuatan profil, file yang dapat dieksekusi sering kali mempertahankan referensi ke HloModule akhir yang dioptimalkan yang dikompilasi dari file tersebut.

Pembuatan file yang dapat dieksekusi akhir diatur oleh compiler untuk setiap backend tertentu. Metode RunBackend dari penerapan Compiler adalah langkah terakhir dalam proses kompilasi, yang memaketkan semua artefak yang dikompilasi ke dalam objek Executable. GpuCompiler dan CpuCompiler masing-masing menargetkan GPU dan CPU.

Saat pengguna memanggil Execute... pada file yang dapat dieksekusi, runtime menggunakan BufferAssignment untuk mengalokasikan memori, lalu memanggil ThunkSequence untuk meluncurkan operasi di perangkat menggunakan kode yang dikompilasi.