We accidentally built an LLVM compiler for Jax

Sometimes when you build software, you shave a yak so deeply that you end up with a surprisingly 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 Catalyst for the quantum software library PennyLane. Our goal was relatively straightforward: optimize large, hybrid quantum-classical workflows in a scalable manner using MLIR. To do this, we needed a fast, robust way to capture classical Python processing (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 computational graph1, its support of the NumPy and SciPy APIs with relatively good coverage, and the fact that it already lowers down to MLIR. We also aimed to address a couple of quality of life improvements, such as the ability to capture native Python control flow2, and support for 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 can feed it pure, standard JAX NumPy code alongside native Python control flow. And when you do, Catalyst completely bypasses the XLA compiler backend, lowers the JAX representations straight 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 from catalyst import qjit @qjit(autograph=True) def iterative_layer(weights, inputs, threshold): x = inputs while jnp.mean(x) > threshold: 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 to take a slight detour through the world of quantum compilation and quantum gradients. Our goal here was three-fold: We wanted to leverage existing, mature classical compilation tooling (LLVM) and add quantum support to it, rather than independently building out our own quantum infrastructure and slowly adding functionality from the classical world. We are experts in quantum computing, why should we 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 ability to naturally intertwine array manipulation, quantum instructions, loops, and if statements. This helps to both preserve a compact representation of the program (we naturally write quantum algorithms with loops and if statements, we should compile with these structures preserved otherwise compilation will never scale), and support natively dynamic parts of the 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 to figure out how to compute gradients of quantum programs on quantum hardware itself, and wanted 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 to even construct the circuits in the first place. This is before we even get to quantum error correction, where the classical control hardware surrounding the QPU needs to calculate parameters and gates at a massive scale, feed them to a QPU, perform measurements, and react in real-time with ultra-low, deterministic latency. The natural conclusion: we can’t just compile the quantum instructions; we need to compile the entire hybrid execution graph into a single, standalone executable that executes as