In diesem Dokument wird der Weg eines High Level Optimizer-Moduls (HLO) von seinem ursprünglichen Zustand bis zu einem endgültigen ausführbaren Programm beschrieben. Manchmal lassen wir das „Modul“ weg und bezeichnen es nur als „HLO“.
HLO vor der Optimierung
Wir beginnen mit dem HLO-Modul vor der Optimierung. Das HLO vor der Optimierung enthält keine Operationen (Ops), die als intern für XLA gelten, z. B. fusion oder bitcast. Vorgänge haben in dieser Phase kein Layout. Falls doch, wird es ignoriert. HLO vor der Optimierung wird in der Regel von Frameworks auf höherer Ebene wie TensorFlow und JAX erzeugt. Wenn Sie das XLA-Flag -xla_dump_to verwenden, wird das HLO vor der Optimierung in eine Datei mit dem Dateinamen-Suffix „before_optimizations.txt“ geschrieben.
HLO-Modul optimieren
In der XLA:GPU-Pipeline wird das HLO vor der Optimierung durch eine Reihe von Durchläufen in optimiertes HLO umgewandelt. Die Durchläufe können semantisch gruppiert und in der folgenden Reihenfolge ausgeführt werden:
Sharding von Karten/Tickets
Dazu gehören Passes wie der Shardy Partitioner oder die für SPMD-Sharding.
Optimierungsdurchläufe
Dazu können sowohl Legalisierungsvorgänge als auch Vereinfachungsvorgänge gehören.
Kollektive Optimierungsdurchläufe
Ähnlich wie Optimierungsdurchläufe, aber mit Fokus auf kollektive Vorgänge.
Layoutzuweisungspässe
Jeder HLO-Vorgang hat ein Layout, das Teil der Anweisungsform ist. Das Layout steuert, wie der Tensor physisch im Speicher angeordnet wird.
Beispiel für eine Form mit Layout:
f32[10,20,30]{2,0,1}
Nach dem Elementtyp folgen die logischen Dimensionen der Form und dann die Layoutpermutation in aufsteigender Reihenfolge. In diesem Beispiel ist die kleinste Dimension 30, die zweitkleinste Dimension 10 und die größte Dimension 20.
Ziel der Layoutzuweisung ist es, die Anzahl der erforderlichen physischen Transpositionen mithilfe einer Greedy-Strategie zu minimieren. Sie beginnt mit bestimmten Layoutbeschränkungen (z.B. erwarten cuDNN-/cuBLAS-Bibliotheken aufeinanderfolgende Dimensionen) und überträgt Layouts „nach unten“ und dann „nach oben“ im HLO-Diagramm. Am Ende der Layoutweitergabe können für einige Anweisungen widersprüchliche Layouts vorhanden sein, eines, das von einem Operanden weitergegeben wurde, und eines, das von einem Nutzer weitergegeben wurde. Um diesen Konflikt zu beheben, wird eine copy-HLO-Anweisung eingefügt, die das Layout vom Operandenlayout zum Anweisungslayout ändert.
Layoutnormalisierungsdurchläufe
Da es etwas schwierig ist, die physische Form zu ermitteln, wird versucht, die Form durch die Layoutnormalisierung so umzuschreiben, dass das Standardlayout {rank-1, rank-2, …, 0} verwendet wird. Im obigen Beispiel wäre die normalisierte Form f32[20,10,30]{2,1,0}. Kopieroperationen, die Layouts ändern, werden als Kombination aus transpose und bitcast neu geschrieben. Da wir derzeit nicht alle Vorgänge normalisieren können, gibt es immer noch einige Vorgänge, die möglicherweise nicht standardmäßige Layouts haben, insbesondere gather und dot. An den Grenzen zwischen normalisierten und nicht normalisierten Operationen gibt es bitcast-Operationen, die eine Transponierung darstellen, d.h. eine Transponierung mit einem zugewiesenen Layout, das sie physisch zu einer No-Op macht.
Durch die Layoutnormalisierung werden auch einige implizite Transponierungen explizit gemacht. Das ist wichtig, da der Codegenerator explizite Transponierungen mit einem speziellen Emitter verarbeiten kann. Beispielsweise darf eine Umformung technisch gesehen ein anderes physisches Layout zwischen Operand und Ergebnis haben (z.B. aufgrund eines unterschiedlichen Rangs). Beim ReshapeDecomposer-Durchlauf, der als Teil der Layoutnormalisierungsdurchläufe ausgeführt wird, wird eine Reshape-Operation in eine Sequenz aus transpose, Reshape bitcast und transpose umgewandelt.
Optimierungsdurchläufe nach der Layoutzuweisung
Die wichtigsten Optimierungen sind Triton-Fusions (GEMM-Fusions + Softmax-/Layernorm-Fusions) oder das Umschreiben in Bibliotheksaufrufe. In diesem Schritt wird auch die automatische Optimierung ausgeführt. Dabei wählt XLA zwischen verschiedenen Emitters aus, wählt den besten Algorithmus für Faltungen oder Punkte aus und findet die beste Kachelung für Fusionen, die vom Triton-Emitter verarbeitet werden.
Fusion-Karten/Tickets
Die beiden wichtigsten Durchläufe sind die PriorityFusion- und die Multi-Output-Zusammenführung.
In PriorityFusion werden Fusions auf Grundlage des Kostenmodells gebildet. Beim Zusammenführen würden wir das Duplizieren von Vorgängen mit mehreren Nutzern zulassen, wenn der Vorgang in alle Nutzer zusammengeführt werden kann. Wir würden auch die Erweiterung vorhandener Triton-Softmax-Fusions zulassen, sofern dies möglich ist.
Multi-Output-Fusion ist ein separater Durchlauf, mit dem Vorgänge/Fusionen zusammengeführt werden können, die einen Operanden gemeinsam nutzen. Außerdem können Operanden/Operandenfusionen in Nutzer ohne Duplizierung zusammengeführt werden, indem zusätzliche Ausgaben hinzugefügt werden. Andere Nutzer des zusammenzuführenden Vorgangs können dann zu diesen Ausgaben umgeleitet werden. Bei diesem Durchlauf muss darauf geachtet werden, dass keine Zyklen in den HLO-Graphen eingeführt werden.
Nach der Multi-Output-Zusammenführung wird die Eliminierung gemeinsamer Teilausdrücke (HloCSE-Pass) ausgeführt. Dabei werden möglicherweise zuvor duplizierte Vorgänge wieder zusammengeführt, wenn sie in derselben Zusammenführung enden.
Mehrere Durchgänge nach der Fusion
Mehrere Durchläufe im Zusammenhang mit Kollektiven (z. B. Umwandlung in asynchrone Kollektive oder Erzwingen einer bestimmten relativen Reihenfolge von Kollektiven).
Schließlich führen wir CopyInsertion aus, wobei Kopien hinzugefügt werden, um zu verhindern, dass Daten, die noch an anderer Stelle benötigt werden, durch In-Place-Vorgänge überschrieben werden.
Am Ende der Optimierung wird das optimierte HLO in eine Datei mit dem Dateinamesuffix „after_optimizations.txt“ geschrieben, wenn das Flag -xla_dump_to verwendet wird. Wenn Sie das HLO nach Zwischenübergängen ausgeben möchten, die das HloModule tatsächlich ändern, können Sie das Flag -xla_dump_hlo_pass_re=.* verwenden (oder einen bestimmten regulären Ausdruck, um es auf bestimmte Übergänge zu beschränken).
Wird geplant
Ein HLO-Modul ohne Zeitplan hat immer noch einen gewissen Spielraum bei der Reihenfolge, in der Vorgänge verarbeitet werden. Jede topologische Sortierung, die die Beziehungen zwischen Operanden und Ergebnissen sowie die Steuerungsabhängigkeiten berücksichtigt, ist gültig. Durch die Planung wird festgelegt, welche spezifische Reihenfolge verwendet werden soll. In dieser Phase geht es hauptsächlich um den maximalen Arbeitsspeicherverbrauch, der von der Lebensdauer der Tensoren abhängt. In einem ersten Schritt testen wir verschiedene Scheduler-Algorithmen und wählen den Zeitplan aus, der den maximalen Speicherverbrauch minimieren sollte. Hinweis: An dieser Stelle arbeiten wir noch nicht mit physischen Puffern (das erfolgt bei der Pufferzuweisung), sondern simulieren die Speichernutzung.
Anschließend wird der LatencyHidingScheduler-Pass ausgeführt und versucht, die Überschneidung zwischen Berechnung und Kommunikation zu maximieren. Dadurch kann die Arbeitsspeichernutzung jedoch wieder steigen.
Wenn der maximale Arbeitsspeicherverbrauch höher ist als der verfügbare Arbeitsspeicher, führen wir HloRematerialization aus. Bei diesem Durchlauf wird versucht, die Arbeitsspeichernutzung auf Kosten der Leistung zu verringern. So werden beispielsweise einige Fusionsvorgänge aufgeteilt und einige Vorgänge dupliziert, um die Lebensdauer von Puffern zu verkürzen. Wenn eine Rematerialisierung erfolgt, kann es sinnvoll sein, die Speicheranforderungen auf der Modellseite zu reduzieren, z.B. durch die Verwendung kleinerer Batchgrößen.
Pufferzuweisung
Unmittelbar bevor wir zu LLVM IR wechseln, führen wir die Pässe für die Pufferzuweisung aus, mit denen jeder Anweisung im HLO-Diagramm Puffersegmente zugewiesen werden. Die Pufferzuweisung erfolgt in mehreren Schritten:
Mit
HloDataflowAnalysiswerdenHloValues(im Wesentlichen logische Puffer) Anweisungen zugewiesen. Bei In-Place-Vorgängen kann dasHloValueeines Operanden wiederverwendet werden. Ein Vorgang kann mehr als einHloValuedefinieren, z.B. mit einer Tupel-Ergebnisform.HloAliasAnalysisversucht, Puffer für Aliasing-Vorgänge zu kombinieren, und berechnet eine Zuordnung vonHloValuezuHloBuffer.BufferAssignmentberechnet eine Zuordnung vonHloBufferszu Puffer-Slices in einem großen Puffer, sodass derselbe Puffer-Slice nicht für verschiedeneHloBuffersmit sich überschneidenden Lebenszyklen verwendet wird. Bei Vorgängen, die Aliasse verwenden können, ist eine leichte Überschneidung in Ordnung. Die Endzeit des einenHloBufferkann mit der Startzeit des anderenHloBufferübereinstimmen. Wenn Sie das Flag-xla_dump_toverwenden, werden einige Informationen zur Pufferzuweisung in eine Datei mit dem Namenssuffix „after_optimizations-buffer-assignment.txt“ geschrieben.
Thunks
Nachdem ein HLO-Diagramm optimiert und geplant wurde, wird es in eine lineare Sequenz von Thunks für ein bestimmtes Backend (CPU oder GPU) umgewandelt.
In XLA ist ein Thunk eine Abstraktion einer in sich geschlossenen Arbeitseinheit, die von der Laufzeit ausgeführt wird. Das kann ein kompilierter Kernel-Start, ein bestimmter Vorgang, ein Bibliotheksaufruf, ein Kontrollflusskonstrukt, eine kollektive Kommunikation usw. sein. Eine Thunk-Sequenz stellt die gesamte ausführbare Datei für ein bestimmtes Back-End dar.
Thunk-Emission
Der Prozess der Umwandlung einer geplanten HLO-Berechnung in eine Thunk-Sequenz wird als „Thunk-Ausgabe“ bezeichnet. Dies wird von einer speziellen Emitter-Klasse in jedem Backend übernommen.
Für das GPU-Backend wird dies von IrEmitterUnnested übernommen.
EmitHloComputation durchläuft die geplante Liste der HLO-Anweisungen in einer Berechnung und leitet an eine spezielle Emit...-Methode weiter (z.B.
EmitFusion, EmitConvolutionThunk, EmitWhile). Mit jeder dieser Methoden werden die entsprechenden Thunk-Objekte erstellt und an die Thunk-Sequenz angehängt.
Für das CPU-Backend übernimmt ThunkEmitter diese Rolle und ist ähnlich organisiert. Die endgültige ThunkSequence ist in die CpuExecutable eingebettet.
Jede Anweisung in der Eintragungsberechnung eines HLO-Moduls kann keinem (kTuple, kConstant, ..), einem oder mehreren (z. B. Sortieranweisung) Thunks in der endgültigen Thunk-Sequenz entsprechen.
Befehlspuffer: Ausführung auf der GPU optimieren
Moderne GPU-Hardware ermöglicht es, eine Sequenz von GPU-Vorgängen (Kernel-Starts, Speicherkopien usw.) einmal aufzuzeichnen und dann mit minimalem CPU-Overhead mehrmals wiederzugeben. Dies ist eine wichtige Leistungsoptimierung, insbesondere für Arbeitslasten mit vielen kleinen, schnell startenden Kernels. XLA verwendet Command Buffer als Abstraktion von CUDA- oder HIP-Graphen. Die Kernschnittstelle ist in GpuCommandBuffer definiert.
Ein Befehlspuffer wird in einer Thunk-Sequenz durch CommandBufferThunk dargestellt.
Der Emitter generiert diesen Thunk nicht direkt aus HLO-Anweisungen. Stattdessen wird dies von CommandBufferConversionPass übernommen, das für die ThunkSequence selbst ausgeführt wird.
Der Durchlauf identifiziert zusammenhängende Untersequenzen kompatibler Chunks (z.B. eine Reihe von KernelThunks und GemmThunks). Anschließend wird die gefundene Teilsequenz durch ein einzelnes CommandBufferThunk ersetzt. Der neue Thunk kapselt die Logik der ursprünglichen Thunks als Liste von einfachen CommandBufferCmd-Objekten.
Wenn ein CommandBufferThunk zum ersten Mal in einem bestimmten GPU-Stream ausgeführt wird, wird die Befehlsfolge in einem Hardware-Befehlspuffer „aufgezeichnet“. Bei allen nachfolgenden Ausführungen wird einfach ein einzelner Befehl an die GPU gesendet, um die aufgezeichnete Sequenz „wiederzugeben“. Dadurch wird der CPU-Overhead beim Starten der einzelnen Kernel vermieden.
Ausführbar
Das Endprodukt der XLA-Kompilierungspipeline ist eine eigenständige, plattformspezifische ausführbare Datei. Dieses Objekt kapselt alle Informationen, die zum Ausführen des kompilierten Programms auf einem Zielgerät wie einer CPU oder GPU erforderlich sind. Sie ist die Brücke zwischen dem Compiler und der Laufzeit. Moderne Laufzeiten wie PJRT verwenden Abstraktionen auf etwas höherer Ebene (siehe PjRtExecutable), aber diese umschließen letztendlich eine backend-spezifische ausführbare Datei.
Eine Executable enthält mehrere wichtige Informationen, die während der Kompilierung generiert werden. Die genauen Inhalte variieren je nach Backend, umfassen aber in der Regel Folgendes:
Kompilierter Code: Dies ist der Low-Level-Maschinencode, der auf dem Gerät ausgeführt wird. Bei CPUs sind das in der Regel ein oder mehrere Objektdateien. Bei GPUs ist dies der kompilierte Gerätecode im PTX- oder HSACO-Format, der zur Laufzeit auf die GPU geladen wird.
Ausführungsplan (ThunkSequence): Das Herzstück der Laufzeitlogik. Dies ist eine lineare Sequenz von Thunk-Objekten. Jeder Thunk stellt eine einzelne Arbeitseinheit dar, z. B. das Starten eines Kernels, das Aufrufen einer Bibliotheksfunktion (z. B. cuBLAS) oder die Verarbeitung des Kontrollflusses. Die Laufzeit führt das Programm aus, indem sie diese Sequenz durchläuft.
Speicherlayout (BufferAssignment): Diese wichtige Metadatenkomponente, die vom BufferAssigner erstellt wird, beschreibt das vollständige Speicherlayout für die Berechnung. Sie gibt die Größe jedes Puffers an und legt fest, wie Arbeitsspeicher für Parameter, Ausgaben und temporäre Werte zugewiesen und wiederverwendet wird. Die Laufzeit verwendet dies, um Gerätespeicher zuzuweisen und die richtigen Zeiger an jeden Thunk zu übergeben.
(Optional) HLO-Modul: Zur Fehlerbehebung und Profilerstellung behält die ausführbare Datei oft einen Verweis auf das endgültige, optimierte HloModule bei, aus dem sie kompiliert wurde.
Die Erstellung der endgültigen ausführbaren Datei wird für jedes spezifische Backend vom Compiler orchestriert. Die Methode RunBackend einer Compiler-Implementierung ist der letzte Schritt im Kompilierungsprozess. Dabei werden alle kompilierten Artefakte in einem Executable-Objekt zusammengefasst.
GpuCompiler und CpuCompiler sind jeweils für GPU und CPU vorgesehen.
Wenn ein Nutzer Execute... für eine ausführbare Datei aufruft, verwendet die Laufzeit BufferAssignment, um Speicher zuzuweisen, und ruft dann ThunkSequence auf, um die Vorgänge auf dem Gerät mit dem kompilierten Code zu starten.