Diese Präsentation erläutert die Integration des oneAPI-Backends in XLA, um JAX-Workloads auf Intel-GPUs auszuführen, am Beispiel der Battlemage-Architektur (Arc Pro B70). Es werden die Hardware-Komponenten, das oneAPI-Software-Stack (oneMKL, oneDNN, oneCCL), die XLA-Kompilierungspipeline (inklusive oneDNN-Pfad, MLIR-Codegen und experimentellem Triton-Support), die Runtime-Ausführung über Level Zero und SYCL sowie die Multi-GPU-Skalierung via oneCCL detailliert beleuchtet. Zudem werden der PJRT-basierte JAX-Plugin-Status und die Roadmap für Upstream-Optimierungen vorgestellt.
Wichtigste Punkte & Zeitstempel
0:09 Intel Battlemage-Architektur: Vorstellung der Arc Pro B70 Workstation-GPU basierend auf der Xe2-Architektur mit 32 GB GDDR6 VRAM, Xe-Cores (je 8 Vektor- und Matrix-Engines) und 256 KB L1/lokalem Speicher pro Core.
0:26 oneAPI-Framework: Einbindung des plattformübergreifenden Standards mit Kernbibliotheken wie oneMKL (Mathematik), oneDNN (Deep Learning) und oneCCL (kollektive Kommunikation) für heterogene Hardware.
0:40 XLA-Kompilierung für GEMM und Convolutions: Umleitung von HLO-Dot- und Convolutions-Operationen über den oneDNN-Pfad in einzeln gefuste SYCL-Kernels via JIT-Kompilierung, wodurch Overhead eliminiert wird.
6:26 MLIR-Codegen & SPIR-V: Verarbeitung von Elementwise- und Reduktions-Operationen über MLIR-Transformatoren, Zielselektion zu SPIR-V-Intrinsics und JIT-Übersetzung in native ISA durch den Intel Graphics Compiler via Level Zero.
8:32 Triton-Codegen: Dreiphasiger Kompilierungspfad (Triton IR, TTIG-Dialekt, Triton-GenOps) mit integriertem Autotuner für das Benchmarken und Optimieren von GEMM-Tiling-Konfigurationen.
10:57 Runtime-Infrastruktur: Initialisierung des SYCL-Executors und Laden von SPIR-V-Binärdateien in Level-Zero-Module zur asynchronen Ausführung auf GPU-Streams.
12:09 Multi-GPU-Skalierung: Nutzung von oneCCL zur Bereitstellung hardwarebeschleunigter Kommunikationsprimitiven (All-Reduce, Broadcast) für verteiltes Training.
13:27 JAX-Integration: Bereitstellung eines PJRT-basierten oneAPI-Plugins mit drei Wheels (jaxlib, jax-oneapi-plugin, jax-oneapi-pjrt) für validierte Modelle wie Gemma 2, SDXL und GPT-2.
14:14 Roadmap & Upstream-Status: Laufende Upstream-Integration von SPIR-V-Codegen und PJRT; geplante Features umfassen PTI-Profiling, Triton-Fusionskernel, SYCL-Graph-Command-Buffer und Unterstützung zukünftiger Rechenzentrum-GPUs wie Crescent Island.
Glossar
XLA (Accelerated Linear Algebra): Ein Domain-spezifischer Optimierungscompiler für lineare Algebra, der Berechnungs-Graphen aus Frameworks wie TensorFlow oder JAX in maschinennahen Code übersetzt.
JAX: Ein Python-Framework für Hochleistungs-Numerik und maschinelles Lernen, das automatische Differenzierung und JIT-Kompilierung unterstützt.
Battlemage / Arc Pro B70: Intels diskrete Grafikprozessor-Generation (Xe2-Architektur) für professionelle Workstations, optimiert für lokale KI-Inferenz und Grafikberechnungen.
oneAPI: Ein offener, plattformübergreifender Programmierstandard von Intel zur einheitlichen Softwareentwicklung über CPUs, GPUs und FPGAs hinweg.
oneMKL / oneDNN / oneCCL: Intel-Bibliotheken für optimierte mathematische Routinen (Math Kernel Library), Deep-Neural-Network-Primitive bzw. kollektive Multiknoten-Kommunikation.
SYCL: Eine C++-basierte Programmierschnittstelle für heterogenes Computing, die die Entwicklung von Code für verschiedene Hardware-Beschleuniger ermöglicht.
MLIR (Multi-Level Intermediate Representation): Eine modulare Compiler-Infrastruktur, die strukturierte Zwischenrepräsentationen auf verschiedenen Abstraktionsebenen ermöglicht.
SPIR-V: Ein offener Standard-Zwischencode (Bytecode) für parallele Shader und Kernel, der von Khronos-Standards und Intel Level Zero genutzt wird.
Level Zero: Eine Low-Level-Hardwareschnittstelle von Intel zur direkten, effizienten Steuerung von GPU-Ressourcen, Queues und Speicher-Managements ohne Treiber-Overhead.
PJRT: Eine einheitliche Laufzeitschnittstelle (Plugin Device Runtime) in JAX, die die Anbindung verschiedener Hardware-Beschleuniger standardisiert.
Triton: Ein von OpenAI entwickelter Compiler und eine Programmiersprache zur automatischen Erstellung hochleistungsfähiger GPU-Kernel.