XLA-Tools verwenden

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