Compiler concepts for physicists: introduction to MLIR dialects with xDSL and PennyLane

This is the first entry of a blog post series introducing compiler concepts for quantum physicists, based on a series of tutorials by Erick Ochoa Lopez. In the first entry of this series, we will explore the concept of MLIR dialects, using PennyLane and xDSL.

"MLIR? LLVM? SSA? What is that and what has it got to do with quantum?"

If you're a frequent user of quantum software, dabble in building your own quantum software, or are a quantum algorithms specialist, then you aware of the importance of quantum compilation. Quantum compilation is critical in order to transform our quantum programs to efficiently execute on simulators or hardware.

Over the last couple of years, this importance has been reflected in quantum software, and there are multiple frameworks and tools for writing quantum compilation passes — from PennyLane's transforms, to TKET and the Qiskit transpiler.

However, in parallel — and over the course of decades! — the classical world has been building out classical compilation tooling, theory, and frameworks. These tools — established toolchains such as LLVM, and new-kids-on-the-block such as MLIR — are incredibly powerful, and quantum software packages have started to realize that incorporating these toolchains and ideas could transform (pardon the pun) how we develop quantum compilation tooling.

In this blog post series, we'll explore classical compilation concepts (such as dialects, Single Static Assignment (SSA) and more) for physicists, using PennyLane and xDSL to illustrate these ideas.

Today, we'll cover the concept of dialects, what they are, and how we define and use them.

The role of intermediate representations (IRs)

Compilers rarely translate code in one go. Instead of directly translating a high-level circuit into low-level pulses, they use stepping stones called Intermediate Representations (IRs).

The first IR might closely resemble your quantum operator concepts. Over the course of multiple passes, the compiler translates ("lowers") your logic from operator concepts to pulse instructions by transforming one IR to a more targeted, more optimized, IR.

What exactly is an IR?

An IR is fundamentally a data structure used to represent a program internally within the compiler. It must be organized in a way that allows for efficient access and manipulation. IRs are used to:

• Transform and Optimize: Change the program structure (e.g., merging gates, reordering operations).

• Prove properties: Ensure correctness (e.g., type-check the program).

Different stages of a compilation pipeline use different kinds of IRs. Some are simple lists (like a PennyLane quantum tape), others represent programs as trees (like Abstract Syntax Trees (ASTs)), and still others use graphs (like Control Flow Graphs). These IRs can also be nested. For instance, a high-level Call Graph might only show function relationships, while a lower-level Control Flow Graph is needed to detail the flow within a single function.

Let's enter xDSL (eXtendable Structured Denotational Language), a powerful Python-native compiler framework. xDSL is built to help us define these advanced IRs, giving us the tools to transform, analyze, and ultimately supercharge quantum programs. And if you're MLIR curious, xDSL is a great place to start — it is a re-implementation of MLIR in Python!

Dialect Definition

Now that we have defined what an IR is and we have acknowledged that PennyLane uses tapes as an IR, let's talk a little bit more about xDSL and Dialects.

xDSL (and MLIR) define the concept of Dialect as a subset of an IR. By allowing you to define your own Dialect, xDSL allows you to create IRs that are composed of a set of Dialects. This allows division of labour and work reuse.

In xDSL and MLIR, a Dialect is composed of Attributes and Operations:

• An Attribute is something that is known at compile time.
• An Operation is something that defines zero or more values.

Values are not part of the dialect, they are part of xDSL itself. However, Types (which are an attribute) may be defined by xDSL's users and values may be typed with custom types. You may already see the benefit of this if one wanted to define a qubit type for example.

Let's start by installing xDSL and PennyLane, and creating a dialect by defining quantum operations that are in the set of {H, S, T, CNOT}.

Installing PennyLane and xDSL

To get started, we need to install PennyLane and xDSL — we will then use xDSL to build a simple dialect representing PennyLane's quantum instruction set.

Both PennyLane and xDSL can be installed via pip:

pip install pennylane xdsl

Defining an operation.

The xDSL way to define a Type is through the use of the @irdl.irdl_op_definition decorator and the irdl.IRDLOperation class. All classes extending irdl.IRDLOperation and decorated with @irdl.irdl_op_definition must have a field named name. The name should start with the name of the dialect (let's make that qml continue with a dot and finish with the name of the operation). This field will be used when representing the IR as a string.

from xdsl.dialects import builtin
from xdsl import irdl

@irdl.irdl_op_definition
class HadamardOp(irdl.IRDLOperation):
    name = "qml.Hadamard"

That's it! With this we are able to represent a Hadamard operation. Let's inspect the textual representation and see what it can tell us about the limitations of our first operation.

>>> print(HadamardOp())
"qml.Hadamard"() : () -> ()

xDSL provides a generic string representation for IRs that do not define their own custom string representation. It starts with the operation's dialect (in this case qml) and continues with a dot . and the dot is followed by the operation's name. Next, a set of parentheses follows, these parentheses denote the operation's operands. In this case, we can see that there are no operands. It is then followed by a colon :. The colon denotes the end of the operation's values and the start of the type of the values. The first parentheses denotes the types of the operands. And the second set of parentheses denotes the types of the outputs.

So now, we can analyze and we know that this operation we have defined is incomplete. There is no way for us to specify which wire the Hadamard operation is applied to.

Properties

Let's try to get it close to PennyLane's Hadamard operation. Even though PennyLane's operations are able to take a wide variety of types as operands (e.g., JAX tracers, class Wire, any hashable type), let's focus now only on int, and make all wires be known statically. This means that they are going to be an attribute/property and not an operand.

To create a property, we use the irdl.prop_def function and optionally pass in a constraint on the attribute. The constraint in this case is that the wire property must be an integer.

@irdl.irdl_op_definition
class HadamardOp(irdl.IRDLOperation):
    name = "qml.Hadamard"

    # Wire is a property of this operation.
    # Properties mean that the value of the wire is known.
    # Properties are typed, but since one can create their own types in xDSL
    # this is very flexibly.
    wire = irdl.prop_def(builtin.i64)

    def __init__(self, wires: list[int]):
        # Some sanitization
        assert type(wires) == list
        assert len(wires) == 1
        wire = wires[0]
        assert type(wire) == int

        # The important part
        wire = builtin.IntegerAttr(wire, builtin.i64)
        properties = {"wire" : wire}
        kwargs = {"properties" : properties}
        super().__init__(**kwargs)

>>> print(HadamardOp(wires=[0]))
"qml.Hadamard"() <{wire = 0 : i64}> : () -> ()

To instantiate an object of the HadamardOp class, one needs to call the constructor of the base class. The constructor of the base class takes named arguments, they are properties : dict [str, ir.Attribute], operands: list[ir.SSAValue], results: list[ir.TypeAttribute].

We now see that the textual representation contains an attribute dictionary. <{wire = 0 : i64}> before the colon. It denotes that this operation contains an attribute (named wire), with the value 0 and a type i64.

Pretty Printing

xDSL allows one to specify a pretty printing format. Notice that we can access the value of the wire attribute by accessing the .value.data field access path.

from xdsl.printer import Printer

@irdl.irdl_op_definition
class HadamardOp(irdl.IRDLOperation):
    name = "qml.Hadamard"

    wire = irdl.prop_def(builtin.i64)

    def __init__(self, wires: list[int]):
        assert type(wires) == list
        assert len(wires) == 1
        wire = wires[0]
        assert type(wire) == int

        wire = builtin.IntegerAttr(wire, builtin.i64)
        properties = {"wire" : wire}
        kwargs = {"properties" : properties}
        super().__init__(**kwargs)

    def print(self, printer: Printer):
        printer.print_string("(")
        printer.print_string("wires")
        printer.print_string("=")
        printer.print_string("[")
        printer.print_string(str(self.wire.value.data))
        printer.print_string("]")
        printer.print_string(")")

>>> print(HadamardOp(wires=[0]))
qml.Hadamard(wires=[0])

Perfect! now we have a pretty printer that looks similar enough to PennyLane that we can build upon.

Next, let's add some logic to define and validate operator wires:

import builtins

def eager_map(func, seq):
    return list(map(func, seq))

def sanitize_wire(wire: int | builtin.IntegerAttr[builtin.I64]):
    match type(wire):
        case builtins.int: return builtin.IntegerAttr(wire, builtin.i64)
        case builtin.IntegerAttr: return wire
    msg = f"Invalid wire type {type(wire)}"
    raise TypeError(msg)

def sanitize_wires(wires: list[int | builtin.IntegerAttr[builtin.I64]]):
    return eager_map(sanitize_wire, wires)

We now have all the pieces to define a base class for PennyLane unary operations (that is, operations that take no parameters and act on only a single wire):

class UnaryOperation(irdl.IRDLOperation):
    wire = irdl.prop_def(builtin.i64)

    def __init__(self, wires: list[int | builtin.IntegerAttr[builtin.I64]], **kwargs):
        assert len(wires) == 1
        wires = sanitize_wires(wires)
        wire = wires[0]
        properties = {"wire": wire}
        kwargs["properties"] = properties
        super().__init__(**kwargs)

    def print(self, printer: Printer):
        printer.print_string("(")
        printer.print_string("wires")
        printer.print_string("=")
        printer.print_string("[")
        printer.print_string(str(self.wire.value.data))
        printer.print_string("]")
        printer.print_string(")")

With our base class defined, this makes it easy for us to create our Hadamard, S, and T operations:

@irdl.irdl_op_definition
class HadamardOp(UnaryOperation):
    name = "qml.Hadamard"

@irdl.irdl_op_definition
class SOp(UnaryOperation):
    name = "qml.S"

@irdl.irdl_op_definition
class TOp(UnaryOperation):
    name = "qml.T"

We can also now define our CNOT operation. Note that here we cannot use the UnaryOperation base class, but instead need to define a new operation as CNOT acts on two wires.

@irdl.irdl_op_definition
class CNOTOp(irdl.IRDLOperation):
    name = "qml.CNOT"
    wire = irdl.prop_def(builtin.i64)
    target = irdl.prop_def(builtin.i64)

    def __init__(self, wires: list[int | builtin.IntegerAttr[builtin.I64]]):
        assert len(wires) == 2
        wires = sanitize_wires(wires)
        wire = wires[0]
        target = wires[1]
        properties = {"wire": wire, "target":target}
        kwargs = {"properties": properties}
        super().__init__(**kwargs)

    def print(self, printer: Printer):
        printer.print_string("(")
        printer.print_string("wires")
        printer.print_string("=")
        printer.print_string("[")
        wires = [self.wire.value.data, self.target.value.data]
        wires = [str(wire) for wire in wires]
        printer.print_list(wires, printer.print_string)
        printer.print_string("]")
        printer.print_string(")")

>>> print(HadamardOp(wires=[0]))
qml.Hadamard(wires=[0])
>>> print(SOp(wires=[1]))
qml.S(wires=[1])
>>> print(TOp(wires=[2]))
qml.T(wires=[2])
>>> print(CNOTOp(wires=[3, 4]))
qml.CNOT(wires=[3, 4])

We are still missing the concept of measurement. Let's define a measurement operation as well.

One thing that is different about measurement operations from all the previous operations is that a measurement will produce a value. Similarly to how a property is defined, a result is defined through the irdl.result_def function and we can also use a constraint to tell it that we want the output to be a builtin.i64.

@irdl.irdl_op_definition
class measureOp(irdl.IRDLOperation):
    name = "qml.measure"
    wire = irdl.prop_def(builtin.i64)
    output = irdl.result_def(builtin.i64)

    def __init__(self, wires: list[int | builtin.IntegerAttr[builtin.I64]]):
        assert len(wires) == 1
        wires = sanitize_wires(wires)
        wire = wires[0]

        # we will be passing two kwargs: properties and result_types
        properties = {"wire": wire}
        result_types = [builtin.i64]
        kwargs = {"properties": properties, "result_types": result_types}
        super().__init__(**kwargs)

    def print(self, printer: Printer):
        printer.print_string("(")
        printer.print_string("wires")
        printer.print_string("=")
        printer.print_string("[")
        printer.print_string(str(self.wire.value.data))
        printer.print_string("]")
        printer.print_string(")")

>>> print(measureOp(wires=[5]))
%0 = qml.measure(wires=[5])

And let's also create another operation that corresponds to a PennyLane quantum device, whether a simulator or a hardware device:

@irdl.irdl_op_definition
class deviceOp(irdl.IRDLOperation):
    name = "qml.device"
    device_name = irdl.prop_def(builtin.StringAttr)
    wire_count = irdl.prop_def(builtin.i64)

    def __init__(self, device_name, wires):
        if type(device_name) == str:
            device_name = builtin.StringAttr(device_name)
        if type(wires) == int:
            wires = builtin.IntegerAttr(wires, builtin.i64)
        properties = {"device_name" : device_name, "wire_count" : wires}
        kwargs = {"properties": properties}
        super().__init__(**kwargs)

    def print(self, printer: Printer):
        printer.print_string("(")
        printer.print_string_literal(self.device_name.data)
        printer.print_string(",")
        printer.print_string("wires")
        printer.print_string("=")
        wires = [self.wire_count.value.data]
        wires = [str(wire) for wire in wires]
        printer.print_list(wires, printer.print_string)
        printer.print_string(")")

>>> print(deviceOp("default.qubit", wires=1))
qml.device("default.qubit",wires=1)

Finally, a Dialect class can be created with a set of operations and attributes.

from xdsl import ir

operations = [HadamardOp, SOp, CNOTOp, measureOp, deviceOp]
attributes = []
QuantumDialect = ir.Dialect("qml", operations, attributes)

And with that, we have defined our first simple quantum dialect using xDSL!

Stay tuned for the next blog post in this series, where you will learn how to 'capture' a structured quantum program (a program with both quantum instructions and classical control flow) from Python and represent it in our xDSL dialect. From there, we'll dive into writing our own compilation passes, understanding Single Static Assignment (SSA), and more!

In the meantime, check out our Catalyst compiler which utilize MLIR and LLVM for quantum compilation, and also explore some of our compilation demos, such as Quantum just-in-time compiling Shor's algorithm with Catalyst.