Neues Back-End für XLA entwickeln

Dieser Leitfaden richtet sich an Systementwickler, die möchten, dass XLA Programme ausgeben, die ihre Hardware effizient ausrichten. Der Leitfaden enthält keine detaillierte Anleitung. Es werden Kenntnisse über LLVM, Bazel und XLA vorausgesetzt.

XLA bietet eine abstrakte Schnittstelle, die eine neue Architektur oder ein neuer Beschleuniger implementieren kann, um ein Back-End zum Ausführen von ML-Programmen zu erstellen, die von XLA ausgegeben werden. Das Retargeting von XLA sollte wesentlich einfacher und skalierbarer sein als die Implementierung jeder vorhandenen Operation aus einem Front-End-Framework wie PyTorch oder TensorFlow für neue Hardware.

Für die meisten Implementierungen gilt eines der folgenden Szenarien:

  1. Vorhandene CPU-Architektur, die noch nicht offiziell von XLA unterstützt wird, mit oder ohne vorhandenes LLVM-Back-End.
  2. Nicht CPU-ähnliche Hardware mit einem vorhandenen LLVM-Back-End.
  3. Nicht CPU-ähnliche Hardware ohne vorhandenes LLVM-Back-End.

Szenario 1: Vorhandene CPU-Architektur, die noch nicht offiziell von XLA unterstützt wird

Sehen Sie sich in diesem Szenario zuerst das vorhandene XLA-CPU-Back-End an. XLA erleichtert das Ansteuern verschiedener CPUs mithilfe von LLVM, da der Hauptunterschied zwischen XLA-Back-Ends für CPUs der von LLVM generierte Code ist.

Wenn der Hardwareanbieter ein LLVM-Back-End für seine Hardware hat, ist es einfach, das Back-End mit der mit XLA erstellten LLVM zu verbinden. Im JIT-Modus gibt das XLA-CPU-Back-End Code für die Host-CPU aus. Für die vorzeitige Kompilierung kann xla::AotCompilationOptions ein LLVM-Dreifach zur Konfiguration der Zielarchitektur bereitstellen.

Wenn zwar kein LLVM-Back-End, aber eine andere Art von Codegenerator vorhanden ist, sollte es möglich sein, den Großteil des vorhandenen CPU-Back-Ends wiederzuverwenden.

Szenario 2: Nicht CPU-ähnliche Hardware mit einem vorhandenen LLVM-Back-End

Es ist möglich, eine neue xla::Compiler-Implementierung auf den vorhandenen xla::CPUCompiler- und xla::GPUCompiler-Klassen zu modellieren, da diese bereits LLVM-IR ausgeben. Je nach Art der Hardware müssen viele Aspekte der LLVM-IR-Generierung möglicherweise geändert werden, es kann jedoch viel Code mit den vorhandenen Back-Ends geteilt werden.

Ein gutes Beispiel ist das GPU-Back-End von XLA. Das GPU-Back-End ist auf ein CPU-ähnliches ISA ausgerichtet. Daher sind einige Aspekte der Codegenerierung nur in der GPU-Domain vorhanden. Andere Arten von Hardware, z. B. DSPs wie Hexagon (mit einem vorgelagerten LLVM-Back-End), können Teile der LLVM-IR-Emissionslogik wiederverwenden. Andere Teile sind jedoch eindeutig.

Szenario 3: Nicht CPU-ähnliche Hardware ohne vorhandenes LLVM-Back-End

Wenn LLVM nicht verwendet werden kann, ist es am besten, ein neues Back-End für XLA für die gewünschte Hardware zu implementieren. Diese Option erfordert den größten Aufwand. Folgende Klassen müssen implementiert werden:

  • StreamExecutor: Für viele Geräte sind nicht alle Methoden von StreamExecutor erforderlich. Weitere Informationen finden Sie in den vorhandenen StreamExecutor-Implementierungen.
  • xla::Compiler: Diese Klasse kapselt die Kompilierung einer HLO-Berechnung in eine xla::Executable.
  • xla::Executable: Diese Klasse wird verwendet, um eine kompilierte Berechnung auf der Plattform zu starten.
  • xla::TransferManager: Diese Klasse ermöglicht Back-Ends, plattformspezifische Mechanismen zum Erstellen von XLA-Literaldaten aus bestimmten Gerätespeicher-Handles bereitzustellen. Mit anderen Worten, es hilft dabei, die Übertragung von Daten vom Host zum Gerät und zurück zu kapseln.