GoKawiilSometimes 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 graph, 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 flow, and support for dynamically-shaped arrays.
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?
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).
... continue reading