Alur kerja pengembangan XLA biasanya berpusat pada
HLO IR, yang mewakili komputasi fungsional
terisolasi yang diberikan ke compiler. XLA dilengkapi dengan beberapa alat command line
(dijelaskan di bawah) yang menggunakan HLO dan menjalankannya, atau menyediakan
tahap kompilasi menengah. Penggunaan alat tersebut sangat berharga untuk siklus iterasi compile->modify->run
yang cepat, karena HLO dapat divisualisasikan dan dapat diretas, serta mengubah dan menjalankannya secara iteratif sering kali menjadi cara tercepat untuk memahami dan memperbaiki performa atau perilaku XLA.
Cara termudah untuk mendapatkan HLO untuk program yang dikompilasi dengan XLA adalah
dengan menggunakan variabel lingkungan XLA_FLAGS
:
XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
yang menyimpan semua file HLO sebelum pengoptimalan dalam folder yang ditentukan, beserta banyak artefak berguna lainnya.
Menjalankan cuplikan HLO: run_hlo_module
Alat run_hlo_module
beroperasi pada HLO pra-pengoptimalan, dan secara default
mengelompokkan paket, berjalan, dan dibandingkan dengan penerapan penafsir
referensi. Misalnya, pemanggilan umum untuk menjalankan file input
computation.hlo
pada GPU NVIDIA dan memeriksa ketepatannya adalah:
run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Seperti semua alat lainnya, --help
dapat digunakan untuk mendapatkan daftar lengkap opsi.
Menjalankan cuplikan HLO dengan dukungan SPMD: multihost_hlo_runner
Runner HLO multihost adalah alat yang sangat mirip, dengan peringatan bahwa runner HLO tersebut mendukung SPMD, termasuk komunikasi lintas host. Pemanggilan umum terlihat seperti:
hlo_runner_main --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo
Tahap/tahap berjalan pada kompilasi HLO: hlo-opt
Saat men-debug atau memahami cara kerja compiler, sebaiknya dapatkan ekspansi untuk hardware tertentu pada titik tertentu dalam pipeline (baik itu HLO, HLO, TritonIR, atau LLVM yang dioptimalkan), untuk input HLO (Stabil) tertentu.
hlo-opt
mendukung beberapa tahap output: baik itu PTX, HLO setelah pengoptimalan,
LLVM IR sebelum pengoptimalan, maupun TritonIR. Serangkaian stage persis yang didukung
bergantung pada platform (misalnya, PTX khusus untuk NVIDIA), dan dapat dilihat menggunakan
perintah --list-stages:
$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx
Setelah memilih tahap, pengguna dapat menulis hasil konversi untuk platform tertentu ke aliran data tertentu:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
yang akan mencetak dump ke stdout (atau ke file tertentu jika -o
ditentukan).
Penggunaan Tanpa Perangkat
Akses ke GPU tidak diperlukan untuk sebagian besar kompilasi, dan dengan menentukan spesifikasi GPU pada command line, kita bisa mendapatkan, misalnya, output PTX tanpa akses ke akselerator:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_80.txtpb input.hlo
Spesifikasi untuk GPU populer dikirimkan dengan compiler, dan file yang diberikan adalah serialisasi string device_description.proto
:
gpu_device_info {
cuda_compute_capability {
major: 8
minor: 0
}
threads_per_block_limit: 1024
threads_per_warp: 32
shared_memory_per_block: 127152
shared_memory_per_core: 65536
threads_per_core_limit: 2048
core_count: 6192
fpus_per_core: 64
block_dim_limit_x: 2147483647
block_dim_limit_y: 65535
block_dim_limit_z: 65535
memory_bandwidth: 2039000000000
l2_cache_size: 4194304
clock_rate_ghz: 1.1105
device_memory_size: 79050250240
}
platform_name: "CUDA"
Kompilasi tanpa perangkat dapat mengalami masalah jika penyesuaian otomatis diperlukan. Untungnya, kita juga dapat memberikannya pada baris perintah:
hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=gpu_specs/a100_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo
File autotune adalah serialisasi teks autotune_results.proto
, dengan
contoh yang terlihat seperti:
version: 2
results {
device: "sm_8.0 with 42331013120B RAM, 108 cores, 1410000KHz clock, 1215000KHz mem clock, 41943040B L2$"
hlo: "{\n tmp_0 = f16[1,16,17,3]{3,2,1,0} parameter(0)\n tmp_1 = f16[16,51]{1,0} bitcast(f16[1,16,17,3]{3,2,1,0} tmp_0)\n tmp_2 = s8[16,17,3]{2,1,0} parameter(1)\n tmp_3 = s8[51,16]{0,1} bitcast(s8[16,17,3]{2,1,0} tmp_2)\n tmp_4 = f16[51,16]{0,1} convert(s8[51,16]{0,1} tmp_3)\n tmp_5 = f16[16,16]{1,0} dot(f16[16,51]{1,0} tmp_1, f16[51,16]{0,1} tmp_4), lhs_contracting_dims={1}, rhs_contracting_dims={0}\n ROOT tmp_6 = f16[1,16,16]{2,1,0} bitcast(f16[16,16]{1,0} tmp_5)\n}"
result {
run_time {
nanos: 31744
}
triton {
block_m: 32
block_n: 32
block_k: 32
split_k: 1
num_stages: 1
num_warps: 4
}
}
}
Database penyesuaian otomatis dapat diserialisasi menggunakan
XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>
Menjalankan Satu Kartu Compiler
Tanda dari XLA_FLAGS
juga didukung, sehingga alat ini dapat digunakan untuk menguji
cara menjalankan satu penerusan:
hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo