This presentation by Google's Robert Hunt provides a comprehensive update on the OpenXLA compiler project, structured around the themes of Performance, Portability, and PyTorch. Hunt details significant performance gains, particularly in the CPU backend, which is now a key focus for host offloading. For portability, the talk introduces major developments, including the evolution of the Stable HLO interface to support symbolic shapes for dynamic workloads and a novel technique to "raise" existing CUDA kernels from HPC applications into the XLA graph, making them portable to hardware like TPUs. The presentation also outlines a new strategy for PyTorch integration, moving towards a fast eager execution mode that minimizes user code changes, with a long-term vision of building a full Just-in-Time (JIT) compiler. Finally, Hunt discusses the project's future, emphasizing a move towards a professional, versioned SDK and addressing key challenges like reducing compiler flag complexity and managing the proliferation of hand-written kernels in favor of more scalable, compiler-driven solutions.
OpenXLA Project Update: Performance, Portability, and PyTorch
0:54 OpenXLA Overview: OpenXLA is a compiler toolchain that takes high-level models from frameworks like JAX, TensorFlow, and PyTorch and optimizes them to run on various hardware backends (CPU, GPU, custom accelerators) through its Stable HLO interface.
3:51 CPU Performance is Now a Priority: The CPU backend, once a "forgotten stepchild," has seen significant performance improvements (e.g., 50x faster parallel compile times). This is crucial for modern use cases like host offloading, where parts of a model run on the host CPU.
4:51 Advancements in Distributed Computing: Shardy, a convenient wrapper for expressing sharding, is being enhanced. Support for MPMD (Multiple Program, Multiple Data) is being added to the infrastructure, allowing for more complex distributed computing patterns beyond SPMD.
5:48 Stable HLO to Support Symbolic Shapes: To improve portability, especially for mobile and other dynamic environments, Stable HLO is evolving to support symbolic shapes. This avoids the need to fork the interface and allows a single, unified standard.
8:00 A New Path for HPC: Raising CUDA to HLO: A novel technique allows existing CUDA kernels to be automatically wrapped in loops and "raised" to HLO. This makes a vast ecosystem of High-Performance Computing (HPC) code, like the "Oceananigans" climate model, portable to the XLA stack and its supported hardware.
9:50 High-Performance Julia Language Integration: A new, comprehensive binding for the Julia programming language allows it to effectively target XLA, unlocking significant performance gains for scientific models written in Julia.
11:09 A New Strategy for PyTorch: The project is moving away from the old lazy_tensor binding. The new plan, requested by Meta, involves creating a fast eager mode that allows PyTorch code to run on XLA backends (like TPUs) with minimal changes, aiming for performance within 2-3x of fully compiled code but with no compilation overhead.
13:23 The Ultimate Goal: A PyTorch JIT Compiler: The long-term vision is to create a full Just-in-Time (JIT) compiler for PyTorch, similar to a Java JIT, that can perform dynamic optimizations like recompilation and specialization to achieve maximum performance.
14:16 Professionalizing with a Versioned SDK: To move beyond a "compiler of the day" approach, OpenXLA will be delivered as a professional SDK with structured, versioned releases. This will provide stability and make it easier for users and partners to manage dependencies and deprecations.
16:56 Key Challenge: Reducing Flag Complexity: Hunt highlights the need to reduce the large number of compiler flags. These flags push complexity onto the user, making it difficult to achieve optimal performance and triage issues.
18:01 Key Dilemma: Containing Kernel Proliferation: The speaker expresses concern over the explosion of custom, hand-written kernels (e.g., Triton, Mosaic). While they provide short-term performance wins, he argues they are equivalent to assembly code—not scalable or portable—and that solving these optimization problems is fundamentally the compiler's job.