Dieser Vortrag von Bavani (Softwareingenieurin bei Intel) beschreibt die Integration des oneAPI-Backends in den XLA-Compiler (Accelerated Linear Algebra), wodurch JAX-Workloads auf Intel-GPUs der Battlemage-Architektur ausgeführt werden können. Die Präsentation beleuchtet die Hardware-Spezifikationen der ARC-Workstation-GPUs, den oneAPI-Software-Stack, die zweigleisige Kompilierungspipeline (über oneDNN und MLIR/SPIR-V sowie experimentellen Triton-Codegen), die Laufzeitausführung via Level-Zero-Schnittstelle, die Multi-GPU-Skalierung mittels oneCCL sowie die Struktur des PJRT-basierten JAX-Plugins.
Wichtigste Punkte & Zeitstempel
0:09 Intel Battlemage GPU-Architektur: Vorstellung der professionellen Workstation-GPU ARC Pro B70 auf Basis der XE2-Architektur mit 32 GB GDDR6 VRAM, 32 XC-Cores (jeweils 8 Vektor- und Matrix-Engines) und einem dynamisch aufteilbaren 256 KB L1-Cache.
0:26 oneAPI Software-Stack: Erläuterung von oneAPI als plattformübergreifendes Programmierframework, das durch Bibliotheken wie oneMKL (Mathematik), oneDNN (Deep Learning) und oneCCL (Kollektivkommunikation) die plattformunabhängige Hardware-Nutzung ermöglicht.
0:40 HLO-Kompilierungspipeline: Trennung der Codegenerierung: Convolutions und Dot-Operationen werden über GEM-Custom-Calls an oneDNN übergeben, während elementweise Operationen und Reduktionen über MLIR-Infrastruktur zu SPIR-V-Binärdateien gewandelt werden.
0:8:32 Experimenteller Triton-Codegen: Dreiphasige Triton-Pipeline (Triton IR, TTIG-Dialekt und Triton GenOps), die mithilfe von Autotuning-Mechanismen optimale Kachelungskonfigurationen (Tiling Configs) für rechenintensive Matrixoperationen ermittelt.
0:10:57 Laufzeitausführung und Level-Zero: Einbindung des SYCL-Executors zur Speicher- und Kernelverwaltung, Laden von SPIR-V-Binärdateien in die Level-Zero-Laufzeitumgebung und JIT-Kompilierung (Just-in-Time) in die native ISA über den Intel Graphics Compiler (IGC).
0:12:09 Multi-GPU-Skalierung mit oneCCL: Nutzung der oneAPI Collective Communications Library zur Abwicklung verteilter Trainingsvorgänge und Synchronisation über hardwarebeschleunigte Kollektivprimitive (z. B. All-Reduce, Broadcast).
0:13:27 JAX-Integration und Validierung: Bereitstellung eines PJRT-basierten oneAPI-Plugins, aufgeteilt in drei Python-Wheels (jaxlib, oneAPI-Plugin, oneAPI-PJRT), mit erfolgreichen Validierungen für Modelle wie Gemma 2, Stable Diffusion XL, BERT und GPT-2.
0:14:37 Zukunftsarbeit und Upstream-Roadmap: Laufende Upstream-Integrationen in XLA/JAX, darunter Profiling-Tools (PTI), erweiterte Multi-GPU-Skalierung, graphbasierte SYCL-Kommandozeilenpuffer sowie zukünftige Unterstützung für Intel Data Center GPUs (z. B. Crescent Island).
Glossar
XLA (Accelerated Linear Algebra): Ein Domain-spezifischer Compiler für lineare Algebra, der Rechengraphen in Machine-Learning-Frameworks (wie JAX und TensorFlow) optimiert und beschleunigt.
JAX: Ein Hochleistungs-Framework für maschinelles Lernen und numerische Berechnungen, das automatische Differenzierung und JIT-Kompilierung unterstützt.
Battlemage / XE2: Die zweite Generation von Intels Grafikprozessor-Architektur (Xe High Performance Graphics) für Gaming und professionelle Workstations.
GDDR6 VRAM: Graphics Double Data Rate 6 Video Random Access Memory – ein extrem schneller Grafikspeicher, der große KI-Modelle und Daten lokal puffert.
XC-Core: Der grundlegende Rechenbaustein der Intel-Xe-Architektur, der Threads steuert und Vektor- sowie Matrix-Berechnungen durchführt.
oneAPI: Ein offenes, herstellerübergreifendes Software-Entwicklungsframework von Intel zur einheitlichen Programmierung verschiedener Hardware-Architekturen (CPUs, GPUs, FPGAs).
oneMKL (oneAPI Math Kernel Library): Eine optimierte Softwarebibliothek für mathematische Kernroutinen wie lineare Algebra und Fast-Fourier-Transformationen (FFTs).
oneDNN (oneAPI Deep Neural Network Library): Eine leistungsoptimierte Bibliothek mit Grundbausteinen für neuronale Netze (wie Faltungen, Aktivierungsfunktionen und Normalisierungen).
oneCCL (oneAPI Collective Communications Library): Eine Bibliothek für verteilte Systemkommunikation, die synchrone Operationen über mehrere GPUs hinweg beschleunigt.
HLO (High-Level Optimizer / High-Level Operations): Die abstrakte Zwischenstufe (Intermediate Representation) im XLA-Compiler, auf der mathematische Operationen optimiert werden.
MLIR (Multi-Level Intermediate Representation): Eine modulare Compiler-Infrastruktur, die verschiedene Abstraktionsebenen (Dialekte) nutzt, um Code schrittweise für spezifische Hardware zu übersetzen.
SPIR-V: Ein offener, standardisierter Zwischencode-Standard der Khronos Group für parallele Compute-Kernel und Shader.
SYCL: Eine offene C++-Standarderweiterung für heterogenes Computing, mit der Entwickler plattformübergreifenden Code für CPUs und Beschleuniger schreiben können.
Level Zero: Eine schlanke, hardwarenahe Low-Level-Programmierschnittstelle (API) zur direkten Steuerung von Intel-Grafikprozessoren mit minimalem Overhead.
ISA (Instruction Set Architecture): Der konkrete maschinennahe Befehlssatz der zugrundeliegenden Hardware, in den der Programmcode übersetzt wird.
PJRT: Eine standardisierte Plug-in-Schnittstelle in JAX, die die Anbindung unterschiedlicher Hardware-Beschleuniger (wie GPUs oder TPUs) vereinheitlicht.
Triton: Ein Python-basiertes Programmiermodell und ein Compiler, mit dem effiziente GPU-Kernel erstellt werden können, ohne Low-Level-Code schreiben zu müssen.
Autotuning: Ein automatisierter Prozess während der Kompilierung, bei dem verschiedene Ausführungsvarianten getestet werden, um die performanteste Konfiguration auszuwählen.