Introducing Catalyst: quantum just-in-time compilation

Just-in-time compilation, allowing rapid prototyping without sacrificing performance, has taken the Python and machine learning world by storm. What if we could bring that to quantum?

Today, we're excited to announce the beta availability of our next-generation compilation framework, code-named Catalyst.

A complete rethinking of a quantum execution pipeline, Catalyst enables you to write hybrid quantum-classical programs that are automatically compiled for execution on CPUs, GPUs, and QPUs — leading to significant performance and scalability improvements, and the ability to unlock new ways of programming quantum computers.

Best of all, with just-in-time (JIT) compilation and PennyLane compatibility, you can keep using the same PennyLane user interface you know and love (with a few additional goodies, such as advanced quantum control), and a significant speed boost. And of course, it is all built for differentiability.

Catalyst is currently open for preview — install it alongside PennyLane and JAX and give it a whirl, or join us directly on this development journey. Catalyst is under heavy development — we'd love to hear about how you're using the library, collaborate on development, or integrate additional devices and frontends.

In the meantime, read on to learn more about what Catalyst unlocks, and what we have planned.

Contents

• The past (and future) of quantum programming
• Catalyst: JIT, but quantum
• Advanced control structures
• Compile your entire workflow
• Behind the scenes
• What's next
• Try it out and get involved

The past (and future) of quantum programming

We open-sourced PennyLane almost four and half years ago, and in this time, we have had over a hundred contributors, published over ninety demonstrations, helped enable countless research papers, and seen a vibrant community of quantum programmers and explorers grow.

The quantum programming landscape has also changed dramatically. While Python and Jupyter remain tools of choice for many scientists and researchers, and we remain in the noisy intermediate-scale quantum (NISQ) era, we've seen a huge increase in the number (and size!) of available quantum hardware systems. Further, quantum hardware today is much more capable — from support for dynamic circuits, to entirely new ways of interfacing with quantum hardware (largely thanks to the many new quantum hardware modalities we find ourselves able to access). Meanwhile, frameworks such as QIR have been changing how we connect quantum hardware with existing classical tooling.

Separately, we've also seen an evolution in machine learning frameworks. There has been a radical transformation from the days of TensorFlow 1.0 (where you would painstakingly build static computational graphs for a performance boost later during execution), to the define-by-run approach of libraries like JAX, which seamlessly meld dynamic Python with under-the-hood just-in-time compilation. This approach — where you can continue to prototype rapidly in near-standard Python, yet have access to high-performance optimization and GPUs — has rapidly caught on, both in PyTorch (via Dynamo) and TensorFlow (via Autograph and XLA).

This evolution has also led to a massive shift in our relationship with Python, as researchers. No longer do we need to reluctantly dust off our C++ or Fortran memories for that much-needed speed boost; we can simply decorate our existing Python code with @jit, and have access to performance, scalability, and even GPUs.

However, while quantum programming has evolved over the last couple of years (alongside more powerful hardware), we are still accessing quantum hardware the same old, slow way — by sending over static circuit descriptions over the internet and waiting in queues.

What if we could bring the recent JIT compilation paradigm shift from machine learning to quantum?

Catalyst: JIT, but quantum

With Catalyst, it is as simple as decorating your existing PennyLane hybrid workflow with @qjit:

import pennylane as qml
from catalyst import qjit
from jax import numpy as jnp

dev = qml.device("lightning.qubit", wires=2, shots=1000)

@qjit
@qml.qnode(dev)
def circuit(x, y, z):
    qml.RX(jnp.sin(x), wires=[y + 1])
    qml.RY(x ** 2, wires=[z])
    qml.CNOT(wires=[y, z])
    return qml.probs(wires=[y + 1])

Under the hood, when circuit is first called, Catalyst compiles the entire quantum function — including both quantum and classical processing — down to a machine binary, by utilizing MLIR and LLVM (via QIR).

On subsequent executions, this pre-compiled binary is called with the new parameter values, avoiding slow Python processing and quantum circuit recompilation.

>>> circuit(1.072, 1, 2)  # circuit is JIT compiled
array([0.64790739, 0.35209261])
>>> circuit(2.65, 0, 2)  # the precompiled binary is executed with new parameters
array([0.94532342, 0.05467658])

All quantum algorithms are 'hybrid' to a degree, and involve some level of classical processing. As algorithms and hardware scales, classical pre- and post-processing in Python will become a significant bottleneck. Catalyst allows us to 'escape' from Python; giving us the performance of C++ from Python.

Keep using the tools you are familiar with for rapid prototyping — Python, Jupyter, and PennyLane — but benefit from increased performance.

Advanced control structures

A significant benefit of Catalyst is the ability to incorporate complex control flow around quantum operations — such as if statements and for loops, and including measurement feedforward — which is preserved during compilation.

This allows for significantly more performant compilation, while also enabling hardware to directly interpret and apply control flow around quantum operations.

from catalyst import for_loop, cond

@qjit
@qml.qnode(dev)
def circuit(n: int, p: float):

    # define a loop function
    def loop_func(i: int, x: float):

        # define a conditional ansatz
        @cond(x > 1.4)
        def ansatz():
            qml.RX(x, wires=0)
            qml.Hadamard(wires=0)

        @ansatz.otherwise
        def ansatz():
            qml.RY(x, wires=0)

        # apply the conditional quantum function
        ansatz()

        # update the value of x for the next iteration
        return x + jnp.pi / 4

    # apply for loop n times
    for_loop(0, n, 1)(loop_func)(p)
    return qml.expval(qml.PauliZ(0))

Compile your entire workflow

Catalyst truly is a hybrid quantum-classical JIT, and really shines when used to compile an entire workflow. For example, rather than just compiling your quantum function within a variational algorithm, you can compile the entire loop:

from catalyst import grad
import jax
import jaxopt

dev = qml.device("lightning.qubit", wires=2)

@qml.qnode(dev)
def circuit(params):
    qml.Hadamard(0)
    qml.RX(jnp.sin(params[0]) ** 2, wires=1)
    qml.CRY(params[0], wires=[0, 1])
    qml.RX(jnp.sqrt(params[1]), wires=1)
    return qml.expval(qml.PauliZ(1))

@qjit
def cost(param):
    diff = grad(circuit, argnum=0)
    return circuit(param), diff(param)[0]

@qjit
def optimization():
    # initial parameter
    params = jnp.array([0.54, 0.3154])

    # define the optimizer
    opt = jaxopt.GradientDescent(cost, stepsize=0.4, value_and_grad=True)
    update = lambda i, args: tuple(opt.update(*args))

    # perform optimization loop
    state = opt.init_state(params)
    (params, _) = jax.lax.fori_loop(0, 100, update, (params, state))

    return params

Note that we have included arbitrary JAX code within our qjit; all JAX code is natively compatible with Catalyst. This includes JAXopt, for JIT-compatible optimization, which works seamlessly with Catalyst's autodifferentiation.

When called, the optimization loop is compiled down to a machine binary, and when executed, it will do so with lightning speed.

>>> min_param = optimization()
>>> cost(min_param)
[array(-0.88926131), array([ 9.93427562e-06, -3.26427774e-05])]

At the moment, Catalyst supports CPUs and the Lightning high-performance simulator, but as we work to add support for more devices — including GPUs and additional quantum hardware — you will be able to deploy your compiled workflow to be executed right beside the hardware.

No more waiting for your quantum jobs to get through a queue, return to you for post-processing, only to generate another job for the queue.

Behind the scenes

Under the hood, Catalyst currently consists of the following components:

• Catalyst Compiler. The core Catalyst compiler is built using MLIR, with the addition of a quantum dialect used to represent quantum instructions. This allows for a high-level intermediate representation of the classical and quantum components of the program, resulting in advantages during optimization. Once optimized, the compiler lowers the representation down to LLVM and QIR, and a machine binary is produced.

• Catalyst Runtime. The runtime is a C++ QIR runtime that enables the execution of Catalyst-compiled quantum programs. Currently, a runtime implementation is available for the state-vector simulator lightning.qubit. A complete list of the quantum instruction set supported by this runtime implementation can be found by visiting the Catalyst documentation.

In addition, we also provide a Python frontend for PennyLane and JAX:

• PennyLane JAX frontend. A Python library that provides a @qjit decorator to just-in-time compile PennyLane hybrid quantum-classical programs. In addition, the frontend package provides Python functions for defining Catalyst-compatible control flow structures.

What's next

Currently, Catalyst is in preview. As we continue to build out Catalyst, the PennyLane frontend will likely be upstreamed into PennyLane proper, providing native JIT functionality built into PennyLane. We'll also be exploring better ways to integrate control flow with existing code, including experimental support for native Python control syntax.

We will also be working hard to build out the compiler stack, and add quantum compilation routines. This includes an API for providing or writing Catalyst-compatible compilation routines. In addition, we will be improving the autodifferentiation support, and adding support for classical autodiff, additional quantum gradients, and quantum-aware optimization methods.

Finally, we will be adding support for more devices, including quantum hardware devices. Our ambitions go even further, however. With Catalyst, we want to build the quantum programming framework of tomorrow, looking towards a future that is heterogeneous — where you can run arbitrary classical and quantum code, right beside CPUs, GPUs, and QPUs working in tandem.

To figure out the details, we need your help — please send us your use cases by starting a conversation, or trying out Catalyst for yourself.

If you are interested in working on connecting a quantum device with Catalyst, contributing quantum compilation routines, or even adding new frontends, please get in touch.

Try it out and get involved

To get started using Catalyst JIT compiler with PennyLane, check out our quick-start guide, as well as the various examples and tutorials in our documentation. For an introduction to quantum computing and quantum machine learning, you can also visit the PennyLane website for tutorials, videos, and demonstrations.

Interested in getting in touch and collaborating with us? Simply head over to our GitHub repository, check out the ongoing work, and join the discussion and development.

And if you are as excited as we are, make sure to keep an eye on the PennyLane Blog and follow us on Twitter and LinkedIn for the latest Catalyst updates throughout 2023.