We accidentally built an LLVM compiler for JAX | josh iza.ac

Sometimes when you build software, you shave a yak so deeply that you end up with a surprisingly<br>nice sweater (while the yak, presumably, wonders why it is suddenly rather drafty).

For the past while, my team and I have been building a quantum compiler called<br>Catalyst for the<br>quantum software library PennyLane. Our goal was relatively<br>straightforward: optimize large, hybrid quantum-classical workflows in a scalable manner using<br>MLIR. To do this, we needed a fast, robust way to capture classical Python processing<br>(including NumPy and associated scientific libraries) and represent it in our intermediate representation.

We chose JAX, for a couple of reasons: its ability to trace through Python functions and capture the<br>computational graph1, its support<br>of the NumPy and SciPy APIs with relatively good coverage, and the fact that it already lowers down<br>to MLIR. We also aimed to address a couple of quality of life improvements, such as the ability to<br>capture native Python control flow2, and support for<br>dynamically-shaped arrays3.

But in the process of wiring JAX to feed our quantum pipeline, we realized something silly.

You don’t actually need to include any quantum instructions in your Catalyst @qjit workflows. You<br>can feed it pure, standard JAX NumPy code alongside native Python control flow. And when you do,<br>Catalyst completely bypasses the XLA compiler backend, lowers the JAX representations straight<br>through to standard MLIR, and compiles it down to machine code via LLVM (with backprop support).

… that is, we built an MLIR compilation pipeline for JAX by accident.

import jax.numpy as jnp<br>from catalyst import qjit

@qjit(autograph=True)<br>def iterative_layer(weights, inputs, threshold):<br>x = inputs

while jnp.mean(x) > threshold:<br>x = jnp.sin(jnp.dot(weights, x))

return x

To understand how we ended up here (otherwise known as, the yak we were originally shaving), we have<br>to take a slight detour through the world of quantum compilation and quantum gradients. Our goal<br>here was three-fold:

We wanted to leverage existing, mature classical compilation tooling (LLVM) and add quantum<br>support to it, rather than independently building out our own quantum infrastructure and slowly<br>adding functionality from the classical world. We are experts in quantum computing, why should we<br>reinvent the wheel when it comes to classical software and infrastructure?4

We wanted to make sure that we could represent quantum programs with structure . That is, the<br>ability to naturally intertwine array manipulation, quantum instructions, loops, and if<br>statements. This helps to both preserve a compact representation of the program (we naturally<br>write quantum algorithms with loops and if statements, we should compile with these structures<br>preserved otherwise compilation will never scale), and support natively dynamic parts of the<br>quantum program (e.g., while loops that represent a quantum-measure-and-repeat-until-success pattern).

Finally, we wanted to keep supporting quantum autodifferentiation. We have put in a lot of effort<br>to figure out how to compute gradients of quantum programs on quantum hardware itself, and wanted<br>to make sure this kept working (alongside the necessary classical autodiff).

I could probably summarize the above three points in a simple statement:

Quantum algorithms are never just quantum5.

There is a mountain of classical processing to prepare the inputs, process the outputs, and honestly<br>to even construct the circuits in the first place. This is before we even get to quantum error<br>correction, where the classical control hardware surrounding the QPU needs to calculate parameters<br>and gates at a massive scale, feed them to a QPU, perform measurements, and react in real-time with<br>ultra-low, deterministic latency. The natural conclusion: we can’t just compile the quantum<br>instructions; we need to compile the entire hybrid execution graph into a single, standalone<br>executable that executes as close to the bare metal as possible.

We decided to build this compiler infrastructure on MLIR. JAX is a natural fit for the Python layer<br>(due to its usage for differentiable programming, and the fact that it already captures classical<br>code and lowers it to MLIR), so we extended JAX to support quantum instruction sets alongside some<br>quality of life improvements (Python control flow, dynamically-shaped arrays).

Then, to handle backprop, we wired up Enzyme to backpropagate directly<br>through the classical LLVM code, while we have custom-written MLIR passes to generate the quantum<br>gradients. Because we compile the classical JAX code straight down to standard LLVM IR, we get to<br>use standard LLVM passes like Enzyme essentially “for free” — something that would be a rather<br>painful exercise to try and wire directly into XLA’s internal pipeline.

XLA is an solid piece of engineering for optimizing linear algebra and targeting GPUs and Google’s<br>own TPUs. Under the hood, it...