Menggunakan alat XLA

Alur kerja pengembangan XLA biasanya berpusat pada IR HLO, yang merepresentasikan fungsi terisolasi komputasi yang diberikan ke kompilator. XLA dilengkapi dengan banyak alat baris perintah (dijelaskan di bawah) yang menggunakan HLO dan menjalankannya, atau memberikan tahap kompilasi perantara. Menggunakan alat-alat tersebut sangat sangat berharga untuk mempercepat compile->modify->run siklus iterasi, karena HLO dapat divisualisasikan dan diretas, dan secara berulang mengubah serta menjalankannya sering kali menjadi cara tercepat untuk memahami dan memperbaiki kinerja atau perilaku XLA.

Cara termudah untuk mendapatkan HLO untuk program yang dikompilasi dengan XLA adalah biasanya 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, berjalan, dan perbandingan dengan penafsir referensi terlepas dari implementasi layanan. Misalnya, pemanggilan biasa untuk menjalankan file input computation.hlo pada GPU NVIDIA dan memeriksa ketepatannya adalah:

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

Seperti alat lainnya, --help dapat digunakan untuk mendapatkan daftar opsi lengkap.

Menjalankan cuplikan HLO dengan dukungan SPMD: multihost_hlo_runner

Multihost HLO Runner adalah alat yang sangat mirip, dengan peringatan yang didukungnya SPMD, termasuk komunikasi lintas {i>host<i}. Lihat Runner HLO Multi-Host untuk mengetahui detailnya.

Putar ulang multi-HLO

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

$ hlo_runner_main /dump/*before_optimizations*

Menjalankan pass/tahap kompilasi HLO: hlo-opt

Saat melakukan debug atau memahami cara kerja compiler, hal ini sering kali berguna untuk mendapatkan ekspansi bagi perangkat keras tertentu pada titik tertentu dalam pipeline (baik itu HLO, HLO yang dioptimalkan, TritonIR, atau LLVM), untuk HLO tertentu (Stabil) input teks.

hlo-opt mendukung beberapa tahap output: baik itu PTX, HLO setelah pengoptimalan, LLVM IR sebelum pengoptimalan, atau TritonIR. Rangkaian tahap persis yang didukung bergantung pada platform (misalnya PTX adalah 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 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 yang bisa kita dapatkan, 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 bersama compiler, dan file yang diberikan akan 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 baris perintah:

$ 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 {i>autotuning<i} dapat diserialisasi menggunakan XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>

Menjalankan Satu Compiler Pass

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

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