Menggunakan alat XLA

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