Korzystanie z narzędzi XLA

Proces tworzenia XLA zazwyczaj koncentruje się HLO – podczerwień – izolowana, funkcjonalna dane przekazywane kompilatorowi. XLA zawiera wiele narzędzi wiersza poleceń (opisane poniżej), które wykorzystują HLO i uruchamiają je lub udostępniają na pośrednim etapie kompilacji. Korzystanie z takich narzędzi jest bezcenne, jeśli chodzi o compile->modify->run, ponieważ HLO jest zarówno widoczny, łatwych do hakowania, iteracyjne zmiany i uruchomienie to często najszybszy sposób zrozumieć i naprawić skuteczność lub zachowanie XLA.

Najłatwiejszym sposobem na uzyskanie HLO dla programu skompilowanego z XLA jest zwykle używaj zmiennej środowiskowej XLA_FLAGS:

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

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

Uruchomione fragmenty kodu HLO: run_hlo_module

Narzędzie run_hlo_module działa na zasadzie HLO przed optymalizacją i domyślnie kompilowanie, uruchamianie i porównywanie pakietów z interpreterem referencyjnym implementacji. Na przykład zwykłe wywołanie do uruchomienia pliku wejściowego computation.hlo na GPU NVIDIA. Aby sprawdzić jego poprawność:

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

Tak jak w przypadku innych narzędzi, możesz użyć narzędzia --help, aby uzyskać pełną listę opcji.

Uruchamianie fragmentów kodu HLO z obsługą SPMD: multihost_hlo_runner

Uruchamiający HLO dla wielu hostów jest bardzo podobnym narzędziem, jednak obsługuje ono SPMD, w tym komunikacja między hostami. Zobacz Uruchamiający HLO na wielu hostach, aby uzyskać szczegółowe informacje.

Ponowne odtwarzanie wielu HLO

Wywołanie z wieloma modułami jest obsługiwane zarówno przez run_hlo_module, jak i hlo_runner_main, co często pozwala odtworzyć wszystkie moduły w zrzucie katalogu:

$ hlo_runner_main /dump/*before_optimizations*

Uruchomione karty/etapy kompilacji HLO: hlo-opt

Przy debugowaniu lub analizowaniu działania kompilatora często przydaje się aby pobrać rozszerzenie dla konkretnego sprzętu w określonym momencie potoku (np. HLO, zoptymalizowany HLO, TritonIR lub LLVM) dla danego (stabilnego) HLO dane wejściowe.

hlo-opt obsługuje wiele etapów danych wyjściowych: PTX, HLO po optymalizacji, LLVM IR przed optymalizacją lub TritonIR. Dokładny zestaw obsługiwanych etapów zależy od platformy (np. pakiet PTX jest związany z rozwiązaniami NVIDIA) i można go wyświetlić za pomocą polecenie --list-stages:

$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx

Po wybraniu etapu użytkownik może zapisać wynik konwersji dla z konkretną platformą do danego strumienia:

$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm

co spowoduje wygenerowanie zrzutu na „stdout” (lub do konkretnego pliku, jeśli określono parametr -o).

Wykorzystanie bez urządzeń

Dostęp do GPU nie jest potrzebny w przypadku większości kompilacji. Dodatkowo musisz określić Specyfikacja GPU w wierszu poleceń, którą możemy uzyskać, np. Wyjście PTX bez dostępu do akcelerator:

$ hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_pcie_80.txtpb input.hlo

Specyfikacje popularnych układów GPU są dostarczane z kompilatorem, a przesłany plik to serializacja 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"

Jeśli jest wymagane automatyczne dostrajanie, kompilacja bez urządzeń może powodować problemy. Na szczęście możemy też podać je w wierszu poleceń:

$ 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 z automatycznym dostrajaniem to serializacja tekstu autotune_results.proto, ze przykład wyglądający tak:

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 zserializować za pomocą XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>

Uruchamianie pojedynczej karty kompilacji

Obsługiwane są również flagi z XLA_FLAGS, więc narzędzie może służyć do testowania w przypadku jednej karnetu:

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