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 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 (due to its usage for differentiable programming, and the fact that it already captures classical code and lowers it to MLIR), so we extended JAX to support quantum instruction sets alongside some quality of life improvements (Python control flow, dynamically-shaped arrays).

Then, to handle backprop, we wired up Enzyme to backpropagate directly through the classical LLVM code, while we have custom-written MLIR passes to generate the quantum gradients. Because we compile the classical JAX code straight down to standard LLVM IR, we get to use standard LLVM passes like Enzyme essentially “for free” — something that would be a rather 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 own TPUs. Under the hood, it takes JAXpr (a representation of the Python program), converts it to the MLIR StableHLO dialect, runs its own heavy optimizations, and eventually uses LLVM to compile down to machine code for CPUs and GPUs (or targets TPUs directly).

So if XLA already uses LLVM, why is our approach different?

We still lean on JAX to do the heavy lifting of lowering linear algebra down to StableHLO. But instead of sending that representation down the traditional XLA path, Catalyst intercepts it and lowers it directly into standard, general-purpose MLIR dialects (like linalg, arith, and scf). From there, we feed it straight to LLVM and Enzyme, completely dropping the XLA runtime.

(Side note: Someday, we’d love to bypass HLO entirely and capture NumPy semantics directly. But for now, letting JAX handle the linear algebra lowering saves us from a lot of unnecessary suffering!)

Because the frontend tracing is completely decoupled from the backend execution, the quantum bits are optional. If you hand our @qjit (Quantum Just-In-Time) decorator6 a pure classical function with zero quantum instructions, it just works:

import jax.numpy as jnp
from catalyst import qjit
import catalyst

# No quantum code here, just pure JAX + native Python control flow!
@qjit(autograph=True)
def iterative_layer(weights, inputs, threshold):
    x = inputs

    # Standard Python loops and conditionals work seamlessly
    # No need for jax.lax.while_loop or jax.lax.cond!
    while jnp.mean(x) > threshold:
        x = jnp.sin(jnp.dot(weights, x))
        catalyst.debug.print('x={x}', x=x) # runtime printing

    return x

So… what is the point?

Honestly? We aren’t entirely sure yet.

Let me be perfectly clear: this is not going to beat XLA for standard deep learning workloads. XLA has years of hyper-specific optimizations for linear algebra on GPUs and TPUs. If you are training a massive transformer, stick to standard JAX.

But what we do think is cool is what happens when you connect JAX directly to the broader LLVM ecosystem and drop the heavy XLA runtime. (Plus, no need to build XLA using Bazel either! You’re welcome.)

This bridge opens up some weird and wonderful hacking opportunities and quality of life improvements:

  • Standalone AOT binaries: XLA expects a Just-In-Time (JIT) execution environment managed by a heavy Python/C++ runtime like PJRT. By bypassing it, you can generate a pure, standalone Ahead-of-Time binary that needs zero external ML dependencies to run.

  • Dynamically-shaped arrays: Anyone who uses JAX knows that XLA is notoriously rigid when it comes to tensor shapes — if your input dimensions change, you usually trigger a frustrating and slow recompilation. Because we map directly to standard MLIR and LLVM, we can support dynamic shapes in our compiled binaries out of the box. This is a massive QoL win if you are handling variable-length sequences.

  • Native Python control flow and NumPy indexing: While, if, for, conditionals, boolean operations, even NumPy array assignments all work out of the box via Autograph/Diastatic Malt (and are differentiable).

  • Bare-Metal and edge deployment: Want to run a JAX-trained reinforcement learning policy on a custom edge-device, an FPGA, or a weird embedded chip? You can’t bring a massive ML runtime with you. Stripping it down to pure LLVM IR makes deploying to the edge (or a toaster???) actually feasible.

  • Custom MLIR dialects and passes: Because the JAX code hits standard MLIR first, you can easily inject your own custom dialects and compiler passes before it ever hits the hardware.

Is any of this useful outside of quantum?

Because we are a quantum computing team, we are hyper-focused on our specific use case (supporting large-scale quantum algorithms). The fact that we built a standalone classical JAX-to-LLVM pipeline is a fun side-effect for us, but it might actually be a powerful tool for someone else.

If you are someone who works in custom hardware, edge ML, or compiler infrastructure: does this unlock anything interesting for you? Is the ability to drop the XLA runtime, write native control flow, and target MLIR with JAX useful or interesting?

  1. We thought long and hard over whether we wanted to trace the computation on input dynamic variables, or parse the Python abstract syntax tree. In the end we went for both (best of both worlds!) but primarily rely on tracing, as it allows for arbitrary Python and 3rd party library code to be used as long as it does not act on dynamic variables. ↩︎
  2. We ended up forking Autograph from TensorFlow into a new package, Diastatic Malt ↩︎
  3. Dynamically-shaped arrays are a huge quality of life improvement, especially in quantum algorithm development where the shape of your arrays might be determined by the output of quantum measurements. ↩︎
  4. Something to note is that quantum algorithms are never purely quantum. There is a tonne of classical processing needed to get inputs ready, post-process outputs, and even construct and represent the quantum circuits themselves. And that’s all ignoring classical processing needed for quantum error correction. ↩︎
  5. Perhaps a misconception exists that you need to know quantum to work in quantum computing. That is incorrect, there is a *tonne* of classical software infrastructure and development needed to support quantum computing. Check out some of [Xanadu’s open roles](https://xanadu.ai/careers) if you are interested. ↩︎
  6. Simply do pip install pennylane-catalyst to run the following example. ↩︎


Tags: quantum, gradients

Related posts