Diese Seite wurde von der Cloud Translation API übersetzt.
Switch to English

MLIR CodeGen für XLA

XLA arbeitet mit HloInstruction und führt viele Optimierungen für diese Darstellung durch, wobei viele davon zwischen Zielgeräten geteilt werden. Irgendwann wird ein linearer Zeitplan berechnet und der Speicherpuffer wird jedem Wert statisch zugewiesen. Der gerätespezifische Codegen durchläuft diese Sequenz und ruft "Emitter" auf, um eine für das Gerät geeignete Darstellung zu generieren (z. B. eine einzelne LLVM-Funktion pro XLA-Berechnung auf der CPU oder eine Sequenz von "Thunks", die GPU-Operationen kapseln und möglicherweise PTX generieren, wenn Targeting GPU).

Als Staging-Schritt sind wir derzeit dabei, den Prozess direkt nach Abschluss der Pufferzuweisungsphase abzufangen und stattdessen ein MLIR-Modul im lhlo Dialekt lhlo . Von dort aus führen wir den Codegen je nach Gerät mit MLIR-Komponenten (hauptsächlich Linalg, affine und GPU-Dialekt) durch.

Nachfolgend finden Sie den Aufzeichnungsplan für die schrittweise Migration von XLA / GPU unter Verwendung von lhlo als Codegen-Eingabe.

Aufgaben

Gastgeber Gerät
Eingabeformat HloInstruction * (Aufgabe 1) HloInstruction * (Aufgabe 1)
Ausgabeformat xla :: Thunk (Aufgabe 2) LLVM IR (Aufgabe 3)
  • Aufgabe 1 ändert sowohl das Host- als auch das Geräteeingabeformat von HloInstruction * in LHLO.
  • Aufgabe 2 ändert das Ausgabeformat des Hosts von Thunks in "Landing Pad für Host" (siehe unten).
  • Task 3 migriert die Geräteausgabe von LLVM IR auf eine Form von MLIR. Es ist für dieses Projekt optional. Weitere Informationen finden Sie im Abschnitt "Migrieren von Geräte-LLVM-IR".

In diesem Projekt wird Wert darauf gelegt, dass End-to-End-Modelle mit möglichst aktivierten LHLO-Emittern so weit wie möglich aktiviert werden. Dies impliziert, dass die folgende Reihenfolge der Ziele nach Priorität geordnet ist:

  • Machen Sie XLA / GPU mit LHLO-Emittern lauffähig, wobei vorhandene Thunks und Emitter unverändert bleiben.
  • Beseitigen Sie von Fall zu Fall die Verweise auf HloInstruction * in LHLO:
    • Wechseln Sie einen Legacy-Emitter zu einem MLIR-basierten Emitter (z. B. Linalg) oder
    • Übersetzen Sie den vorhandenen Emitter mechanisch, um eine MLIR-Darstellung zu erhalten (Migration auf Standard mit GPU-Dialekt).

Thunks migrieren (Aufgabe 2)

xla :: gpu :: Thunk ist eine Datenstruktur, die:

  • Kann vom Host aufgerufen werden (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • Trägt verschiedene Daten in seinen Unterklassen.
  • Interagiert mit BufferAllocation :: Slice und StreamExecutor.
  • Startet Kernel
  • Ruft alle Laufzeitbibliotheken auf.

Die Kosten dafür beinhalten:

  • Darstellung op-spezifischer Konfigurationsdaten (z. B. Faltungskonfigurationen).
  • Migrieren von Operationsformen und Operandenformen.
  • Darstellung eines Baumes von Thunks (während, Zustand usw.).

Die Migrationsarbeit ist unabhängig von der LHLO / Emitter-Migration. Bei begrenzten Ressourcen wird die LHLO / Emitter-Migration priorisiert.

Wir haben verschiedene Möglichkeiten, wie Sie den Host-seitigen Teil von LHLO senken können:

  • TFRT
    • (Pro) großartige CUDA- und HIP-Wrapper zur Verwendung.
    • (Pro) einfach zu implementierende Bibliotheksaufrufe (cuDNN, cuBLAS, cuFFT usw.), da TFRT-Operationen von C ++ - Code interpretiert werden.
    • Die (Con) Host-Seite befindet sich in der Entwicklung und wird nicht getestet.
  • Angepasster CPU-Code
    • (Pro) große geringere Fähigkeit. Erstellen Sie ein paar Schleifen und Bedingungen und fertig.
    • (Con) GPUDialect modelliert noch keine Ketten / Streams / Asynchronität / Gerätezuordnung.
    • (Con) Die CUDA / HIP-Laufzeitunterstützung ist minimal (Toolkit-Pfad, Version, dynamisches Laden usw.).
  • Vorhandene (interpretierende) XLA-Laufzeit

Entscheidung: TFRT übernehmen, aber auch Jitting-CPU-Code in TFRT unterstützen.

Migrieren des Geräte-LLVM-IR (Aufgabe 3)

Ein Elementaremitter erzeugt eine Zieloperation, indem er sie Element für Element füllt. Jedes Ausgabeelement hängt von einer Reihe von Elementen aus den Operanden ab. Alle Elemente werden durch Kombinieren des Puffers mit dynamischen Indizes beschrieben. Es reicht aus, fast alle "mathematischen" Operationen zu beschreiben, aber aus Leistungsgründen wird nur eine große Teilmenge der "mathematischen" Operationen direkt in (Cpu | Gpu) ElementalIrEmitter implementiert.

ElementalIrEmitter ist insofern einzigartig:

  • Ein großer Teil des Codes wird von XLA / GPU und CPU gemeinsam genutzt.
  • Es stellt einen großen Teil der Operationen dar, die in Modellen zu sehen sind, einschließlich aller elementweisen Operationen.
  • Die meisten Fusionen hängen ausschließlich von ElementalIrEmitter ab.
  • Es ist strukturell einfach, da es eine Datenabhängigkeits-DAG zwischen Operationselementen und Operandenelementen beschreibt.
  • Es ist meistens portabel und auf hohem Niveau (z. B. im Gegensatz zu GPU kReduce und GPU kCopy).
  • Die dynamische Formunterstützung ist zumindest für elementweise Operationen einfach.

Nun gibt es für alle Operationen, die elementar emittiert werden oder nicht, verschiedene Varianten des Endzustands jeder XLA-Operation:

  1. Gerätecode bleibt als LLVM IR.
  2. Refactor den alten Emitter wie LHLO -> MLIR LLVM Dialekt:
    • (Kosten) Wird wegwerfbar sein, wenn wir letztendlich auf Standard migrieren wollen.
    • (Vorteil) Es ist einfach und mechanisch. Kann in kurzer Zeit durchgeführt werden.
    • (Nutzen) Es profitiert nicht mehr als (1).
  3. Refactor alte Emitter wie LHLO -> MLIR GPU + Standard + Loops:
    • (Kosten) Das Anheben vorhandener Emittenten auf Standard bringt einige Herausforderungen mit sich. Zeiger und GEPs müssen in MemRefs und SubViews konvertiert werden. Die Vollständigkeit von amdgpu sicherzustellen, ist eine andere Sache.
    • (Kosten) XLA / GPU stützt sich stark auf LLVM-Metadaten:
      • range für Block- / Thread-Indizes.
      • align , dereferenceable , invariant.load , alias.scope , noalias für load / noalias .
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable für sequentielle Schleifen.
    • (Nutzen) Kann langfristig sein. Tragbarer.
  4. Refactor alte Emitter, um LHLO -> Linalg zu sein, und schreibe neue Linalg-Emitter
    • (Kosten) Dies ist von Fall zu Fall. Im Vergleich zu früheren Optionen muss eine neue Implementierung, die der Leistung von XLA entspricht, den Benchmark <-> Optimierungsworkflow durchlaufen, was für einige Operationen erhebliche Kosten verursachen kann.
    • (Vorteil) einheitlicher Stapel; gemeinschaftliche Unterstützung; Portabilität; mehr Optimierungspotentiale.

Schlussfolgerungen:

  • Gehen Sie nicht für (2). (1) oder (3) sind einfach besser als (2). (2) kostet mehr als (1), da es viel mechanisches Refactoring erfordert. Mit (1) können wir immer noch das Ziel erreichen, dass XLA MLIR-Emitter aufnehmen kann. Dies geschieht durch Ausführen von LHLO -> LLVM IR -> Ausführen älterer Geräteemitter.
  • ElementalIrEmitter-Operationen gehen für (4), aber nicht inkrementell. Es gibt keine Möglichkeit, dies op für op zu tun, da alle elementar emittierten Operationen mit demselben Diagramm verbunden sind. Diese Arbeit kann auch als Vereinigungspunkt mehrerer laufender Kräfte dienen (xla / service / mlir_gpu, der Kernelgenerator, Linalg).
  • Alle anderen Operationen gehen für (1). Als Streckenziel können sie nach (3) oder (4) migriert werden.

Priorisierung

Während alle drei oben genannten Aufgaben parallelisierbar sind, müssen sie unter begrenzten Ressourcen serialisiert werden. Die Priorisierung konzentriert sich auf sichtbare Ergebnisse für die Ausführung jeder Aufgabe.

Die Priorisierung lautet: Task1 (LHLO für Legacy-Emitter)> Task 2 (Thunks)> Task 3 (MLIR-Emitter).

Am Ende von Aufgabe 1 können Benutzer von XLA ein LHLO (z. B. einen Kernelgenerator) generieren und ausführen. Das Kompilierungsformat ist kein serialisierbares MLIR.

Am Ende von Aufgabe 2 senkt sich LHLO auf das richtige, serialisierbare MLIR. Dies ermöglicht die Offline-Kompilierung.

Bis zum Ende von Aufgabe 3 sind alle XLA-Emitter in ihrer Implementierung MLIR-basiert.

Detailliertes Design

Schritt 1: (Aufgabe 1) Schließen Sie LHLO ab und lassen Sie Legacy-Emittenten LHLO nehmen

Durch diesen Schritt interagieren alle vorhandenen XLA / GPU-Emitter mit MLIR-Operationen. Dieser Schritt ist reines Refactoring und NFC.

Dieser Schritt ist größtenteils mechanisch, es lohnt sich jedoch, die folgenden Diskrepanzen zwischen einer nicht verschachtelten HloComputation und LHLO zu beachten:

  • Jeder HloInstruction hat direkten Zugriff auf seine Operanden (eine Datenfluss-DAG). Im Gegensatz dazu hat jede LHLO-Operation nur Zugriff auf ihre Operandenpuffer (eine zweiteilige zwischen Operationen und Puffern). LHLO-Ops müssen Use-Def-Ketten durchlaufen, um auf ihre Operanden-Ops zugreifen zu können.
  • Unverschachtelte Legacy-Emittenten greifen empirisch fast nie auf ihre Operanden zu. Die einzige Ausnahme ist kReduce.
  • Nicht verschachtelte Legacy-Emitter greifen nur zum Abrufen von Slices auf BufferAssignment zu, nicht für den Zugriff auf Aux-Datenstrukturen wie dataflow_analysis () oder alias_analysis (). llvm_ir erstellt eine eigene alias_analysis () basierend auf Slice-Informationen.

Die Schlussfolgerung ist, dass LHLO ohne großen Aufwand genau passen sollte.

Schritt 2: (Optionale) Profiling-Unterstützung

Dieser Schritt ist nur erforderlich, wenn wir beginnen, einen Teil der XLA-Thunk-Logik zu verwerfen (siehe nächster Schritt).

Bevor Sie MLIR-basierte Emitter aktivieren, müssen Sie ein Profil für MLIR-basierte Emitter erstellen.

Derzeit führt XLA eine eigene Profilerstellung durch, indem der Timer von StreamExecutor aufgerufen wird. Der Timer unter der Haube fügt zwei Ereignisse vor und nach einem Kernelstart ein und misst die Synchronisierungszeit zwischen diesen beiden Ereignissen.

Es gibt ungefähr drei Ansätze zur Unterstützung der Profilerstellung in MLIR:

  • Führen Sie einen Profiler Ende-zu-Ende aus
  • Fügen Sie mit einem injizierten Profiler für jede Operation in LHLO eine Profiloperation hinzu.

Der "End-to-End" -Ansatz ist für MLIR transparent, weist jedoch das gleiche Problem auf, das XLA dazu veranlasst, ihn überhaupt nicht zu verwenden: Von einem Profiler (nvprof / ...) gesammelte Bibliotheksaufrufe können sich nicht leicht auf HLO beziehen ops. Beispielsweise startet cuDNN mehrere Kernel für jedes HLO, und es ist schwer zu sagen, welche Kernel welchem ​​HLO entsprechen.

Der Ansatz des "injizierten Profilers" erfordert:

  • LHLO nimmt einen Profiler als Parameter.
  • Einfügen von profile.start / profile.end vor und nach jeder Operation.
  • Ein Durchgang von diesem senkt das Profil. {start, end} zu einer C ++ - Implementierung.

Die genaue Profilerstellung ist für MLIR-generierte Operationen nicht einfach durchzuführen, da:

  • MLIR hat keinen Timer und hängt auch nicht von TFRT / StreamExecutor ab.
  • MLIR ruft C-Funktionen mit komplizierten Parametern nicht einfach auf.

Schritt 3: (Aufgabe 2) Migrieren von Thunks

Als Anmerkung gibt es ungefähr drei Arten von Thunks:

  • KernelThunk, der einen Kernel startet.
  • Control Flow Thunks, die über eine Host Control Flow-Logik (bedingt, while, for, sequence) und Startkörper-Kernel verfügen.
  • Bibliotheks-Thunks: cuDNN, cuBLAS, cuFFT, NCCL usw.

Der Plan ist:

  • Machen Sie Thunks (de) serialisierbar.
  • Verbessern Sie TFRT in einem Zustand, in dem diese Semantik unterstützt werden kann.
  • Wenn sich der Status verbessert, migrieren Sie einzelne Thunks schrittweise.

Diese Aktionselemente sind nur teilweise bestellt. Die tatsächliche Ausführungsreihenfolge / Engineering-Parallelität ist zu bewerten.

Schritt 4: (Aufgabe 3) Migrierter ElementalIrEmitter

Sobald die Profilerstellung fertig ist, können wir alle ElementalIrEmitter-basierten Emitter in MLIR vervollständigen und optimieren. Dann schalten wir sie standardmäßig ein, vorausgesetzt, alle diese MLIR-basierten Emitter verwenden einen einzigen Stream.

Beachten Sie, dass es auch von Vorteil ist, den ElementalIrEmitter von XLA / CPU zu migrieren, da diese einen großen Teil des Codes gemeinsam nutzen.

Nachdem alle Benchmarking- und Leistungssuchvorgänge durchgeführt wurden (TODO: Definieren der Leistungsparität), aktivieren wir den neuen MLIR-basierten Elementaremitter und löschen den alten ElementalIrEmitter.

Dieser Schritt bietet auch einfache Fusionsübergänge (verschachtelte Operationen) für die spätere Migration.

Schritt 5: Multi-Stream-Unterstützung oder Drop

Wir können einige der Emitter erst löschen, wenn wir sie in MLIR unterstützen oder die Funktion löschen. Es ist eine relativ große Menge an Arbeit in MLIR und eine kleine Menge an Gewinn für XLA. Wir sollten aktuelle Benutzer von Multi-Stream-XLA / GPU-Benutzern untersuchen und versuchen, diese Funktion zu löschen, wenn dies zumutbar ist.

Schritt 6: (Aufgabe 3) Migrierte Geräteoperationen

Dieser Schritt migriert alle nicht verschachtelten Operationen, dann können wir alle nicht verschachtelten Emitter löschen.

Dies erfordert ein Umschreiben / Refaktorieren für kCopy und kReduce. An kReduce wird bereits viel gearbeitet, daher bleibt abzuwarten, wie viel Arbeit tatsächlich erledigt werden muss.