Der XLA-Entwicklungsworkflow konzentriert sich normalerweise auf HLO IR, was eine isolierte funktionale Berechnung des Compilers darstellt. XLA wird mit mehreren Befehlszeilentools (siehe unten) geliefert, die HLO verarbeiten 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 visualisiert als auch hackbar ist. Die iterative Änderung und Ausführung ist oft der schnellste Weg, 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, in der Regel mit der Umgebungsvariable XLA_FLAGS:
XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
mit dem alle HLO-Dateien vor der Optimierung zusammen mit vielen anderen nützlichen Artefakten im angegebenen Ordner gespeichert werden.
HLO-Snippets werden ausgeführt: run_hlo_module
Das Tool run_hlo_module arbeitet mit HLO vor der Optimierung und bündelt standardmäßig Kompilierung, Ausführung und Vergleich mit der Implementierung des Referenzinterpreters. Der übliche Aufruf zum Ausführen der Eingabedatei computation.hlo auf einer NVIDIA-GPU und zum Prüfen der Richtigkeit lautet beispielsweise:
run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Wie bei allen Tools kann --help verwendet werden, um die vollständige Liste der Optionen abzurufen.
HLO-Snippets mit SPMD-Unterstützung ausführen: multihost_hlo_runner
Der Multihost-HLO-Runner ist ein sehr ähnliches Tool, mit dem Unterschied, dass es SPMD, einschließlich der Cross-Host-Kommunikation, unterstützt. Ein typischer Aufruf sieht so aus:
hlo_runner_main --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo
Pässe/Phasen der HLO-Kompilierung werden ausgeführt: hlo-opt
Wenn Sie Fehler beheben oder die Funktionsweise des Compilers verstehen möchten, ist es oft nützlich, die Erweiterung für eine bestimmte Hardware an einem bestimmten Punkt in der Pipeline (z. B. HLO, optimiert HLO, TritonIR oder LLVM) für eine bestimmte (stabile) HLO-Eingabe zu erhalten.
hlo-opt unterstützt mehrere Ausgabephasen: PTX, HLO nach Optimierungen, LLVM IR vor Optimierungen oder TritonIR. Die genaue Gruppe der unterstützten Phasen hängt von der Plattform ab (z.B. ist PTX NVIDIA-spezifisch) und können mit dem Befehl --list-stages abgerufen werden:
$ 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 in stdout (oder in einer bestimmten Datei, wenn -o angegeben wurde) ausgegeben.
Gerätelose Nutzung
Für die meisten Kompilierungen ist kein Zugriff auf eine GPU erforderlich. Wenn Sie in der Befehlszeile eine GPU-Spezifikation angeben, erhalten Sie z. B. die PTX-Ausgabe ohne Zugriff auf einen Beschleuniger:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_80.txtpb input.hlo
Spezifikationen für gängige GPUs werden mit dem Compiler geliefert. 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 einer gerätelosen Kompilierung können Probleme auftreten, wenn Autotuning erforderlich ist. Glücklicherweise können wir diese auch in der Befehlszeile angeben:
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
Die Autotune-Datei ist die Textserialisierung von autotune_results.proto. Beispiel:
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
}
}
}
Die Autotuning-Datenbank kann mit XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt> serialisiert werden
Einzelnen Compiler-Pass ausführen
Die Flags von XLA_FLAGS werden ebenfalls unterstützt, sodass das Tool zum Testen der Ausführung einer einzelnen Karte / eines einzelnen Ausweises verwendet werden kann:
hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo