Catalyst

Development

Examples

API

  1. Docs
  2. Catalyst Compiler Passes

Catalyst Compiler Passes

The MLIR Transformation Framework

Catalyst program transformations are implemented using the MLIR framework. A good starting point is to understand the IR structure and how MLIR transformations are written. Refer to the resources below for further information beyond this document.

In addition, many transformations can be expressed via pattern rewriting. A pattern is a specific arrangement of a group of operations into a directed acyclic graph (DAG). This allows the framework developer to provide a generic transformation mechanism, for which the compiler developer only needs to provide source and replacement patterns in order to implement a transformation. See the following resource for information on pattern rewrites:

What does Catalyst’s IR look like?

Let’s look at a simple example starting in Python:

def circuit(x: complex):
    qml.Hadamard(wires=0)

    A = np.array([1, 0], [0, np.exp(x)])
    qml.QubitUnitary(A, wires=0)

    B = np.array([1, 0], [0, np.exp(2*x)])
    qml.QubitUnitary(B, wires=0)

    return measure(0)

The corresponding IR might look something like this (simplified):

func.func @circuit(%arg0: complex<f64>) -> i1 {
    %c00 = complex.constant [0.0, 0.0] : complex<f64>
    %c10 = complex.constant [1.0, 0.0] : complex<f64>
    %c20 = complex.constant [2.0, 0.0] : complex<f64>

    %reg = quantum.alloc(1) : !quantum.reg
    %q0 = quantum.extract %reg[0] : !quantum.bit

    %q1 = quantum.custom "Hadamard"() %q0 : !quantum.bit

    %0 = complex.exp %arg0 : complex<f64>
    %A = tensor.from_elements %c10, %c00, %c00, %1 : tensor<2x2xcomplex<f64>>
    %q2 = quantum.unitary %A, %q1 : !quantum.bit

    %1 = complex.mul %arg0, %c20 : complex<f64>
    %2 = complex.exp %0 : complex<f64>
    %B = tensor.from_elements %c10, %c00, %c00, %2 : tensor<2x2xcomplex<f64>>
    %q3 = quantum.unitary %B, %q2 : !quantum.bit

    %m = quantum.measure %q3 : i1

    quantum.dealloc %r : !quantum.reg
    func.return %m : i1
}

A few things you might note in this representation:

  • Qubits need to be extracted from the global register in each scope before they can be used by gates.

  • Most quantum gates currently do not have an explicit operation defined for them, instead they share a common instruction quantum.custom and are identified by a string attribute. (This may change in the future.)

  • Quantum gates not only receive qubits as arguments, but produce new qubit values as results. In this sense, the !quantum.bit data type can be thought of as the pseudo-state of a qubit at a particular point in the program. This allows the connectivity between gates on each qubit (the “circuit wires”) to be explicit in the IR, via these data relations that MLIR calls def-use chains, linking the usage of a value (from an operation argument) to the definition of value (from an operation result).

  • Classical instructions and quantum instructions are freely mixed in the IR, just like in the user source code.

The representation above is very handy for traversing quantum operations, in effect providing the same capabilities of a directed acyclic graph (DAG) encoding of a quantum circuit. Meanwhile, classical instructions as well as side-effecting operations (e.g., print) can also be present in the same representation.

Writing transformations on Catalyst’s IR

We’ll start with DAG-to-DAG transformations, which typically match small pieces of code at a time. In our example above, we might to consider merging the two quantum.unitary applications because they act on the same qubit in immediate succession:

%0 = complex.exp %arg0 : complex<f64>
%A = tensor.from_elements %c10, %c00, %c00, %1 : tensor<2x2xcomplex<f64>>
%q2 = quantum.unitary %A, %q1 : !quantum.bit                                (A)

%1 = complex.mul %arg0, %c20 : complex<f64>
%2 = complex.exp %0 : complex<f64>
%B = tensor.from_elements %c10, %c00, %c00, %2 : tensor<2x2xcomplex<f64>>
%q3 = quantum.unitary %B, %q2 : !quantum.bit                                (B)

Note how the value %q2 links the two operations together from definition (A) to use (B) across several other instructions.

As seen in the pattern rewriter documentation, a new rewrite pattern can be defined as a C++ class as follows, where we will focus on the match and rewrite methods (refer to the link for the full class and up to date information):

struct QubitUnitaryFusion : public OpRewritePattern<QubitUnitaryOp>
{
    ...

    LogicalResult match(QubitUnitaryOp op) const override {
        // The ``match`` method returns ``success()`` if the pattern is a match, failure
        // otherwise.
    }

    void rewrite(QubitUnitaryOp op, PatternRewriter &rewriter) {
        // The ``rewrite`` method performs mutations on the IR rooted at ``op`` using
        // the provided rewriter. All mutations must go through the provided rewriter.
    }

    ...
};

Note that by inheriting from OpRewritePattern instead of the generic RewritePattern, operations will automatically be filtered and our pattern will only be invoked on QubitUnitaryOp objects.

The first step in pattern rewriting is the matching phase. We want to match the following pattern of QubitUnitary operations (represented in graph form, where the first argument is the matrix, and the second is a list of qubits):

QubitUnitary(*, QubitUnitary(*, *))

Let’s implement it in C++:

LogicalResult QubitUnitaryFusion::match(QubitUnitaryOp op)
{
    ValueRange qbs = op.getInQubits();
    Operation *parent = qbs[0].getDefiningOp();

    if (!isa<QubitUnitaryOp>(parent))
        return failure();

    QubitUnitaryOp parentOp = cast<QubitUnitaryOp>(parent);
    ValueRange parentQbs = parentOp.getOutQubits();

    if (qbs.size() != parentQbs.size())
        return failure();

    for (auto [qb1, qb2] : llvm::zip(qbs, parentQbs))
        if (qb1 != qb2)
            return failure();

    return success();
}

Note that we have used a couple of functions (like getInQubits and getOutQubits) from the definition of the QubitUnitaryOp class. Since we define our operations in the declarative ODS (tablegen) format, the corresponding C++ classes are automatically generated. This is the definition for the QubitUnitaryOp from the QuantumOps.td file:

def QubitUnitaryOp : Gate_Op<"unitary"> {
    let summary = "Apply an arbitrary fixed unitary matrix";
    let description = [{
        The `quantum.unitary` operation applies an arbitrary fixed unitary matrix to the
        state-vector. The arguments are a set of qubits and a 2-dim matrix of complex numbers
        that represents a Unitary matrix of size 2^(number of qubits) * 2^(number of qubits).
    }];

    let arguments = (ins
        2DTensorOf<[Complex<F64>]>:$matrix,
        Variadic<QubitType>:$in_qubits
    );

    let results = (outs
        Variadic<QubitType>:$out_qubits
    );

    let assemblyFormat = [{
        `(` $matrix `:` type($matrix) `)` $in_qubits attr-dict `:` type($out_qubits)
    }];
}

MLIR will automatically generate canonical get* methods for attributes like in_qubits, out_qubits, and matrix. When in doubt it’s best to have a look at the generated C++ files in the build folder, named QuantumOps.h.inc and QuantumOps.cpp.inc in this instance.

Alright, now that we have the matching part, let’s implement the actual transformation via the rewrite method. All we need to do is replace the original pattern with the following:

QubitUnitary(A, QubitUnitary(B, Q))  -->  QubitUnitary(AxB, Q)

In C++ it will look as follows:

void QubitUnitaryFusion::rewrite(QubitUnitaryOp op, PatternRewriter &rewriter)
{
    ValueRange qbs = op.getInQubits();
    QubitUnitaryOp parentOp = cast<QubitUnitaryOp>(qbs[0].getDefiningOp());

    Value m1 = op.getMatrix();
    Value m2 = parentOp.getMatrix();

    linalg::MatmulOp matmul = rewriter.create<linalg::MatmulOp>(op.getLoc(), {m1, m2}, {});
    Value res = matmul.getResult(0);

    rewriter.updateRootInPlace(op, [&] {
        op->setOperand(0, res);
    });
    rewriter.replaceOp(parentOp, parentOp.getResults());
}

When writing transformations, the rewriter is the most important tool we have. It can create new operations for us, delete others, or change the place in the IR where we are choosing to make changes (also called the insertion point). Let’s have look at some of these elements:

  • Constructing new operations:

    New operations are created via the rewriter.create method. Here we want to generate a matrix multiplication instruction from the linalg dialect. C++ namespaces usually correspond to the dialect name. The first thing the rewriter needs is always a location object, which is used in debugging to refer back to the original source code line, for example. Following this, we need to provide the right arguments to instantiate the operation. So-called operation builders are automatically defined for this purpose, whose source can be referenced to consult which arguments are required. Looking into LinalgStructuredOps.h.inc for example reveals the following options:

    static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ValueRange inputs, ValueRange outputs, ArrayRef<NamedAttribute> attributes = {});
static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, TypeRange resultTensorTypes, ValueRange inputs, ValueRange outputs, ArrayRef<NamedAttribute> attributes = {});
static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, TypeRange resultTensorTypes, ValueRange operands, ArrayRef<NamedAttribute> attributes = {});
static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, TypeRange resultTensorTypes, ValueRange inputs, ValueRange outputs, Attribute cast, ArrayRef<NamedAttribute> attributes = {});

    We can always ignore the first two arguments, odsBuilder and odsState, but the remaining ones are the arguments we’ll need to provide to the rewriter. We chose the simplest one which only requires specifying a range of values for the operation inputs (two to be precise). We can ignore outputs argument for now as it is a peculiarity of the linalg dialect. If necessary, the result types of an operation may be specified as can be seen in the second version, but for matmul the result types can be automatically deduced.

  • Removing operations:

    We can remove operations via the rewriter.replaceOp method (among others). The reason we don’t straight up delete operations is that that would break the def-chains in the IR. Instead, we always need to provide replacement values for the results that the operation to be deleted defined. In this case, we simply replace the output qubit values with the input qubit values to maintain the correct “wire” connections. We would thus change

    %q2 = quantum.unitary %A, %q1 : !quantum.bit
%q3 = quantum.unitary %B, %q2 : !quantum.bit

    into

    %q3 = quantum.unitary %B, %q1 : !quantum.bit

    Note how the argument of the second unitary op was automatically swapped from %q2 to %q1.

  • Updating operations:

    Operation arguments and attributes can also be modified in-place (without creating a new operation). We use this to replace the matrix argument of our operation with the result of the multiplication. Since this mechanism doesn’t go through the rewriter, he have to notify it explicitly that we are making changes to an operation:

    rewriter.updateRootInPlace(op, [&] {
    op->setOperand(0, res);
});

    Note that in order to change to results on an operation you will need to create a copy of it and erase the existing operation, they cannot be modified in-place.

Invoking transformation patterns

IR changes are always effected by a transformation pass. Many compilers are structured around the notion of passes, where the program is progressively transformed and each pass is responsible for a particular sub-task.

While the transformation pattern we wrote above defines how we want to transform certain aspects of our program, it doesn’t yet specify how the patterns are applied to an input program. For this we need to write a pass.

The simplest approach might be to say we want our transformation pass to look at the entire program, and apply a set of patterns we defined like the one above. We can do so by creating an OperationPass that acts on an MLIR module (remember an MLIR module is an operation that itself contains globals and other function operations, which themselves can contain other operations, and so on):

struct QuantumOptimizationPass : public PassWrapper<QuantumOptimizationPass, OperationPass<ModuleOp>>
{
    void runOnOperation() {
        // Get the current operation being operated on.
        ModuleOp op = getOperation();
        MLIRContext *ctx = &getContext();

        // Define the set of patterns to use.
        RewritePatternSet quantumPatterns(ctx);
        quantumPatterns.add<QubitUnitaryFusion>(ctx);

        // Apply patterns in an iterative and greedy manner.
        if (failed(applyPatternsAndFoldGreedily(op, std::move(quantumPatterns)))) {
            return signalPassFailure();
        }
    }
};

To apply patterns we need a pattern applicator. There a few in MLIR but typically you can just use the greedy pattern rewrite driver (applyPatternsAndFoldGreedily), which will iterative over the IR and apply patterns until a fixed point is reached.

Writing more general transformations

The pattern-based approach to transformations is not limited to small peephole optimizations like the one above, in fact all transformation passes in Catalyst currently use either regular rewrite patterns or dialect conversion patterns. Let’s take a quick look at the finite-difference method in Catalyst for example.

The starting point for the transformation is the differentiation instruction in our gradient dialect (GradOp). It acts like a function call, but instead returns the derivative of the function for some given inputs:

func.func @my_func(f64, f64, f64) -> f64 {
    ...
}

%deriv:3 = gradient.grad "fd" @my_func(%x, %y, %z) : (f64, f64, f64) -> (f64, f64, f64)

We’ll want to replace this with code that implements the finite-difference method. The pass implementation will essentially look like the one above (say GradientPass), but with a different pattern set. This pattern would instead act on all GradOp objects in the program:

struct FiniteDiffLowering : public OpRewritePattern<GradOp>

But since the gradient could be calculated in different ways, we want to filter matches to those gradient ops that specify the finite-difference method, indicated via the "fd" attribute:

LogicalResult FiniteDiffLowering::match(GradOp op)
{
    if (op.getMethod() == "fd")
        return success();

    return failure();
}

For the rewriting part we’ll want to introduce a few new elements, such as looking up symbols (function names), creating new functions, and changing the insertion point.

void FiniteDiffLowering::rewrite(GradOp op, PatternRewriter &rewriter)
{
    // First let's find the function the grad operation is referencing.
    func::FuncOp callee =
        SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(op, op.getCalleeAttr());
    if (!callee)
        return signalPassFailure();

    // Now let's create a new function to place the differentiation code into, so it doesn't
    // pollute the current scope. We'll insert the new function after the callee.
    {
        // Insertion guards are useful to store the current IR position (insertion point) on stack,
        // returning to it upon exiting the C++ scope.
        PatternRewriter::InsertionGuard insertGuard(rewriter);
        rewriter.setInsertionPointAfter(callee);

        // Specify the properties of the new function like name, type signature, and visibility.
        std::string fnName = op.getCallee().str() + ".finitediff";
        StringAttr visibility = rewriter.getStringAttr("private");
        // The function type should be identical to the type signature of the grad operation.
        FunctionType fnType = rewriter.getFunctionType(op.getOperandTypes(), op.getResultTypes());

        gradFn = rewriter.create<func::FuncOp>(op.getLoc(), fnName, fnType, visibility, nullptr, nullptr);

        // Now we just to populate the actual body of the function. First create an empty body.
        Block *fnBody = gradFn.addEntryBlock();
        // Move the insertion point to inside the function body.
        rewriter.setInsertionPointToStart(fnBody);
        // Populate the function body.
        populateFiniteDiffMethod(rewriter, op, gradFn);
    }
}

Symbols are string references to IR objects, which rather than containing a physical reference or pointer to the actual object, only refer to it by name. In order to dereference a symbol we always have to look it up in a symbol table. This means that symbols are a bit more flexible and don’t have the same constraints as the SSA values used everywhere else in the IR.

To help visualize the process, after this step we would have gone from the IR shown above:

func.func @my_func(%x: f64, %y: f64, %z: f64) -> f64 {
    ...
}

%deriv:3 = gradient.grad "fd" @my_func(%x, %y, %z) : (f64, f64, f64) -> (f64, f64, f64)

to the following IR:

func.func @my_func(%x: f64, %y: f64, %z: f64) -> f64 {
    ...
}

func.func @my_func.finitediff(f64, f64, f64) -> (f64, f64, f64) {
    <contents of populateFiniteDiffMethod>
}

%deriv:3 = gradient.grad "fd" @my_func(%x, %y, %z) : (f64, f64, f64) -> (f64, f64, f64)

Let’s fill out the rest of the method. The finite-difference method works by invoking the callee twice with slightly different parameter values, approximating the partial derivative as follows:

\[\frac{d}{dx} f(x, y, z) \approx \frac{f(x + h, y, z) - f(x, y, z)}{h}\]

In code:

void populateFiniteDiffMethod(PatternRewriter &rewriter, GradOp op, func::FuncOp gradFn)
{
    Location loc = op.getLoc();
    ValueRange callArgs = gradFn.getArguments();

    // We can reuse the same f(x, y, z) evaluation for all partial derivatives.
    func::CallOp callOp = rewriter.create<func::CallOp>(loc, callee, callArgs);

    // Loop through x, y, z to collect the partial derivatives.
    std::vector<Value> gradient;
    for (auto [idx, arg] : llvm::enumerate(callArgs)) {

        FloatAttr hAttr = rewriter.getF64FloatAttr(0.1); // or another small fd parameter
        Value hValue = rewriter.create<arith::ConstantOp>(loc, hAttr);

        Value argPlusH = rewriter.create<arith::AddFOp>(loc, arg, hValue);

        // Make a copy of arguments to replace the argument with it's shifted value.
        std::vector<Value> callArgsForward(callArgs.begin(), callArgs.end());
        callArgsForward[idx] = argPlusH;
        func::CallOp callOpForward =
            rewriter.create<func::CallOp>(loc, callee, callArgsForward);

        // Compute the finite difference.
        Value difference = rewriter.create<arith::SubFOp>(loc, callOpForward.getResult(0), callOp.getResult(0));
        Value partialDerivative = rewriter.create<arith::DivFOp>(loc, difference, hValue);
        gradient.push_back(partialDerivative);
    }

    rewriter.create<func::ReturnOp>(loc, gradient);
}

Alright, our function should now look something like this:

func.func @my_func.finitediff(%x: f64, %y: f64, %z: f64) -> (f64, f64, f64) {
    %h = arith.constant 0.1 : f64

    %fres = func.call @my_func(%x, %y, %z) : (f64, f64, f64) -> f64

    %xph = arith.addf %x, %h : f64
    %fxph = func.call @my_func(%xph, %y, %z) : (f64, f64, f64) -> f64
    %diffx = arith.subf %fxph, %fres : f64
    %dx = arith.divf %diffx, %h

    %yph = arith.addf %y, %h : f64
    %fyph = func.call @my_func(%x, %yph, %z) : (f64, f64, f64) -> f64
    %diffy = arith.subf %fyph, %fres : f64
    %dy = arith.divf %diffy, %h

    %zph = arith.addf %z, %h : f64
    %fzph = func.call @my_func(%x, %y, %zph) : (f64, f64, f64) -> f64
    %diffz = arith.subf %fzph, %fres : f64
    %dz = arith.divf %diffz, %h

    func.return %dx, %dy, %dz : f64, f64, f64
}

Finally, we have to amend our rewrite function to invoke the new function we created and delete the GradOp from the IR:

void FiniteDiffLowering::rewrite(GradOp op, PatternRewriter &rewriter)
{
    ...
        populateFiniteDiffMethod(rewriter, op, gradFn);
    }

    rewriter.replaceOpWithNewOp<func::CallOp>(op, gradFn, op.getArgOperands());
}

Note how we can create a new operation, take its results, and use those to replace another operation in one go. This turns the previous IR:

func.func @my_func(%x: f64, %y: f64, %z: f64) -> f64 {
    ...
}

func.func @my_func.finitediff(f64, f64, f64) -> (f64, f64, f64) {
    <contents of populateFiniteDiffMethod>
}

%deriv:3 = gradient.grad "fd" @my_func(%x, %y, %z) : (f64, f64, f64) -> (f64, f64, f64)

into:

func.func @my_func(%x: f64, %y: f64, %z: f64) -> f64 {
    ...
}

func.func @my_func.finitediff(f64, f64, f64) -> (f64, f64, f64) {
    <contents of populateFiniteDiffMethod>
}

%deriv:3 = func.call @my_func.finitediff(%x, %y, %z) : (f64, f64, f64) -> (f64, f64, f64)

Catalyst’s Transformation Library

Why don’t you try writing a pass of your own? Or have a look at our existing transformations from

The pass declarations and headers for transformations are located in the include directory of each dialect: quantum, gradient, and catalyst.

Previous
Next

Downloads

dev/transforms
 
Download Python script
 
Download Notebook
 
View on GitHub