PyTorch-Team erklärt MLIR-Dataflow-Backend, das KTIR in Code für IBM Spyre übersetzt

KTIR: der Vertrag zwischen Inductor und Backend

Laut dem Beitrag entscheidet der Inductor von torch-spyre, wie die Arbeit jedes Kernels auf die Kerne von Spyre verteilt wird, und verwaltet, wo Tensoren liegen: im On-Chip-Scratchpad oder im Off-Chip-Speicher. Der Backend-Compiler übersetzt diese Entscheidungen anschließend in etwas, das die Hardware ausführen kann. Der Vertrag zwischen beiden ist KTIR, kurz für Kernel Tile IR, eine MLIR-Repräsentation, die auf einer datenparallelen Sicht auf den Beschleuniger beruht.

In dieser Sicht kann jede Compute-Tile jeden Speicher sehen, und ein einzelner logischer Tensor kann über die Scratchpads mehrerer Kerne verteilt sein. KTIR ist in RFC 0682 beschrieben und in torch-spyre/ktir-mlir-frontend implementiert. Die Autoren sagen, sie wollten es ausdrucksstark genug machen, um Abbildungen aus zwei unterschiedlichen Quellen zu erfassen: aus Inductor und aus handgeschriebenen Spyre-Kernels in Sprachen wie Triton oder Helion.

Das Hardwarebild hinter KTIR ist bewusst schlicht: Kerne mit einer Recheneinheit und eigenem Scratchpad, verbunden über ein On-Chip-Interconnect, das auch den Off-Chip-Speicher erreicht. Der Beitrag betont, dass dies recht genau zu Spyre passt: 32 Kerne auf einem Ring, jeder mit 2 MB Scratchpad, dahinter der Off-Chip-Gerätespeicher. Daten werden in Sticks bewegt, festen 128-Byte-Blöcken, die 64 fp16-Werte fassen. Details wie FIFOs, Load/Store-Einheiten und entkoppelte Ausführungs-Engines bleiben absichtlich außen vor, denn KTIR beschreibt, was jede Tile tut und wo ihre Daten liegen, nicht wie die Daten dorthin gelangen.

Wie KTIR aufgebaut ist

KTIR verwendet Upstream-MLIR-Dialekte für Arithmetik, Reduktionen und Kontrollfluss wieder: arith, math, linalg, scf, tensor und memref. Ein eigener Dialekt, ktdp, ergänzt explizite Speicher und explizite Zugriffe. Zu seinen Kernoperationen zählen get_compute_tile_id, construct_memory_view, construct_distributed_memory_view, construct_access_tile, construct_indirect_access_tile, load und store sowie inter_tile_produce und inter_tile_reduce für die Kommunikation zwischen Tiles, etwa einen All-Reduce über Teilsummen.

Der Beitrag geht eine elementweise Addition über fp16-Tensoren der Größe 96×64 durch, die auf 32 Compute-Tiles mit je drei Zeilen aufgeteilt wird. Views legen fest, wie Speicher interpretiert wird, Access-Tiles bestimmen, welche Koordinaten eine Tile berührt, Loads und Stores sind die einzige Datenbewegung, und linalg übernimmt die Arithmetik. Die Arbeitsaufteilung steht bereits durch grid = [32] und die Arithmetik auf der Tile-ID fest, und das Backend soll sie einhalten, statt sie zu hinterfragen.

Ein zweiter, deutlich kleinerer Dialekt, spyreop, stellt Spyre-Intrinsics als gewöhnliche MLIR-Operationen bereit. Er deckt elementweise Funktionen und Aktivierungsfunktionen wie exp, gelu, sigmoid, silu, softplus, sqrt, rsqrt, reciprocal und realdiv ab, Bausteine für Normalisierung, Ganzzahl- und Adressarithmetik einschließlich idx32toaddr sowie Vergleich und Auswahl. Außerdem definiert er df16, IBMs 16-Bit-Format DLFloat, das seine Bits im Verhältnis 1-6-9 auf Vorzeichen, Exponent und Mantisse verteilt, gegenüber 1-5-10 bei IEEE-Halbpräzision.

Der Dataflow-Scheduler und SpyreCode

Ein neuer Open-Source-Dataflow-Scheduler liest KTIR ein und macht aus der Arbeit jeder Tile eine Dataflow-Pipeline. Er zerlegt den Kernel in Programme und beginnt mit einer einfachen Pipeline aus Laden, Rechnen und Speichern. Anschließend leitet er jeden Transfer durch die Speicherhierarchie, die eine Architekturbeschreibung vorgibt, und erweitert die Pipeline zu einem verschachtelten, doppelt gepufferten Multi-Instanz-Schedule.

Der Beitrag erklärt, dass Spyre eine Dataflow-Maschine mit entkoppeltem Zugriff und entkoppelter Ausführung ist: Loads, Stores und Berechnungen laufen auf getrennten Engines, die über FIFOs kommunizieren. Der Scheduler startet mit der einfachsten zulässigen Pipeline und verfeinert sie Transformation für Transformation, jede davon ein gut lesbares MLIR-Rewrite.

Das Ergebnis des Schedulers ist DataflowIR, eine MLIR-Zwischenrepräsentation. Die gerätespezifischen Stufen von Spyre nutzen sie, um SpyreCode zu erzeugen, bestehend aus einem symbolischen SpyreBinary und einem Job-Plan. Die Spyre-Runtime bereitet dies einmal vor und startet es vielfach, wobei sie beim Start Tensoradressen und Shapes einsetzt. KTIR, das ScheduleIR des Schedulers und DataflowIR sind allesamt gewöhnliches MLIR auf Basis von Upstream-Dialekten, und der Scheduler samt seinen Dialekten ist unter Apache-2.0 quelloffen. Die gerätespezifische Absenkung für die Codegenerierung beginnt an der Grenze zu DataflowIR.

Was das für Entwickler bedeutet

Der Beitrag bildet den Mittelteil einer dreiteiligen Serie. Die früheren Beiträge behandeln die Erweiterungen von Inductor für Layout, Tiling und Scratchpad-Platzierung sowie die Frage, wie die Runtime kompilierte Artefakte startet. Für Entwickler bedeutet das Design, dass die parallele Zerlegung einmal im Frontend festgelegt und nicht vom Backend neu erschlossen wird. Weil die Schnittstellen MLIR-IRs statt Bibliotheksaufrufe sind und der Scheduler unter Apache-2.0 lizenziert ist, lassen sich die Stufen einsehen und wiederverwenden. KTIR soll zudem handgeschriebene Kernels aus Triton oder Helion neben der Inductor-Ausgabe aufnehmen.