Proces programowania XLA zwykle koncentruje się na podczerwieni HLO, która reprezentuje izolowane obliczenia funkcjonalne przekazywane kompilatorowi. XLA zawiera wiele narzędzi wiersza poleceń (opisanych poniżej), które wykorzystują HLO i uruchamiają je lub zapewniają pośredni etap kompilacji. Korzystanie z tych narzędzi jest nieocenione w przypadku szybkiego, compile->modify->runcyklu iteracji, ponieważ HLO jest zarówno wizualizacja, jak i możliwość hakowania. Jej iteracyjne zmiany i uruchamianie często jest najszybszym sposobem zrozumienia i naprawienia wydajności lub zachowania XLA.
Najłatwiejszym sposobem uzyskania HLO dla programu kompilowanego z użyciem XLA jest zwykle użycie zmiennej środowiskowej XLA_FLAGS:
XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
który przechowuje w określonym folderze wszystkie pliki HLO przed optymalizacją wraz z wieloma innymi przydatnymi artefaktami.
Uruchomiono fragmenty kodu HLO: run_hlo_module
Narzędzie run_hlo_module działa na podstawie wstępnej optymalizacji HLO i domyślnie łączy w pakiety kompilację, uruchamianie i porównanie z implementacją interpretatora plików referencyjnych. Na przykład zwykłe wywołanie pliku wejściowe computation.hlo w GPU NVIDIA i sprawdzanie jego poprawności to:
run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Podobnie jak w przypadku innych narzędzi, możesz użyć --help, aby zobaczyć pełną listę opcji.
Uruchamianie fragmentów kodu HLO z obsługą SPMD: multihost_hlo_runner
Narzędzie do uruchamiania HLO z wieloma hostami jest narzędziem do bardzo podobnych, ale obsługuje SPMD, w tym komunikację między hostami. Typowe wywołanie wygląda tak:
hlo_runner_main --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo
Aktywne karnety/etapy kompilacji HLO: hlo-opt
Podczas debugowania lub analizowania działania kompilatora przydaje się często uzyskanie rozszerzenia dla określonego sprzętu w konkretnym punkcie potoku (np. HLO, zoptymalizowanym HLO, TritonIR lub LLVM) dla danego (stabilnego) wejścia HLO.
hlo-opt obsługuje wiele etapów wyjściowych: PTX, HLO po optymalizacji, IR LLVM przed optymalizacją, a także TritonIR. Dokładny zestaw obsługiwanych etapów zależy od platformy (np. PTX jest specyficzny dla platformy NVIDIA) i można go wyświetlić przy użyciu polecenia --list-stages:
$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx
Po wybraniu etapu użytkownik może zapisać w danym strumieniu wynik konwersji dla danej platformy:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
, który wydrukuje zrzut w postaci stdout (lub do konkretnego pliku, jeśli określono -o).
Wykorzystanie bez urządzeń
W większości kompilacji nie trzeba korzystać z GPU. Jeśli określisz w wierszu poleceń specyfikację GPU, uzyskasz na przykład dane wyjściowe PTX bez dostępu do akceleratora:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_80.txtpb input.hlo
Specyfikacje popularnych procesorów GPU są przesyłane przez kompilator. Przesłany plik to serializacja ciągów 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 wymagane jest autodostrajanie, w przypadku kompilacji bez urządzeń mogą wystąpić problemy. Na szczęście możemy je też podać w wierszu poleceń:
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
Plik automatycznego dostrajania to ciąg tekstowy autotune_results.proto. Przykład powinien wyglądać tak:
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
}
}
}
Bazę danych dostrajania można zserializować za pomocą parametru XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>
Uruchamianie pojedynczego przepustki kompilatora
Obsługiwane są też flagi z XLA_FLAGS, więc narzędzia można używać do testowania uruchamiania pojedynczego karnetu:
hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo