SciPy 2024

PIXIE: Blending Just-in-time and Ahead-of-time compilation in scientific Python applications
07-10, 16:05–16:35 (US/Pacific), Room 315

Speeding up Python code traditionally involves the use of Just-In-Time (JIT) or Ahead-Of-Time (AOT) compilation. There are tradeoffs to both approaches, however. As part of the Numba project's aim to create a compiler toolkit, the PIXIE project is being developed. It offers a multiple-language consuming, extensible toolchain, that produces AOT compiled binary extension modules. These PIXIE based extension modules contain CPU-specific function dispatch for AOT use and also support something similar to Link-Time-Optimization (LTO) for use in situations such as JIT compilation and/or cross module optimization. PIXIE modules are easy to load and call from Python, and can be inlined into Numba JIT compilation, giving Python developers access to the benefits of both AOT and JIT.


Increasing the performance of general numerical applications in Python has traditionally involved creating Ahead-Of-Time (AOT) compiled binary extension modules, with compiler options and target microarchitecture features locked in at compilation time. With the advent of Just-In-Time (JIT) compilers, like Numba, this performance can be gained at runtime with greater opportunity for specialization, but incurs the runtime cost of JIT compilation. This talk showcases PIXIE, a project that lets you create high performance binary extension modules that have automatic CPU-specialized function dispatch and can also take part in further LTO and LTCG based optimization when consumed by LLVM-based compilers. The PIXIE project comprises an extensible toolchain and an output format that focuses on leaving optimization in the hands of the end user. It aims to make it easy to extend the toolchain to new languages, all that’s needed is a frontend that generates LLVM IR. Further, the output format is a true binary extension module that is designed for easy introspection and use, but is also extensible for your use cases!

To demonstrate the power of PIXIE, we will do a demonstration of using multiple input languages to create AOT compiled PIXIE based binary extension modules. These modules will then be “shipped” and later used in some user defined code. The functions in the modules will be called as AOT symbols and also take part in JIT compiler driven LTO. Specifically, we’ll take a simple Newton-Raphson solver written in Python and use a Numba extension to PIXIE to AOT compile this into a binary extension module. We’ll then take some objective functions written in C and Cython and AOT compile these into another PIXIE based binary extension module. Finally, in some Numba JIT driven code, the functions in the binary extension modules will be used in both JIT and AOT compilation contexts. This will demonstrate the optimization composability enabled by PIXIE, multiple languages across multiple modules are combined with a single tool, cross module optimization takes place, and compilation is shown as a spectrum of choices in terms of execution latency and performance.

Stanley Seibert is the Senior Director of Community Innovation at Anaconda, where he manages many of Anaconda's open source project maintainers. He has been a contributor to the Numba Python compiler project for 10 years, and has worked with a wide range of technologies on past projects, including GPU computing, sparse linear algebra for graph analytics, and the MLIR compiler framework. Stanley was previously trained as a high-energy physicist, working on experimental software and hardware for the study of neutrinos and dark matter.