Menggunakan alat XLA

Alur kerja pengembangan XLA biasanya berpusat pada IR HLO, 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 perantara. Penggunaan alat tersebut sangat berharga untuk siklus iterasi compile->modify->run yang cepat, karena HLO dapat divisualisasi dan di-hack, serta mengubah dan menjalankannya secara iteratif sering kali merupakan cara tercepat untuk memahami dan memperbaiki performa atau perilaku XLA.

Cara termudah untuk mendapatkan HLO untuk program yang dikompilasi dengan XLA biasanya adalah menggunakan variabel lingkungan XLA_FLAGS:

$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point

yang menyimpan semua file HLO sebelum pengoptimalan di folder yang ditentukan, bersama dengan banyak artefak berguna lainnya.

Menjalankan cuplikan HLO: run_hlo_module

Alat run_hlo_module beroperasi pada HLO pra-pengoptimalan, dan secara default menggabungkan kompilasi, pengoperasian, dan perbandingan dengan implementasi interpretor referensi. Misalnya, pemanggilan biasa untuk menjalankan file input computation.hlo di GPU NVIDIA dan memeriksa kebenarannya adalah:

$ run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo

Seperti semua alat, --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 pengecualian bahwa alat ini mendukung SPMD, termasuk komunikasi lintas host. Lihat Multi-Host HLO Runner untuk mengetahui detailnya.

Replay multi-HLO

Pemanggilan dengan beberapa modul didukung untuk run_hlo_module dan hlo_runner_main, yang sering kali mudah untuk memutar ulang semua modul dalam direktori dump:

$ hlo_runner_main /dump/*before_optimizations*

Menjalankan tahap/pass kompilasi HLO: hlo-opt

Saat men-debug atau memahami cara kerja compiler, sering kali berguna untuk mendapatkan perluasan untuk hardware tertentu pada titik tertentu dalam pipeline (baik HLO, HLO yang dioptimalkan, TritonIR, maupun LLVM), untuk input HLO (Stabil) tertentu.

hlo-opt mendukung beberapa tahap output: baik PTX, HLO setelah pengoptimalan, LLVM IR sebelum pengoptimalan, atau TritonIR. Kumpulan tahap yang didukung tergantung pada platform (misalnya, PTX khusus 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 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 spec GPU di 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_pcie_80.txtpb input.hlo

Spesifikasi untuk GPU populer dikirimkan dengan compiler, dan file yang disediakan 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 mungkin mengalami masalah jika penyesuaian otomatis diperlukan. Untungnya, kita juga dapat menyediakannya di command line:

$ hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=gpu_specs/a100_pcie_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: 3
results {
  device: "CUDA: 8.0, Cores: 108, GPU clock: 1.41 GHz, Memory bandwidth: 1555 GB/s, L2 cache: 40 MB"
  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 autotuning dapat diserialisasi menggunakan XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>

Menjalankan Satu Pass Compiler

Flag dari XLA_FLAGS juga didukung, sehingga alat ini dapat digunakan untuk menguji proses menjalankan satu kartu:

$ hlo-opt --platform=CUDA --stage=hlo --passes=algebraic_simplifer input.hlo