XLA (Accelerated Linear Algebra) ist ein Compiler für maschinelles Lernen (ML), der lineare Algebra optimiert und so die Ausführungsgeschwindigkeit und Speichernutzung verbessert. Diese Seite bietet einen kurzen Überblick über die Ziele und die Architektur des XLA-Compilers.
Zielsetzungen
Heute unterstützt XLA mehrere ML-Framework-Front-Ends (einschließlich PyTorch, TensorFlow und JAX) und ist Teil des OpenXLA-Projekts – einer Open-Source-Compiler-Technologie für ML, die von führenden ML-Hardware- und -Softwareunternehmen gemeinsam entwickelt wird. Vor der Erstellung des OpenXLA-Projekts wurde XLA innerhalb des TensorFlow-Projekts entwickelt, die grundlegenden Ziele bleiben jedoch dieselben:
Verbessern Sie die Ausführungsgeschwindigkeit. Kompilieren Sie Teilgrafiken, um die Ausführungszeit von kurzlebigen Vorgängen zu reduzieren und den Overhead der Laufzeit zu vermeiden, fusionieren Sie Pipeline-Vorgänge, um den Speicheraufwand zu reduzieren, und spezialisieren Sie bekannte Tensorformen, um eine aggressivere konstante Weitergabe zu ermöglichen.
Arbeitsspeichernutzung verbessern: Analysieren und planen Sie die Arbeitsspeichernutzung und eliminieren Sie viele Zwischenspeicher.
Reduzieren Sie die Abhängigkeit von benutzerdefinierten Vorgängen. Beseitigen Sie die Notwendigkeit vieler benutzerdefinierter Operationen, indem Sie die Leistung der automatisch zusammengeführten Low-Level-Operationen so verbessern, dass sie der Leistung benutzerdefinierter Operationen entsprechen, die ursprünglich von Hand zusammengeführt wurden.
Portabilität verbessern Es sollte relativ einfach sein, ein neues Back-End für neue Hardware zu schreiben, damit ein großer Teil der ML-Modelle unverändert auf dieser Hardware ausgeführt werden kann. Dies steht im Gegensatz zum Ansatz, einzelne monolithische Operationen auf neue Hardware zu spezialisieren, bei denen Modelle neu geschrieben werden müssen, um diese Operationen nutzen zu können.
Funktionsweise
Der XLA-Compiler übernimmt Modelldiagramme aus ML-Frameworks, die in StableHLO definiert sind, und kompiliert sie in Maschinenbefehle für verschiedene Architekturen. StableHLO definiert einen versionierten Vorgangssatz (HLO = High-Level-Vorgänge), der eine Portabilitätsebene zwischen ML-Frameworks und dem Compiler bietet.
Im Allgemeinen umfasst der Kompilierungsprozess, der die Modellgrafik in eine zieloptimierte ausführbare Datei konvertiert, die folgenden Schritte:
XLA führt mehrere integrierte Optimierungs- und Analysedurchgänge für die StableHLO-Grafik aus, die zielunabhängig sind, z. B. CSE, zielunabhängige Vorgangsfusion und Pufferanalyse für die Zuweisung von Laufzeitspeicher für die Berechnung. Während dieser Optimierungsphase konvertiert XLA auch den StableHLO-Dialekt in einen internen HLO-Dialekt.
XLA sendet die HLO-Berechnung für weitere Optimierungen auf HLO-Ebene an ein Back-End, diesmal unter Berücksichtigung zielspezifischer Informationen und Anforderungen. Das GPU-Back-End kann beispielsweise Vorgangsfusionen durchführen, die speziell für das GPU-Programmiermodell vorteilhaft sind, und bestimmen, wie die Berechnung in Streams partitioniert wird. In dieser Phase können Back-Ends auch Muster für bestimmte Vorgänge oder Kombinationen davon mit optimierten Bibliotheksaufrufen verknüpfen.
Das Back-End führt dann eine zielspezifische Codegenerierung durch. Die in XLA enthaltenen CPU- und GPU-Back-Ends verwenden LLVM für Low-Level-IR, Optimierung und Codegenerierung. Diese Back-Ends geben die LLVM-IR aus, die erforderlich ist, um die HLO-Berechnung auf effiziente Weise darzustellen, und rufen dann LLVM auf, um nativen Code aus dieser LLVM-IR zu senden.
Innerhalb dieses Prozesses ist der XLA-Compiler modular, da er einfach in ein alternatives Back-End gesteckt werden kann, um eine neue HW-Architektur anzusteuern. Das GPU-Back-End unterstützt derzeit NVIDIA-GPUs über das LLVM-NVPTX-Back-End. Das CPU-Back-End unterstützt mehrere CPU-ISAs.