Dieser Leitfaden richtet sich an Systementwickler, die mit XLA Programme ausgeben möchten, die auf ihre Hardware effizient ausgerichtet sind. Der Leitfaden enthält keine detaillierte Anleitung. Es werden Kenntnisse von 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 über ein Front-End-Framework wie PyTorch oder TensorFlow für neue Hardware.
Bei den meisten Implementierungen gibt es folgende Szenarien:
- Die vorhandene CPU-Architektur, mit oder ohne vorhandenes LLVM-Back-End, wird von XLA noch nicht offiziell unterstützt.
- Nicht CPU-ähnliche Hardware mit einem vorhandenen LLVM-Back-End.
- Nicht CPU-ähnliche Hardware ohne vorhandenes LLVM-Backend.
Szenario 1: Vorhandene CPU-Architektur, die noch nicht offiziell von XLA unterstützt wird
In diesem Szenario sehen Sie sich zuerst das vorhandene XLA-CPU-Back-End an. XLA erleichtert das Targeting auf verschiedene 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, kann das Back-End einfach mit der mit XLA erstellten LLVM verbunden werden. Im JIT-Modus gibt das XLA-CPU-Back-End Code für die Host-CPU aus. Für eine vorzeitige Kompilierung kann xla::AotCompilationOptions ein LLVM-Dreifaches für die Konfiguration der Zielarchitektur bereitstellen.
Wenn kein LLVM-Back-End vorhanden ist, aber eine andere Art von Codegenerator, sollte es möglich sein, den größten 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 nicht CPU-ähnliches ISA ausgerichtet. Daher gelten einige Aspekte der Codegenerierung nur für die GPU-Domain. 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-Backend
Wenn LLVM nicht verwendet werden kann, ist die beste Option, 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 vonStreamExecutorerforderlich. Weitere Informationen finden Sie in den vorhandenenStreamExecutor-Implementierungen.xla::Compiler: Diese Klasse kapselt die Kompilierung einer HLO-Berechnung in einexla::Executable.xla::Executable: Diese Klasse wird verwendet, um eine kompilierte Berechnung auf der Plattform zu starten.xla::TransferManager: Diese Klasse ermöglicht es 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.