Narzędzia XLA

Proces tworzenia XLA zwykle koncentruje się na HLO IR, który reprezentuje izolowane obliczenia funkcjonalne przekazywane do kompilatora. XLA zawiera kilka narzędzi wiersza poleceń (opisanych poniżej), które wykorzystują HLO i uruchamiają je lub zapewniają pośredni etap kompilacji. Korzystanie z takich narzędzi jest nieocenione w przypadku szybkiego cyklu iteracji, ponieważ HLO można wizualizować i modyfikować, a iteracyjne zmienianie i uruchamianie tego języka jest często najszybszym sposobem na zrozumienie i naprawienie wydajności lub zachowania XLA.compile->modify->run

Najprostszym sposobem uzyskania HLO dla programu kompilowanego za pomocą XLA jest zwykle użycie zmiennej środowiskowej XLA_FLAGS:

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

które przechowuje wszystkie pliki HLO przed optymalizacją w określonym folderze wraz z wieloma innymi przydatnymi artefaktami.

[run_hlo_module] Uruchamianie modułów HLO

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

Narzędzie run_hlo_module działa na podstawie HLO przed optymalizacją i domyślnie łączy kompilację, uruchamianie i porównywanie z implementacją interpretera referencyjnego. Na przykład zwykłe wywołanie, aby uruchomić plik wejściowycomputation.hlo na procesorze GPU NVIDIA i sprawdzić jego poprawność, to:

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

Uruchamianie wielu modułów HLO

Wywoływanie z wieloma modułami HLO jest obsługiwane w przypadku run_hlo_module. Aby uruchomić wszystkie moduły HLO z katalogu:

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

[multihost_hlo_runner] Uruchamianie modułów HLO z obsługą SPMD

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

Narzędzie do uruchamiania HLO na wielu hostach jest bardzo podobne, z tym wyjątkiem, że obsługuje SPMD, w tym komunikację między hostami. Więcej informacji znajdziesz w artykule Multi-Host HLO Runner.

Uruchamianie wielu modułów HLO z obsługą SPMD

Podobnie jak w przypadku run_hlo_module, multihost_hlo_runner obsługuje też wywoływanie z użyciem wielu modułów.

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

[hlo-opt] Kompilacja modułu HLO

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

Podczas debugowania lub analizowania działania kompilatora często przydatne jest uzyskanie rozwinięcia dla konkretnego sprzętu w określonym punkcie potoku (HLO, zoptymalizowany HLO, TritonIR lub LLVM) dla danego wejścia HLO lub StableHLO.

hlo-opt obsługuje wiele etapów wyjściowych: PTX, HLO po optymalizacji, LLVM IR przed optymalizacją lub TritonIR. Dokładny zestaw obsługiwanych etapów zależy od platformy (np. PTX jest specyficzny dla NVIDIA) i można go sprawdzić za pomocą polecenia --list-stages:

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

Po wybraniu etapu użytkownik może zapisać wynik konwersji dla danej platformy w danym strumieniu:

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

które spowoduje wydrukowanie zrzutu na standardowe wyjście (lub do danego pliku, jeśli określono -o).

Kompilacja bez urządzenia na potrzeby GPU

Kompilacja bez urządzenia nie wymaga dostępu do GPU. Kompilacja bez urządzenia umożliwia określenie specyfikacji GPU w wierszu poleceń (--xla_gpu_target_config_filename) na etapach, na których wymagany jest dostęp do GPU, co eliminuje potrzebę korzystania z urządzenia GPU.

Przykład: dane wyjściowe PTX bez dostępu do urządzenia GPU:

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

Specyfikacje popularnych procesorów GPU są dostarczane z kompilatorem, a podany plik jest serializacją ciągu znaków 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"

Więcej informacji o specyfikacjach GPU znajdziesz na stronie /xla/tools/hlo_opt/gpu_specs

Automatyczne dostrajanie

Czasami kompilacja może obejmować automatyczne dostrajanie na podstawie kompilacji --stage. Aby kompilacja bez urządzenia działała, użytkownik musi
wyłączyć automatyczne dostrajanie za pomocą --xla_gpu_autotune_level=0
lub
wczytać istniejące wyniki automatycznego dostrajania za pomocą --xla_gpu_load_autotune_results_from=<filename> (uzyskane za pomocą --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

Plik automatycznego dostrajania to serializacja tekstowa autotune_results.proto. Przykład:

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
    }
  }
}

Bazę danych automatycznego dostrajania można serializować za pomocą funkcji XLA_FLAGS=--xla_gpu_dump_autotune_results_to=<myfile.pbtxt>

[hlo-opt] Tworzenie i debugowanie przepustek 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>

Narzędzie hlo-opt umożliwia wykonanie poszczególnych etapów niezależnie od etapów kompilacji na danej platformie. Ta izolacja pomaga szybko uruchamiać przebiegi w module wejściowym HLO i wskazywać główną przyczynę błędów.

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

Narzędzie hlo-opt obsługuje też DebugOptions XLA_FLAGS.

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

Użyj opcji--list-passes, aby uzyskać ciąg nazwy karty.

hlo-opt --list-passes

Użytkownicy mogą tworzyć własne niestandardowe potoki, określając więcej niż 1 przekazywanie w opcji --passes.

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

Pomoc w opracowywaniu nowego karnetu HLO

  1. Najpierw napisz pozwolenie.
  2. Zarejestruj nową kartę w rejestrze kart narzędzi hlo-opt.

    RegisterPass<FooPass>(FooPassInputOptions)
    

    W zależności od rodzaju karty wybierz jedną z tych lokalizacji rejestracji:
    opt_lib.cc Karty niezależne od sprzętu.
    cpu_opt.cc Dokumenty dotyczące konkretnych procesorów.
    gpu_opt.cc Dokumenty dotyczące konkretnych procesorów graficznych.
    compiled_opt.cc Testy wspólne dla procesora, GPU i XPU.
    Nie zapomnij dodać zależności kompilacji.

    Uwzględnij rejestrację karty w swoim PR(przykład), aby karta była dostępna dla wszystkich użytkowników hlo-opt.

  3. Odbuduj narzędzie hlo-opt, sprawdź, czy rejestracja karty została przeprowadzona prawidłowo, używając opcji --list-passes, a następnie uruchom kartę za pomocą opcji --passes.

    $ hlo-opt --passes=foo-pass input.hlo
    
  4. Pisanie testów jednostkowych dla przekazywania? Więcej informacji znajdziesz na stronie https://openxla.org/xla/test_hlo_passes.

Pomiar czasu działania karty

W przypadku dużych modeli pełna kompilacja może trwać kilka minut, co utrudnia wykrywanie subtelnych regresji wydajności. Natomiast pojedyncze przebiegi testów z użyciem hlo-opt umożliwiają precyzyjny pomiar wydajności i łatwe wykrywanie nawet niewielkich wzrostów czasu wykonania spowodowanych nowymi zmianami w kodzie.

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

[hlo-opt] Konwertowanie formatów modułów HLO

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

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

Przekonwertuj HLO Text –> HLO Proto

hlo-opt --emit-proto input.hlo

Przekształć HLO Proto lub HLO Proto Binary -> HLO Text

hlo-opt input.pbtxt or input.pb

[ptx-opt] Kompilator LLVM Module do PTX

Narzędzie uruchomi potok optymalizacji LLVMIR, a następnie wywoła funkcję CompileToPtx.

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

Narzędzie może też zrzucać LLVMIR po każdej ścieżce.

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