Dieser Vortrag von Tomas behandelt Mosaic, einen Compiler, der als Bindeglied zwischen dem Programmiermodell Palace und zielspezifischen Backends (wie Tensor Cores und Sparse Cores) für TPUs (Tensor Processing Units) dient. Mosaic verarbeitet MLIR (Multi-Level Intermediate Representation) mit Kern- und TPU-spezifischen Dialekten und ist primär für die Verwaltung, Transformation und Optimierung von Vektorkonstrukten (Vector Layouts) verantwortlich. Der Vortrag erläutert die vier Komponenten von Vektorlayouts (Bitbreite, Tiling, Offsets mit Replikation und implizite Dimensionen), die Handhabung von Layout-Konvertierungen (Relayouts) sowie die drei zentralen Compiler-Phasen: Infer Vector Layout, Einfügen von Relayouts und Apply Vector Layout (Unrolling). Abschließend werden Unterschiede zu Sparse-Core-Architekturen sowie die Verarbeitung gepackter Sub-32-Bit-Datentypen (z. B. FP8, BF16) erörtert.
Wichtigste Punkte & Zeitstempel
0:16 Mosaic-Rolle im Compiler-Stack: Mosaic agiert als Zwischenschicht-Compiler zwischen Palace und TPU-Backends, verarbeitet MLIR-Eingaben und führt hardwarenahe Transformationen für Tensordaten durch.
1:16 TPU-Vektorregister (VREGs): TPU-Vektorregister sind zweidimensionale Strukturen aus 8 Sublanes und 128 Lanes; Tensoren müssen für Operationen zwingend in diese Register geladen werden.
3:43 Vektor-Layouts (Vector Layouts): Layouts definieren, wie logische N-dimensionale Tensordaten in physischen Registern repräsentiert werden, gesteuert über vier Kernkomponenten.
4:39 Tiling-Mechanismus: Tiling bestimmt die Aufteilung von Tensoren in Register-Slices (z. B. 8x128 oder 4x256), bedingt durch Speicherstrukturen von XLA und zur Minimierung von Padding-Overhead.
6:46 Offsets und Replikation: Offsets verwalten Padding-Daten aus Slicing-Operationen oder nicht ausgerichteten Speicherzugriffen; Replikations-Flags ermöglichen identische Werte über Dimensionen hinweg für effiziente Reduktionen.
9:41 Implizite Dimensionen: Historisch für komplexe Operationen (Reduktion gefolgt von Reshape) eingeführt, erlauben implizite Dimensionen das Vermeidung kostspieliger Transpositionen zwischen Sublanes und Lanes.
11:51 Relayout-Operationen: Da binäre Operationen identische Layouts erfordern, fügt Mosaic automatische Konvertierungen (Relayouts) ein, deren Aufwand optimiert werden muss.
13:48 Die drei Hauptphasen des Compilers: Der Optimierungsprozess gliedert sich in Infer Vector Layout (globale/lokale Layout-Zuweisung), das Einfügen von Relayouts und Apply Vector Layout (Unrolling der Operationen für das Ziel-Backend).
16:36 Sparse Core (Sparsore) Unterstützung: Sparse Cores nutzen kleinere Registerstrukturen (z. B. 8x32 oder 16x32 Bit) und verlangen separate Layout-Logiken für nicht-zusammenhängend gepackte Datenelemente.
19:04 C++-Nutzung und Layout-Heuristiken: Fortgeschrittene Anwender nutzen Mosaic direkt via C++, um Python-Schichten zu umgehen; bestehende Layout-Heuristiken sind teilweise lokal und werden kontinuierlich verbessert.
21:41 Gepackte Datentypen (Packed Types): Sub-32-Bit-Typen wie FP8, BF16 oder INT16 packen mehrere Elemente in ein Vektorregister, wodurch die logische Registerkapazität entsprechend multipliziert wird.
Glossar
Mosaic: Ein spezialisierter Compiler im TPU-Software-Stack zur Steuerung von Speicher- und Register-Layouts zwischen High-Level-APIs und Hardware-Backends.
Palace: Ein High-Level-Programmiermodell für TPUs (oft in Python oder C++), das MLIR-Code für Mosaic generiert.
MLIR (Multi-Level Intermediate Representation): Ein modulares Compiler-Infrastruktur-Framework, das als Eingabeformat für Mosaic dient und domänenspezifische Dialekte verwendet.
TPU (Tensor Processing Unit): Ein von Google entwickelter, kundenspezifischer Prozessor (ASIC), der speziell für die Beschleunigung von Deep-Learning- und Tensor-Berechnungen optimiert ist.
VREG (Vector Register): Ein Hardware-Vektorregister auf der TPU, das als 2D-Matrix (z. B. 8 Sublanes mal 128 Lanes) organisiert ist.
Tiling: Die Zerlegung großer, mehrdimensionaler Tensoren in kleinere Blöcke (Tiles), um sie effizient in Hardware-Registern zu verarbeiten.
Padding / Offsets: Unnütze oder mit Platzhalterdaten gefüllte Speicherbereiche in Registern, die durch Slicing oder unaligned Memory Loads entstehen.
Unrolling: Ein Compiler-Optimierungsschritt, bei dem abstrakte Operationen oder Schleifen in explizite, sequenzielle Maschineninstruktionen aufgelöst werden.
Tensor Core / Sparsore: Spezifische Hardware-Ausführungseinheiten (Backends) der TPU für dichte bzw. dünnbesetzte (sparse) Matrixoperationen.
FP8 / BF16 / INT16: Numerische Datenformate mit reduzierter Bitbreite (8-Bit-Gleitkomma, 16-Bit-Bfloat, 16-Bit-Ganzzahl) zur Steigerung von Rechen- und Speichereffizienz.
Sublanes und Lanes: Die physischen Dimensionen der internen TPU-Vektorregisterstruktur, die die parallele Datenverarbeitung steuern.