XLA-Tools

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

  1. Schreiben Sie zuerst Ihren Pass.
  2. 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.cc Hardwareunabhängige Karten oder Tickets.
    cpu_opt.cc CPU-spezifische Durchläufe.
    gpu_opt.cc GPU-spezifische Durchläufe.
    compiled_opt.cc Fü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.

  3. Erstellen Sie das hlo-opt-Tool neu, validieren Sie die erfolgreiche Registrierung des Passes mit der Option --list-passes und führen Sie den Pass dann mit der Option --passes aus.

    $ hlo-opt --passes=foo-pass input.hlo
    
  4. Informationen 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>