Der XLA-Entwicklungs-Workflow dreht sich in der Regel um die HLO-IR, die eine isolierte funktionale Berechnung darstellt, die dem Compiler übergeben wird. XLA enthält 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 visualisiert als auch manipuliert werden kann. Das iterative Ändern und Ausführen ist oft der schnellste Weg, um die Leistung oder das Verhalten von XLA zu verstehen und zu beheben.
Am einfachsten erhalten Sie das HLO für ein Programm, das mit XLA kompiliert wird, in der Regel über die Umgebungsvariable XLA_FLAGS:
$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
Dadurch werden alle HLO-Dateien vor der Optimierung im angegebenen Ordner gespeichert, zusammen mit vielen anderen nützlichen Artefakten.
[run_hlo_module] HLO-Module ausführen
bazel run //xla/tools:run_hlo_module -- [flags] <filename>
Das Tool run_hlo_module arbeitet mit HLO vor der Optimierung und umfasst standardmäßig die Kompilierung, Ausführung und den Vergleich mit der Referenzimplementierung des Interpreters. Der übliche Aufruf zum Ausführen einer 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
Mehrere HLO-Module ausführen
Der Aufruf mit mehreren HLO-Modulen wird für run_hlo_module unterstützt. So führen Sie alle HLO-Module aus einem Verzeichnis aus:
bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations*
[multihost_hlo_runner] HLO-Module mit SPMD-Unterstützung ausführen
# Note: Binary name is `hlo_runner_main`.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] <filename>
Der Multihost-HLO-Runner ist ein sehr ähnliches Tool, mit dem Unterschied, dass er SPMD einschließlich der hostübergreifenden Kommunikation unterstützt. Weitere Informationen finden Sie unter Multi-Host HLO Runner.
Mehrere HLO-Module mit SPMD-Unterstützung ausführen
Ähnlich wie bei run_hlo_module wird bei multihost_hlo_runner auch der Aufruf mit mehreren Modulen unterstützt.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] /dump/*before_optimizations*
[hlo-opt] HLO-Modul kompilieren
bazel run //xla/tools:hlo-opt -- --platform=[gpu|cpu|...] [more flags] <filename>
Beim Debuggen oder beim Versuch, die Funktionsweise des Compilers zu verstehen, ist es oft hilfreich, die Erweiterung für eine bestimmte Hardware an einem bestimmten Punkt in der Pipeline (HLO, optimiertes HLO, TritonIR oder LLVM) für eine bestimmte HLO- oder StableHLO-Eingabe zu erhalten.
hlo-opt unterstützt mehrere Ausgabestufen, z. B. PTX, HLO nach Optimierungen, LLVM IR vor Optimierungen oder TritonIR. Die genaue Anzahl der unterstützten Phasen hängt von der Plattform ab (z. B. ist PTX NVIDIA-spezifisch). Sie kann mit dem Befehl „--list-stages“ angezeigt werden:
hlo-opt --platform=CUDA --list-stages
buffer-assignment
hlo
hlo-backend
html
llvm
llvm-after-optimizations
llvm-before-optimizations
ptx
Nachdem der Nutzer eine Phase ausgewählt hat, kann er das Ergebnis der Konvertierung für eine bestimmte Plattform in einen bestimmten Stream schreiben:
hlo-opt --platform=cpu --stage=hlo input.hlo
Dadurch wird der Dump in stdout (oder in eine bestimmte Datei, wenn -o angegeben wurde) ausgegeben.
Gerätelose Kompilierung für GPU
Für die gerätelose Kompilierung ist kein Zugriff auf eine GPU erforderlich. Mit der gerätelosen Kompilierung können Sie die GPU-Spezifikation in der Befehlszeile (--xla_gpu_target_config_filename) für Phasen angeben, in denen Zugriff auf die GPU erforderlich ist. Dadurch ist kein GPU-Gerät mehr erforderlich.
Beispiel: PTX-Ausgabe ohne Zugriff auf ein GPU-Gerät:
hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=/xla/tools/hlo_opt/gpu_specs/a100_pcie_80.txtpb input.hlo
Spezifikationen für beliebte GPUs werden mit dem Compiler ausgeliefert. 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"
Weitere GPU-Spezifikationen finden Sie unter /xla/tools/hlo_opt/gpu_specs.
Automatische Abstimmung
Manchmal kann die Kompilierung Autotuning auf Grundlage einer Kompilierung --stage umfassen.
Damit die gerätelose Kompilierung funktioniert, muss der Nutzer entweder
Autotuning deaktivieren mit --xla_gpu_autotune_level=0
oder
vorhandene Autotuning-Ergebnisse laden mit
--xla_gpu_load_autotune_results_from=<filename> (erhalten mit
--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
Die Datei für die automatische Optimierung ist eine Textserialisierung von autotune_results.proto. Ein Beispiel:
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 Abstimmung kann mit XLA_FLAGS=--xla_gpu_dump_autotune_results_to=<myfile.pbtxt> serialisiert werden.
[hlo-opt] HLO-Pass-Entwicklung und ‑Debugging
# 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>
Mit dem hlo-opt-Tool können einzelne Durchläufe unabhängig von den angegebenen Phasen der Plattformkompilierung ausgeführt werden. Diese Isolation hilft, schnell Durchläufe für das HLO-Eingabemodul auszuführen und die Ursache von Fehlern zu ermitteln.
hlo-opt --passes=schedule-aware-collective-cse input.hlo
Das hlo-opt-Tool unterstützt auch DebugOptions XLA_FLAGS.
hlo-opt --passes=schedule-aware-collective-cse
--xla_gpu_experimental_collective_cse_distance_threshold=20 input.hlo
Verwenden Sie die Option--list-passes, um den Namen des Passes als String abzurufen.
hlo-opt --list-passes
Nutzer können eine eigene benutzerdefinierte Pipeline erstellen, indem sie mehr als einen Durchlauf für die Option --passes angeben.
hlo-opt --passes=pass1,pass2,pass3 input.hlo
Unterstützung bei der Entwicklung neuer HLO-Pässe
- Schreiben Sie zuerst Ihren Pass.
Registrieren Sie die neue Karte bzw. das neue Ticket in der
hlo-opt-Tool-Registry.RegisterPass<FooPass>(FooPassInputOptions)Wählen Sie je nach Art der Karte oder des Tickets einen der folgenden Registrierungsorte aus:
opt_lib.ccHardwareunabhängige Karten oder Tickets.
cpu_opt.ccCPU-spezifische Durchläufe.
gpu_opt.ccGPU-spezifische Durchläufe.
compiled_opt.ccFür CPU, GPU und XPU gemeinsame Durchläufe.
Vergessen Sie nicht, eine Build-Abhängigkeit hinzuzufügen.Fügen Sie die Registrierung von Karten/Tickets in Ihren PR ein(Beispiel), damit die Karte/das Ticket für alle
hlo-opt-Nutzer verfügbar ist.Erstellen Sie das
hlo-opt-Tool neu, validieren Sie die erfolgreiche Registrierung des Passes mit der Option--list-passesund führen Sie den Pass dann mit der Option--passesaus.$ hlo-opt --passes=foo-pass input.hloInformationen zum Schreiben von Unittests für den Pass finden Sie unter https://openxla.org/xla/test_hlo_passes.
Laufzeitmessung für Karten
Bei großen Modellen kann die vollständige Kompilierung einige Minuten dauern, was es schwierig macht, subtile Leistungsregressionen zu erkennen. Im Gegensatz dazu ermöglichen einzelne Durchläufe mit hlo-opt eine genaue Leistungsmessung und die einfache Erkennung selbst kleiner Steigerungen der Ausführungszeit, die durch neue Codeänderungen verursacht werden.
time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier
--xla_reduce_window_rewrite_base_length=128 input.hlo
[hlo-opt] HLO-Modulformate konvertieren
# Use the light weight version of the `hlo-opt` tool.
bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>
HLO Text → HLO Proto konvertieren
hlo-opt --emit-proto input.hlo
HLO Proto oder HLO Proto Binary –> HLO Text konvertieren
hlo-opt input.pbtxt or input.pb
[ptx-opt] Compiler LLVM Module down to PTX
Das Tool führt die LLVMIR-Optimierungspipeline aus und ruft dann „CompileToPtx“ auf.
bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 <filename>
Das Tool kann auch LLVMIR nach jedem Pfad ausgeben.
bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to=<path> --xla_gpu_dump_llvmir <filename>