Dieser Vortrag von Bavani (Software Engineer bei Intel und OpenXLA-Contributor) behandelt die Integration des Intel oneAPI-Backends in XLA, um die Ausführung von JAX-Workloads auf Intel-GPUs zu ermöglichen. Der Schwerpunkt liegt auf der Hardware-Architektur der Intel Arc Pro B70 (Xe2-Architektur) sowie auf dem Zusammenspiel des oneAPI-Software-Stacks (1DNN, 1MKL, 1CCL) mit der Kompilierungs- und Laufzeitpipeline von XLA. Es wird erläutert, wie HLO-Operationen über spezialisierte Pfade – 1DNN für Matrixmultiplikationen und Faltungen, MLIR/SPIR-V für Elementwise- und Reduktions-Ops sowie eine dreistufige Triton-CodeGen-Pipeline – abgesenkt werden. Abschließend werden die Multi-GPU-Skalierung über 1CCL, das PJRT-basierte JAX-Plugin-System und der aktuelle Status im Upstream-Repository samt zukünftiger Roadmap (u. a. SYCL Graph Command Buffers und Crescent Island Support) vorgestellt.
Wichtigste Punkte & Zeitstempel
0:00 oneAPI-XLA-Integration: Überblick über die Einbindung des oneAPI-Backends in XLA zur Ausführung von JAX auf Intel-GPUs.
0:56 Battlemage-GPU-Architektur: Die Workstation-GPU Arc Pro B70 basiert auf der Xe2-Architektur, bietet 32 GB GDDR6-VRAM für lokale Inferenz großer KI-Modelle und unterstützt Multi-GPU-Deployments unter Linux.
1:40 Xe2-Core-Spezifikationen: Ein Xe-Stack kombiniert 32 Xe-Cores; jeder Core besitzt 8 Vektor-Engines (ALUs), 8 Matrix-Engines (XMX) sowie 256 KB dynamisch aufteilbaren L1- und Shared Local Memory.
2:26 oneAPI-Software-Stack: Bietet die Voraussetzungen für das JAX-Plugin über plattformübergreifende Bibliotheken: 1MKL (BLAS/FFT), 1DNN (Deep-Learning-Basisfunktionen) und 1CCL (Kollektive Kommunikation).
4:07 HLO-Kompilierungspipeline: HLO-Faltungen und Matrixmultiplikationen (Dots) laufen über 1DNN oder optional Triton-Code-Generierung, während Elementwise- und Reduktions-Ops die MLIR-Pipeline durchlaufen.
4:46 1DNN-Absenkung & Epilog-Fusion: Der GEMM-Rewriter wandelt Dots in Custom Calls um; der Thunk-Emitter erzeugt Thunks, die zur Laufzeit von 1DNN als gefügte SYCL-Kernels mit Epilog-Post-Ops JIT-kompiliert werden.
6:26 MLIR & SPIR-V-CodeGen: Fusionierte Knoten werden in MLIR-Dialekte (XLA, SCF, Arith) übersetzt, optimiert, über LLVM in SPIR-V-Binärdateien konvertiert und mittels Level Zero / Intel Graphics Compiler (IGC) in nativen ISA-Code übersetzt.
8:32 Triton-CodeGen-Pipeline: Dreistufiger Ablauf vom geräteagnostischen Triton-IR über den hardwarenahen TTIG-Dialekt (Triton Intel GPU) und Triton GenOps bis hin zur SPIR-V-Generierung.
9:58 Triton-Autotuning: Ein GEMM-Fusion-Pass entscheidet per Autotuner zwischen 1DNN und Triton und wählt durch Tiling-Benchmarking zur Kompilierzeit die optimale Konfiguration aus.
10:57 SYCL-Laufzeitausführung: Der SYCL-Executor verwaltet Kernel- sowie Speicherlebenszyklen, lädt SPIR-V-Binärdateien in Level-Zero-Module und stellt SYCL-Ressourcen für 1DNN-Thunks bereit.
12:09 Multi-GPU-Ausführung via 1CCL: 1CCL-Communicators verwalten Device-Ranks und führen GPU-Kollektivoperationen (z. B. AllReduce, Broadcast) durch direkte Einbindung der SYCL-Queue aus dem XLA-Stream aus.
13:21 JAX PJRT-Plugin & Modell-Validierung: Die Anbindung erfolgt über ein dreiteiliges Wheel-Set (jaxlib, jax-oneapi-plugin, jax-oneapi-pjrt); erfolgreich getestete Modelle umfassen MaxText (Llama 2 2B), MaxDiffusion (SDXL), BERT, ViT und Whisper Tiny.
14:15 Upstream-Status & Roadmap: Upstream-XLA unterstützt bereits StreamExecutor und SPIR-V-Codegen; in Entwicklung befinden sich PTI-Profiling-Tools, Scaleout-Multi-GPU, Low-Precision-Datentypen, SYCL Graph Command Buffers und Crescent Island Support.
Glossar
XLA (Accelerated Linear Algebra): Ein spezialisierter Domänen-Compiler für Machine Learning (u. a. für TensorFlow und JAX), der mathematische Graphen in hochoptimierten Maschinencode für Beschleuniger übersetzt.
JAX: Ein von Google entwickeltes Python-Framework für wissenschaftliches Rechnen und Machine Learning mit automatischer Differenzierung und XLA-Kompilierung.
oneAPI: Ein von Intel initiiertes, offenes und plattformübergreifendes Programmiermodell, das eine einheitliche Entwicklung über CPUs, GPUs und FPGAs hinweg ermöglicht.
SYCL: Ein herstellerunabhängiger C++-Standard für heterogene Systeme, auf dem der oneAPI-Laufzeitstack aufbaut.
1DNN (oneAPI Deep Neural Network Library): Eine hochoptimierte Bibliothek von Intel, die grundlegende Bausteine für neuronale Netze (z. B. Faltungen, Aktivierungsfunktionen, Matrixmultiplikationen) bereitstellt.
1MKL (oneAPI Math Kernel Library): Eine Bibliothek mit mathematischen Routinen für lineare Algebra, Fast Fourier Transformations (FFT) und Zufallsgeneratoren.
1CCL (oneAPI Collective Communications Library): Eine Bibliothek für verteilte Trainingsinfrastrukturen, die primitive Kommunikationsmuster (z. B. AllReduce) über mehrere GPUs hinweg beschleunigt.
MLIR (Multi-Level Intermediate Representation): Ein flexibles Compiler-Framework unter dem LLVM-Schirm, das Programmcode schrittweise über verschiedene Abstraktionsebenen (Dialekte) transformiert.
SPIR-V: Ein plattformunabhängiges binäres Zwischenformat für Grafik- und Rechen-Kernels, das von Treibern direkt in nativen GPU-Code kompiliert werden kann.
Level Zero: Eine hardwarenahe, extrem leichtgewichtige Treiberschnittstelle von Intel zur Steuerung von GPU-Ressourcen und zur JIT-Kompilierung.
PJRT (Plugin Joint Runtime): Die standardisierte Laufzeitschnittstelle in XLA, über die verschiedene Hardware-Backends nahtlos an JAX angebunden werden.
Triton: Eine Open-Source-Programmiersprache und ein Compiler, mit denen hochleistungsfähige GPU-Kernels (insbesondere für Matrixmultiplikationen) vereinfacht geschrieben werden können.
TTIG (Triton Intel GPU Dialect): Eine spezielle MLIR-Zwischenrepräsentation für Intel-Hardware, die Matrix-Engines und Subgruppenstrukturen im Triton-Compiler abbildet.
Xe2-Architektur / XMX (Xe Matrix eXtensions): Intels GPU-Architektur der zweiten Generation mit dedizierten Hardware-Beschleunigern (Matrix-Engines) für KI-Berechnungen.