We accidentally built an LLVM compiler for Jax

In the realm of quantum computing, a team faced an unexpected outcome while developing a compiler called Catalyst for the PennyLane quantum software library. Their objective was to optimize hybrid quantum-classical workflows using MLIR, a compilation framework. To achieve this, they harnessed the capabilities of JAX, a Python library known for its ability to trace through functions and capture computational graphs, along with its support for NumPy and SciPy APIs.

Initially, the team aimed to capture native Python control flow and support dynamically-shaped arrays within Catalyst's quantum workflows. However, while wiring JAX to feed the quantum pipeline, they discovered an accidental development: pure JAX NumPy code and native Python control flow could be fed directly into Catalyst without the need for quantum instructions.

Astonishingly, Catalyst bypassed the XLA compiler backend, lowered JAX representations to standard MLIR, and compiled it using LLVM, incorporating backpropagation support.

This accidental creation marked the birth of an MLIR compilation pipeline for JAX, an unexpected outcome. The team's motivation stemmed from their desire to leverage existing classical compilation tooling, such as LLVM, and incorporate quantum support without starting from scratch. They aimed to represent quantum programs with structure, permitting the natural integration of array manipulation, quantum instructions, loops, and if statements.

This approach also supported dynamically-shaped arrays and maintained the support for quantum autodifferentiation.

The decision to build the compiler infrastructure on MLIR stemmed from the need to compile the entire hybrid execution graph into a single, standalone executable, executing as closely to the bare metal as possible. By compiling the classical JAX code directly to LLVM IR, Catalyst could utilize standard LLVM passes, such as Enzyme, for backpropagation without the complexities of integrating them directly into XLA's internal pipeline.

This approach leveraged standard LLVM optimizations while still targeting specific quantum instructions when required.