Perkakas XLA

Alur kerja pengembangan XLA biasanya berpusat pada IR HLO, yang merepresentasikan 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 divisualisasikan dan dimodifikasi, 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 dengan 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, beserta banyak artefak berguna lainnya.

[run_hlo_module] Menjalankan Modul HLO

bazel run //xla/tools:run_hlo_module -- [flags] <filename>

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

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

Menjalankan Beberapa Modul HLO

Pemanggilan dengan beberapa modul HLO didukung untuk run_hlo_module. Untuk menjalankan semua modul hlo dari direktori:

bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations*

[multihost_hlo_runner] Menjalankan Modul HLO dengan Dukungan SPMD

# Note: Binary name is `hlo_runner_main`.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] <filename>

Peluncur HLO multihost adalah alat yang sangat mirip, dengan catatan bahwa alat ini mendukung SPMD, termasuk komunikasi lintas host. Lihat Multi-Host HLO Runner untuk mengetahui detailnya.

Menjalankan Beberapa Modul HLO dengan Dukungan SPMD

Serupa dengan run_hlo_module, multihost_hlo_runner juga mendukung pemanggilan dengan beberapa modul.

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] /dump/*before_optimizations*

[hlo-opt] Kompilasi Modul HLO

bazel run //xla/tools:hlo-opt -- --platform=[gpu|cpu|...] [more flags] <filename>

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

hlo-opt mendukung beberapa tahap output: baik PTX, HLO setelah pengoptimalan, LLVM IR sebelum pengoptimalan, atau TritonIR. Kumpulan tahap yang didukung bergantung pada platform (misalnya, PTX khusus untuk NVIDIA), dan dapat dilihat menggunakan perintah --list-stages:

hlo-opt --platform=CUDA --list-stages
buffer-assignment
hlo
hlo-backend
html
llvm
llvm-after-optimizations
llvm-before-optimizations
ptx

Setelah memilih tahap, pengguna dapat menulis hasil konversi untuk platform tertentu ke aliran tertentu:

hlo-opt --platform=cpu --stage=hlo input.hlo

yang akan mencetak dump ke stdout (atau ke file tertentu jika -o ditentukan).

Kompilasi Tanpa Perangkat untuk GPU

Kompilasi tanpa perangkat tidak memerlukan akses ke GPU. Kompilasi Tanpa Perangkat menyediakan cara untuk menentukan spesifikasi GPU di command line (--xla_gpu_target_config_filename) untuk tahap yang memerlukan akses ke GPU, sehingga tidak memerlukan perangkat GPU.

Contoh: Output PTX tanpa akses ke perangkat GPU:

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=/xla/tools/hlo_opt/gpu_specs/a100_pcie_80.txtpb input.hlo

Spesifikasi untuk GPU populer dikirimkan bersama 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"

Spesifikasi GPU lainnya dapat dilihat di /xla/tools/hlo_opt/gpu_specs

Penyesuaian otomatis

Terkadang kompilasi dapat melibatkan penyetelan otomatis berdasarkan kompilasi --stage. Agar kompilasi tanpa perangkat berfungsi, pengguna harus
menonaktifkan penyetelan otomatis dengan --xla_gpu_autotune_level=0
atau
memuat hasil penyetelan otomatis yang sudah ada dengan --xla_gpu_load_autotune_results_from=<filename> (diperoleh dengan --xla_gpu_dump_autotune_results_to=<filename>).

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 penyetelan otomatis 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 penyetelan otomatis dapat diserialisasi menggunakan XLA_FLAGS=--xla_gpu_dump_autotune_results_to=<myfile.pbtxt>

[hlo-opt] Pengembangan dan Proses Debugging Kartu HLO

# If you are working with hardware independent passes from the
# `xla/hlo/transforms/` directory, prefer light-weight version
# of the `hlo-opt` tool with fewer dependencies:

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

# Otherwise, for hardware independent and CPU, GPU passes use
# the same binary from "Compile HLO Modules" section above:

bazel run //xla/tools:hlo-opt -<- [flags>] filename

Alat hlo-opt memungkinkan eksekusi setiap lulus secara independen dari tahap kompilasi platform yang diberikan. Isolasi ini membantu menjalankan lulus dengan cepat pada modul hlo input dan menunjukkan akar penyebab kegagalan.

hlo-opt --passes=schedule-aware-collective-cse input.hlo

Alat hlo-opt juga mendukung DebugOptions XLA_FLAGS.

hlo-opt --passes=schedule-aware-collective-cse
--xla_gpu_experimental_collective_cse_distance_threshold=20 input.hlo

Gunakan opsi --list-passes untuk mendapatkan string nama kartu.

hlo-opt --list-passes

Pengguna dapat membuat pipeline kustom mereka sendiri dengan menentukan lebih dari satu opsi lulus ke --passes.

hlo-opt --passes=pass1,pass2,pass3 input.hlo

Membantu Pengembangan Kartu HLO Baru

  1. Pertama, tulis kartu izin lewat Anda.
  2. Daftarkan kartu baru ke registry kartu alat hlo-opt.

    RegisterPass<FooPass>(FooPassInputOptions)
    

    Berdasarkan jenis kartu, pilih salah satu lokasi berikut untuk pendaftaran:
    opt_lib.cc Kartu yang tidak bergantung pada hardware.
    cpu_opt.cc Penerusan khusus CPU.
    gpu_opt.cc Penerusan khusus GPU.
    compiled_opt.cc Lulus tes umum untuk CPU, GPU, XPU.
    Jangan lupa menambahkan dependensi build.

    Sertakan pendaftaran kartu sebagai bagian dari PR Anda(contoh) agar kartu tersedia untuk digunakan oleh semua pengguna hlo-opt.

  3. Bangun kembali alat hlo-opt, validasi pendaftaran kartu yang berhasil menggunakan opsi --list-passes, lalu gunakan opsi --passes untuk menjalankan kartu.

    $ hlo-opt --passes=foo-pass input.hlo
    
  4. Menulis pengujian unit untuk pass? Lihat https://openxla.org/xla/test_hlo_passes untuk mengetahui detail selengkapnya.

Pengukuran Waktu Proses Kartu

Untuk model besar, proses kompilasi penuh dapat memakan waktu hingga beberapa menit, sehingga sulit untuk mendeteksi regresi performa yang tidak terlalu terlihat. Sebaliknya, proses pengujian individual menggunakan hlo-opt memungkinkan pengukuran performa yang akurat dan deteksi yang mudah bahkan untuk peningkatan kecil dalam waktu eksekusi yang disebabkan oleh perubahan kode baru.

time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier
--xla_reduce_window_rewrite_base_length=128 input.hlo

[hlo-opt] Mengonversi Format Modul HLO

# Use the light weight version of the `hlo-opt` tool.

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

Konversi HLO Text -> HLO Proto

hlo-opt --emit-proto input.hlo

Konversi HLO Proto atau HLO Proto Binary -> HLO Text

hlo-opt input.pbtxt or input.pb

[ptx-opt] Compiler Modul LLVM ke PTX

Alat ini akan menjalankan pipeline pengoptimalan LLVMIR, lalu memanggil CompileToPtx.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 <filename>

Alat ini juga dapat membuang LLVMIR setelah setiap jalur.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to=<path> --xla_gpu_dump_llvmir <filename>