Dieser Vortrag von Kevin Gleason und Peter Gavin (Google) bietet eine tiefgehende architektonische Übersicht über das XLA-Compiler-System (Accelerated Linear Algebra) und dessen Zwischendarstellungen (Intermediate Representations, IR). Nachgezeichnet wird die Evolution der IR-Schichten von der ursprünglichen HLO (2017) über MHLO und CHLO bis hin zum standardisierten StableHLO und dessen versionierter Form VHLO innerhalb des OpenXLA-Ökosystems. Neben der geräte- und framework-agnostischen Entkopplung wird die C-Schnittstelle PJRT für Laufzeit-, Kompilierungs- und Puffer-Management erläutert. Im zweiten Teil stehen die Compiler-Interna im Fokus: XLA nutzt rekursionsfreie, datenflussbasierte Graphen, um eine vollständig statische Speicherallokation zur Laufzeit zu garantieren. Durch Techniken wie minor-to-major Layout-Transformationen, Vektor-Tiling und primär die Instruktions-Fusion (Loop Fusion) werden Bandbreitenengpässe beim DRAM-Zugriff minimiert. Abschließend werden Mechanismen zur Steuerung numerischer Präzision (z. B. BFloat16-Emulation) und Ursachen hoher Kompilierungszeiten adressiert.
Wichtigste Punkte & Zeitstempel
0:05 Einführung und Agenda: Vorstellung des XLA-Compiler-Stacks durch Kevin Gleason (StableHLO, PJRT) und Peter Gavin (XLA TPU Compiler) von Google.
0:54 Entstehung von HLO: HLO (High Level Optimizer) wurde 2017 als framework- und hardware-unabhängige Zwischendarstellung mit ca. 120 Basis-Operationen entwickelt, um TensorFlow von der TPU-Hardware zu entkoppeln.
2:07 Einführung von MHLO: MHLO (Meta HLO) integrierte HLO 2019 in das MLIR-Ökosystem, um domänenspezifische Features wie dynamische Shapes und Quantisierung abzubilden.
3:31 StableHLO & OpenXLA: 2022 als stabiles Kern-IR für das OpenXLA-Projekt eingeführt (~100 Ops), um Kompatibilität zwischen heterogenen Frameworks (JAX, PyTorch) und Hardware-Backends (TPU, AWS Neuron, Apple Metal) zu sichern.
5:22 VHLO & Kompatibilitätsmanagement: VHLO (Versioned StableHLO) nutzt MLIR-Bytecode sowie In-Memory-Upgrades und -Downgrades, um eine zweijährige Vorwärts- und Rückwärtskompatibilität zu gewährleisten.
6:35 Client HLO (CHLO): High-Level-Op-Set zur Erhaltung mathematischer Abstraktionen; erlaubt wahlweise direkte Backend-Abbildung oder gezielte Dekposition in StableHLO für einheitliche Numerik.
7:42 Lowering & Custom Kernels: Frameworks lowern Primitive zu StableHLO; benutzerdefinierte Kernels (z. B. JAX Palace) werden über custom_call-Ops als opake MLIR-Blobs an backend-spezifische Compiler übergeben.
9:21 PJRT-API-Schnittstelle: Abstraktionsschicht mit standardisierten C-APIs für Kompilierung, Programmausführung und Host-Device-Speichertransfers.
10:31 HLO-Modulstruktur & Grapheneigenschaften: HLO-Module enthalten datenflussbasierte Instruktionsgraphen (Computations). Rekursion ist strikt verboten, um statische Speichergrenzen vorab zu berechnen.
12:05 Mehrdimensionale Arrays, Layouts & Tiling: Speicherlayouts werden von minor-to-major definiert. Tiling strukturiert Daten in mehrdimensionalen Blöcken (z. B. $8 \times 8$), angepasst auf Vektorregister der Zielhardware.
18:26 Instruktions-Fusion: Kernoptimierung von XLA. Führt Untergraphen (elementweise Operationen) in einheitliche Kernel zusammen, um Zwischenergebnisse direkt in Registern zu halten und DRAM-Lese-/Schreibzugriffe zu eliminieren.
19:45 Compiler-Frontend vs. Backend: Das Frontend führt geräteunabhängige Optimierungen (CSE, Fusionsentscheidungen) auf HLO-Ebene durch; das Backend übersetzt gefuste Graphen in Maschinencode (z. B. via LLVM).
24:00 Q&A – Custom Ops & Fusion: Custom Ops agieren standardmäßig als Blackboxes und werden vom generischen Fusion-Pass nicht automatisch mit anderen Operationen verschmolzen.
25:50 Q&A – Statische Speicherallokation & Kompilierungszeit: Das Fehlende dynamischer Speicherallokation erzwingt vollständiges Unrolling von Schleifen/Graphen, was bei tiefen Netzen zu langen Kompilierungszeiten führt.
27:00 Q&A – Numerische Präzision & Flag-Steuerung: Steuerung von Rundungsverhalten und Emulation niedriger Präzision (z. B. BFloat16 via F32) über XLA_FLAGS sowie precision_config bei Matrixmultiplikationen.
Glossar
XLA (Accelerated Linear Algebra): Domänenspezifischer Compiler für lineare Algebra, der Graphstrukturen aus Machine-Learning-Frameworks in effizienten Maschinencode für Beschleuniger übersetzt.
HLO (High Level Optimizer): Die primäre Graphen-Zwischendarstellung (IR) innerhalb des XLA-Compilers, in der geräteunabhängige Optimierungen durchgeführt werden.
MHLO (Meta HLO): Ein auf MLIR (Multi-Level Intermediate Representation) basierender Dialekt von HLO, der erweiterte Sprachfeatures wie dynamische Dimensionen unterstützt.
StableHLO: Ein durch das OpenXLA-Konsortium standardisierter, hochkompatibler IR-Op-Set-Standard, der als Schnittstelle zwischen ML-Frameworks und Compilern dient.
VHLO (Versioned StableHLO): Versionierte Repräsentation von StableHLO, die Bytecode-Stabilität und Konvertierungen zwischen verschiedenen Dialektversionen ermöglicht.
CHLO (Client HLO): High-Level-Dialekt auf Framework-Ebene, der komplexe Operationen vor der Dekposition in elementare StableHLO-Ops kapselt.
PJRT (Plugin Joint RunTime): Einheitliche C-API zur Anbindung von Hardware-Backends an ML-Frameworks zur Steuerung von Kompilierung, Execution und Speicherverwaltung.
Instruction Fusion (Instruktions-Fusion): Compiler-Pass, der mehrere aufeinanderfolgende Operationen in einer einzigen Kernel-Schleife zusammenfasst, um Zugriffe auf den Hauptspeicher (DRAM) zu vermeiden.
Tiling: Verfahren zur Umstrukturierung und Ausrichtung von Datenarrays in feste Speicherblöcke (Tiles), um die Auslastung von Vektor- und Matrix-Auszuführungseinheiten zu maximieren.
TPU (Tensor Processing Unit): Anwendungspezifischer Integrierter Schaltkreis (ASIC) von Google, optimiert für Vektor- und Matrixberechnungen im Deep Learning.