XLA-Architektur

Warum haben wir XLA gebaut?

Wir hatten mehrere Ziele für XLA, mit TensorFlow zusammenzuarbeiten:

  • Verbessern Sie die Ausführungsgeschwindigkeit. Kompilieren Sie Untergraphen, um die Ausführungszeit kurzlebiger Operationen zu verkürzen, um Overhead aus der TensorFlow-Laufzeit zu vermeiden, verschmelzen Sie Pipeline-Operationen, um den Speicher-Overhead zu reduzieren, und spezialisieren Sie sich auf bekannte Tensorformen, um eine aggressivere konstante Ausbreitung zu ermöglichen.

  • Verbessern Sie die Speichernutzung. Analysieren und planen Sie die Speichernutzung und eliminieren Sie im Prinzip viele Zwischenspeicherpuffer.

  • Reduzieren Sie die Abhängigkeit von benutzerdefinierten Operationen. Beseitigen Sie die Notwendigkeit vieler benutzerdefinierter Ops, indem Sie die Leistung von automatisch fusionierten Ops auf niedriger Ebene verbessern, um sie an die Leistung von benutzerdefinierten Ops anzupassen, die von Hand fusioniert wurden.

  • Reduzieren Sie den mobilen Platzbedarf. Beseitigen Sie die TensorFlow-Laufzeit, indem Sie den Untergraphen vorab kompilieren und ein Objekt / Header-Dateipaar ausgeben, das direkt mit einer anderen Anwendung verknüpft werden kann. Die Ergebnisse können den Platzbedarf für mobile Inferenz um mehrere Größenordnungen verringern.

  • Verbessern Sie die Portabilität. Machen Sie es relativ einfach, ein neues Backend für neuartige Hardware zu schreiben. Ab diesem Zeitpunkt wird ein großer Teil der TensorFlow-Programme unverändert auf dieser Hardware ausgeführt. Dies steht im Gegensatz zu dem Ansatz, einzelne monolithische Ops auf neue Hardware zu spezialisieren, bei dem TensorFlow-Programme neu geschrieben werden müssen, um diese Ops nutzen zu können.

Wie funktioniert XLA?

Die Eingabesprache für XLA heißt "HLO IR" oder einfach HLO (High Level Operations). Die Semantik von HLO wird auf der Seite Operationssemantik beschrieben. Es ist am bequemsten, sich HLO als Compiler-IR vorzustellen.

XLA verwendet in HLO definierte Diagramme ("Berechnungen") und kompiliert sie zu Maschinenanweisungen für verschiedene Architekturen. XLA ist modular in dem Sinne, dass es einfach ist, ein alternatives Backend einzusetzen, um auf eine neuartige HW-Architektur abzuzielen. Das CPU-Backend für x64 und ARM64 sowie das NVIDIA-GPU-Backend befinden sich im TensorFlow-Quellbaum.

Das folgende Diagramm zeigt den Kompilierungsprozess in XLA:

XLA enthält mehrere zielunabhängige Optimierungen und Analysedurchläufe, z. B. CSE , zielunabhängige Operationsfusion und Pufferanalyse zum Zuweisen von Laufzeitspeicher für die Berechnung.

Nach dem zielunabhängigen Schritt sendet XLA die HLO-Berechnung an ein Backend. Das Backend kann weitere Optimierungen auf HLO-Ebene durchführen, diesmal unter Berücksichtigung zielspezifischer Informationen und Anforderungen. Beispielsweise kann das XLA-GPU-Backend eine Operationsfusion durchführen, die speziell für das GPU-Programmiermodell von Vorteil ist, und bestimmen, wie die Berechnung in Streams aufgeteilt werden soll. In diesem Stadium können Backends auch bestimmte Operationen oder Kombinationen davon mit Mustern an optimierte Bibliotheksaufrufe anpassen.

Der nächste Schritt ist die zielspezifische Codegenerierung. Die in XLA enthaltenen CPU- und GPU-Backends verwenden LLVM für Low-Level-IR, Optimierung und Codegenerierung. Diese Backends senden das LLVM-IR aus, das erforderlich ist, um die XLA-HLO-Berechnung auf effiziente Weise darzustellen, und rufen dann LLVM auf, um nativen Code von diesem LLVM-IR auszugeben.

Das GPU-Backend unterstützt derzeit NVIDIA-GPUs über das LLVM-NVPTX-Backend. Das CPU-Backend unterstützt mehrere CPU-ISAs.