Der XLA-Entwicklungsablauf konzentriert sich in der Regel auf die HLO-IR, die eine isolierte funktionale Berechnung darstellt, die dem Compiler übergeben wird. XLA bietet mehrere Befehlszeilentools (siehe unten), die HLO verwenden und entweder ausführen oder eine Zwischenkompilierungsphase bereitstellen. Die Verwendung solcher Tools ist für einen schnellen compile->modify->run-Iterationszyklus von unschätzbarem Wert, da HLO sowohl visualisierbar als auch hackbar ist. Die iterative Änderung und Ausführung ist oft die schnellste Methode, um die Leistung oder das Verhalten von XLA zu verstehen und zu korrigieren.
Am einfachsten erhalten Sie die HLO für ein Programm, das mit XLA kompiliert wird, mit der Umgebungsvariablen XLA_FLAGS:
$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
Hier werden alle HLO-Dateien vor der Optimierung zusammen mit vielen anderen nützlichen Artefakten im angegebenen Ordner gespeichert.
Ausgeführte HLO-Snippets: run_hlo_module
Das Tool run_hlo_module arbeitet mit HLO vor der Optimierung und bündelt standardmäßig die Kompilierung, Ausführung und den Vergleich mit der Referenz-Interpreterimplementierung. So wird beispielsweise eine Eingabedatei computation.hlo auf einer NVIDIA-GPU ausgeführt und auf Richtigkeit geprüft:
$ run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Wie bei allen Tools können Sie mit --help eine vollständige Liste der Optionen aufrufen.
HLO-Snippets mit SPMD-Unterstützung ausführen: multihost_hlo_runner
Der Multihost-HLO-Runner ist ein sehr ähnliches Tool, mit dem Unterschied, dass er SPMD unterstützt, einschließlich der geräteübergreifenden Kommunikation. Weitere Informationen finden Sie unter Multi-Host HLO Runner.
Multi-HLO-Wiedergabe
Die Ausführung mit mehreren Modulen wird sowohl für run_hlo_module als auch für hlo_runner_main unterstützt. Das ist oft praktisch, um alle Module in einem Dumpverzeichnis noch einmal abzuspielen:
$ hlo_runner_main /dump/*before_optimizations*
Ausgeführte Durchläufe/Phasen der HLO-Kompilierung: hlo-opt
Beim Debuggen oder Verstehen der Funktionsweise des Compilers ist es oft nützlich, die Erweiterung für eine bestimmte Hardware an einem bestimmten Punkt in der Pipeline (HLO, optimierte HLO, TritonIR oder LLVM) für einen bestimmten (stabilen) HLO-Eingang abzurufen.
hlo-opt unterstützt mehrere Ausgabephasen: PTX, HLO nach Optimierungen, LLVM IR vor Optimierungen oder TritonIR. Welche genauen Phasen unterstützt werden, hängt von der Plattform ab (z. B. ist PTX NVIDIA-spezifisch). Sie können sie mit dem Befehl „–list-stages“ aufrufen:
$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx
Nach der Auswahl einer Phase kann der Nutzer das Ergebnis der Conversion für eine bestimmte Plattform in einen bestimmten Stream schreiben:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
Dadurch wird der Dump auf stdout ausgegeben (oder in eine angegebene Datei, wenn -o angegeben wurde).
Nutzung ohne Gerät
Für den Großteil der Kompilierung ist kein Zugriff auf eine GPU erforderlich. Durch Angabe einer GPU-Spezifikation in der Befehlszeile können wir beispielsweise PTX-Ausgabe ohne Zugriff auf einen Beschleuniger erhalten:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_pcie_80.txtpb input.hlo
Die Spezifikationen für gängige GPUs sind im Compiler enthalten und die bereitgestellte Datei ist eine String-Serialisierung von 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"
Bei der gerätelosen Kompilierung können Probleme auftreten, wenn eine automatische Optimierung erforderlich ist. Glücklicherweise können wir diese auch über die Befehlszeile bereitstellen:
$ 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
Die Autotune-Datei ist eine Textserialisierung von autotune_results.proto. Ein Beispiel sieht so aus:
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
}
}
}
Die Datenbank für die automatische Optimierung kann mithilfe von XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt> serialisiert werden.
Einen einzelnen Compilerdurchlauf ausführen
Die Flags von XLA_FLAGS werden ebenfalls unterstützt, sodass das Tool zum Testen der Ausführung eines einzelnen Durchlaufs verwendet werden kann:
$ hlo-opt --platform=CUDA --stage=hlo --passes=algebraic_simplifer input.hlo