XLA (Accelerated Linear Algebra) ist ein Compiler für maschinelles Lernen (ML), der die 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, einem Ökosystem von Open-Source-Compiler-Technologien für ML, die von führenden ML-Hardware- und -Softwareunternehmen gemeinsam entwickelt wurde. Vor der Erstellung des OpenXLA-Projekts wurde XLA innerhalb des TensorFlow-Projekts entwickelt, die grundlegenden Ziele bleiben jedoch die gleichen:
Verbessern Sie die Ausführungsgeschwindigkeit. Kompilieren Sie Teilgrafiken, um die Ausführungszeit von kurzlebigen Vorgängen zu reduzieren und einen Overhead aus der Laufzeit zu vermeiden, führen Sie Pipeline-Vorgänge zusammen, 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, indem Sie viele Zwischenspeicher zwischenspeichern.
Reduzieren Sie die Abhängigkeit von benutzerdefinierten Vorgängen. Beseitigen Sie die Notwendigkeit vieler benutzerdefinierter Operationen, indem Sie die Leistung von automatisch zusammengeführten Low-Level-Operationen verbessern, um der Leistung benutzerdefinierter Operationen zu 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 auf dieser Hardware unverändert ausgeführt werden kann. Dies steht im Gegensatz zum Ansatz, einzelne monolithische Vorgänge für neue Hardware zu spezialisieren, für die Modelle neu geschrieben werden müssen, um diese Vorgänge zu nutzen.
Funktionsweise
Der XLA-Compiler nimmt Modellgrafiken aus ML-Frameworks, die in StableHLO definiert wurden, und kompiliert sie in Maschinenanweisungen 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 durch, 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 von Vorteil sind, und bestimmen, wie die Berechnung in Streams partitioniert wird. In dieser Phase können Back-Ends auch bestimmte Vorgänge oder Kombinationen davon mit optimierten Bibliotheksaufrufen kombinieren.
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-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-Datei auszugeben.
Innerhalb dieses Prozesses ist der XLA-Compiler in dem Sinne modular, dass 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.