Data

Tools

Indexes

Applications

Experiments

Awesome

An open API service indexing awesome lists of open source software.

https://github.com/cqcl/simple_tket2_mbqcification

Implementation of a simple MBQCification pass in TKET2.
https://github.com/cqcl/simple_tket2_mbqcification

Last synced: about 1 year ago
JSON representation

Implementation of a simple MBQCification pass in TKET2.

Awesome Lists containing this project

README 

          
This is a toy example of the use of TKET2 to do passes with local rewrites that include classical and quantum wires.

This was developed as a hack month project, see its Confluence page [here](https://cqc.atlassian.net/wiki/spaces/HM2/pages/2852159499/HM2-51+MBQCification+pass+on+TKET2).



# Getting started

- Install TKET2 by following the steps in https://github.com/CQCL/tket2/blob/main/DEVELOPMENT.md.

- Clone this repository so that both the folder containing the `tket2` repository and the folder containing this repository live in the same directory.

- Enter the root folder of this repository and run `cargo run`.



# High-level explanation of the code



I hope that this toy project can be used as a reference implementation for people to learn about TKET2/HUGR features. 

To that end, this section provides a high-level explanation of how these features were used in the project.

You should probably read the [Simple approach](https://cqc.atlassian.net/wiki/spaces/HM2/pages/2852159499/HM2-51+MBQCification+pass+on+TKET2#Simple-approach) section from the Confluence page I wrote, which explains the algorithm and the different steps to be implemented.



## A preamble on Rust



If you are new to Rust (like me), here are some tips I gathered while doing this project:

- The [Rust book](https://doc.rust-lang.org/book/?search=object) is the go to reference. You should probably read the "Getting Started" section. These are the other sections I suggest you read:

  - Variables can't change unless you explicitly declare them as mutable. See [Variables and Mutability](https://doc.rust-lang.org/book/ch03-01-variables-and-mutability.html).

  - Having a look at the [Data Types](https://doc.rust-lang.org/book/ch03-02-data-types.html#data-types) section may be useful. In particular, note that indexing of vectors uses the syntax `my_vec[0]` but indexing of tuples uses `my_tuple.0`.

  - In Rust, you usually specify the return value of a function by not adding a `;` on the last line of your function (there's a `return` keyword you can use as well, though). See [Functions with Return Values](https://doc.rust-lang.org/book/ch03-03-how-functions-work.html?highlight=semicolon#functions-with-return-values).

  - If you get confused by things like `&`, `*` and `mut`, read the [Reference and Borrowing](https://doc.rust-lang.org/book/ch04-02-references-and-borrowing.html) section from the Rust book.

  - If you get confused by things like `.unwrap()` and `?`, read the [Recoverable Errors with Result](https://doc.rust-lang.org/book/ch09-02-recoverable-errors-with-result.html?highlight=unwrap#recoverable-errors-with-result) section.

  - If you get confused by things like `|x| x+1`, know that this is just the way Rust lets you define "anonymous functions", i.e. `f(x) = x+1` in this case. For a more in-depth explanation see [this section](https://doc.rust-lang.org/book/ch13-01-closures.html?highlight=anonymous%20functions#closures-anonymous-functions-that-capture-their-environment) (but it's not an essential read in my opinion).

- [rust-analyzer](https://rust-analyzer.github.io/) is your friend. In my code (and also TKET2/HUGR code) type annotations only appear in function signatures, but the types of all variables can be inferred automatically and displayed when using rust-analyzer.

  - Tip: if you feel like your code gets too cluttered with all of the types annotations added in, there's an option (at least in VS Code) to hide them by default and only show them when you press Ctrl+Alt.

- Using `dbg!(my_value)` will print a string representation of `my_value`. You can't apply it on everything (only on basic types or structs that implement the [Debug trait](https://doc.rust-lang.org/book/appendix-03-derivable-traits.html?highlight=Debug%20trait#debug-for-programmer-output)), but it works for most TKET2/HUGR objects.



## Code structure



- `Cargo.toml` provides metadata for the project's crate and its dependencies. This is automatically generated by calls to `cargo` (although changing it by hand is also find).

- `src/main.rs` contains a simple example circuit and the main function that calls the implementation of the steps described in [the Confluence page](https://cqc.atlassian.net/wiki/spaces/HM2/pages/2852159499/HM2-51+MBQCification+pass+on+TKET2#Simple-approach).

- `src/utils.rs` provides a function `viz_hugr` for visualisation of HUGRs, and a function `apply_rules_exhaustively` that applies all specified rewrite rules to a given HUGR until no more can be applied. The rewrite rules are specified by providing a list (vector) of pairs `(LHS, RHS)` where both elements of the tuple are HUGRs.

- `src/rewrites.rs` provides the implementation of steps 1-3 described in [the Confluence page](https://cqc.atlassian.net/wiki/spaces/HM2/pages/2852159499/HM2-51+MBQCification+pass+on+TKET2#Simple-approach). Each one is a rewrite pass that is implemented by calling `apply_rules_exhaustively` from `utils.rs`.

- `src/patterns.rs` provides functions to build each of the HUGRs acting as the LHS and RHS for the rewrite rules.

- `src/mbqc_ops.yaml` defines an MBQC extension for HUGR, including a custom `MyBool` type and custom operations such as classically controlled Paulis, destructive measurements and XOR logical gates.



## Building HUGRs



All HUGRs in this project are built using the `DFGBuilder` from the `hugr` crate. At the time of writing, there's also a `CircuitBuilder` available in `hugr` whose interface is more closely related to that of TKET1; however, in its current stage it was not flexible enough to add qubit allocation operations which were essential for this project.

The `DFG` in `DFGBuilder` stands for "Data Flow Graph", which means that we're going to be building a graph whose (directed) edges carry data from one node (operation) to another.



A good example to get started with is `circ_example` from `main.rs`. The first line in it initialises the builder and indicates the input and output types of the HUGR we're building.

```

let mut h = DFGBuilder::new(FunctionType::new(vec![QB_T; 4], vec![QB_T; 4]))?;

```

In this case, we're specifying that the HUGR will have four input qubits and four output qubits. This is described by the input to the initialiser `DFGBuilder::new` which needs to be of type `FunctionType`, i.e. a function signature: in our case, of one that goes from a vector of four qubits (input) to another vector of four qubits (output). This will create the following HUGR:

![image](https://github.com/CQCL/simple_tket2_mbqcification/assets/104848389/9c69f903-dff9-43a4-8cde-b5a1c8da509d)



**Note:** You can visualise the HUGR at any point while we build it by adding a call to `viz_hugr(h.hugr());` (this will open up an image in your browser).

As expected, our HUGR now has an `Input` and `Output` node, both with four qubits. The tree on the right is describing the hierarchy of our HUGR and we do not need to worry about it: it's just saying that node 0 (the data flow graph itself) "contains" the input and output node we just added. All of the nodes we add to this `h` are going to be at the same level of the hierarchy.



Before we start adding quantum gates to the HUGR we need to fetch the qubits from the `Input` node:

```

let mut inps = h.input_wires();

let q0 = inps.next().unwrap();

let q1 = inps.next().unwrap();

let q2 = inps.next().unwrap();

let q3 = inps.next().unwrap();

```

The expression `h.input_wires()` returns an iterator over the qubits and we can fetch each of them by calling `next()` on this iterator. The output of `next()` is an `Option`, since it may return a `Wire` or "nothing" if there are no more elements in the iterator. Since we are certain that there are exactly four qubits, we can just call `.unwrap()` to get the `Wire` out of it (see [more details](#a-preamble-on-rust)).



We can now add our first quantum gate to the HUGR:

```

let res = h.add_dataflow_op(Tk2Op::H, [q3])?;

let q3 = res.out_wire(0);

```

We are adding a Hadamard gate to the last qubit. The first argument to `h.add_dataflow_op` specifies the operation that we are adding; all quantum gates in TKET1 will eventually be available through `Tk2Op::name`. The second argument is the list of wires the operation acts on. The result `res` we get back is a handle that we can use to retrieve the output wires of the node we just added. We do so by calling `res.out_wire(i)`, which returns the wire in the `i`-th position of the list of outputs of the node. In this case, a Hadamard gate has a single output, so we retrieve our qubit wire by calling `res.out_wire(0)`.

Two-qubit gates can be added in a similar fashion:

```

let res = h.add_dataflow_op(Tk2Op::CZ, [q2, q3])?;

let q2 = res.out_wire(0);

let q3 = res.out_wire(1);

```

Currently, our HUGR looks like this:

![image](https://github.com/CQCL/simple_tket2_mbqcification/assets/104848389/d2edc356-bf18-49d9-a852-9dbdf7e6daa9)



When we are done adding operations, we need to connect the final qubit wires to the output node.

```

h.finish_hugr_with_outputs([q0, q1, q2, q3], &PRELUDE_REGISTRY)

```

The first input is the list of wires to be connected to the output node. The second is a reference to a data structure `PRELUDE_REGISTRY` that we import from the `hugr` crate without altering it. In short, the registry lets the builder know which are the HUGR extensions we have used to build our HUGR. This is necessary (among other reasons?) so that the builder knows the type signature of each operation in the HUGR that it's building, so that it can check it is valid. In the section on [using HUGR extensions](#using-hugr-extensions) we'll see how we can update this data structure to add information about our custom HUGR extension. 



The call to `h.finish_hugr_with_outputs` returns a `Result`, which matches the output of `circ_example()`. To obtain the HUGR we call

```

let mut circ = circ_example().unwrap();

```

in our `main` function. The reason why we make it mutable is that the rewrite passes we apply next will modify the HUGR in place. We can visualise the final circuit by calling `viz_hugr(&circ);`.

![image](https://github.com/CQCL/simple_tket2_mbqcification/assets/104848389/90dba7cc-744f-44ba-bf9e-781276906920)



All HUGRs in `patterns.rs` are built in the same way. However, in many cases the operations added don't come just from `Tk2Op`, but some come from the custom `ExtMBQC` extension I have defined in `mbqc_ops.yaml`. We delve into this in a following [section](#using-hugr-extensions).



### HUGR validation and debugging



When building a HUGR with `DFGBuilder`, validation is done when calling `h.finish_hugr_with_outputs`. This means that any error such as a leaving qubit wire unconnected will only be detected at Rust runtime, and the error will point to the line calling `h.finish_hugr_with_outputs`. For instance, here is an error you may encounter:

```

thread 'main' panicked at src/main.rs:37:80:

called `Result::unwrap()` on an `Err` value: InvalidHUGR(UnconnectedPort { node: Node(1), port: Port(Outgoing, 3), port_kind: Value(Type(Extension(CustomType { extension: IdentList("prelude"), id: "qubit", args: [], bound: Any }), Any)) })

```

This error is telling us that the HUGR is invalid because there is an `UnconnectedPort`; specifically: `node: Node(1), port: Port(Outgoing, 3)`. We can then call `viz_hugr(h.hugr())` just before `h.finish_hugr_with_outputs` and find the node and port it is referring to. Notice that since we are calling `viz_hugr(h.hugr())` before `h.finish_hugr_with_outputs` the final wires have not yet been connected to the output node.

![image](https://github.com/CQCL/simple_tket2_mbqcification/assets/104848389/3e9a2bb2-b9fe-462c-9afe-6c1bd730c479)



Indeed, the error was caused because I missed adding input to the `H` gate, so `q3` is not connected to anything.



## Using HUGR extensions



In this project I defined a custom HUGR extension to define operations that are not supported in `Tk2Op`. There were two motivations:

1. Extend the features of the HUGR. In particular, I added a `Copy` and `DiscardSignal` nodes for classical wires, which were not available in the `PRELUDE_REGISTRY`. At some point these nodes may not be necessary, since classical wires can be flagged as "copyable", so that a single wire can be connected to multiple (or no) inputs. However, adding these nodes explicitly was useful for the sake of rewrites.

2. Simplify the HUGR by defining primitives that would otherwise require multiple nodes to define using operations from `Tk2Op`. An example of this is `MeasureX` which represents a destructive measurement on the `X` basis. If described with `Tk2Op` nodes this would require a `Measure` after an `H` gate to change the basis, followed by a `QFree` to discard the qubit output. Creating a single node that captures this semantics not only makes the HUGR smaller, but it also simplifies rewriting. This was particularly useful in the case of classically-controlled Pauli corrections, which would have otherwise required me to do rewriting on a multi-level HUGR with `Conditional` nodes (e.g. when pushing `X` corrections through a `CZ` gate). At the end of our passes we can apply a final rewrite pass to "compile down" each of these abstract nodes down to its `Tk2Op` constituents.



The extension `ExtMBQC` is defined in the human-readable format in the YAML file `mbqc_ops.yaml`. Its syntax is self-explanatory. Note that `Q` is the current keyword that maps to the HUGR wire type `QB_T` used for wires (I expect this name inconsistency will be fixed in later versions of the YAML parser). Since at the moment of writing there was no special keyword to map to the classical wire type `BOOL_T`, I had to add my own type `MyBool`. All values of the field `name` are chosen by the user and will be used in the Rust code to refer to the extension/operations/types. The `bound` field when defining a type determines which operations are allowed on wires of the given type; in this case I indicate `MyBool` can be copied by using the special keyword `Copyable`.



To use the extension defined in `mbqc_ops.yaml` you first need to load it. This is done in the `main` function as follows:

```

let file = Path::new("./src/mbqc_ops.yaml");

let mut reg = PRELUDE_REGISTRY.clone();

load_extensions_file(file, &mut reg).unwrap();

```

where we first indicate the location of the YAML file, then creates a save a copy of the `PRELUDE_REGISTRY` in the variable `reg` to which we will add the extension using `load_extensions_file`. The resulting registry in `reg` must be passed to any function that builds or rewrites a HUGR that contains nodes from the `ExtMBQC` extension we just loaded.



The function `xcorr_h` in `patterns.rs` builds a HUGR that contains operations and types from the `ExtMBQC` extension. It uses the `ExtensionRegistry` passed to it as input (e.g. via `xcorr_h(&reg)` from `main`) to create the necessary instances of custom wires and nodes:

```

pub fn xcorr_h(registry: &ExtensionRegistry) -> Result {

    let extension = registry.get("ExtMBQC").unwrap();

    let my_bool = Type::new_extension(extension.get_type("MyBool").unwrap().instantiate([]).unwrap());

    let x_corr = extension.instantiate_extension_op("CorrectionX", [], registry).unwrap();

```

Notice that the syntax to instantiate types and operations are different. This interface may change in the future, but the key idea should remain: you can use the `name` field defined in the YAML file to create an instance that can be used by the HUGR builder.



Continuing with the example of `xcorr_h`, we can use an extension type in the signature of a HUGR:

```

let mut h = DFGBuilder::new(FunctionType::new(vec![QB_T, my_bool], vec![QB_T]))?;

```

In this case, we are representing a circuit whose input is a qubit and a classical wire and whose output is just a single qubit. We can unpack the inputs as usual:

```

let mut inps = h.input_wires();

let q = inps.next().unwrap();

let c = inps.next().unwrap();

```

and add our custom `ExtensionOp` node in the same way we'd add a `Tk2Op` node:

```

let res = h.add_dataflow_op(x_corr, [q, c])?;

let q = res.out_wire(0);

let res = h.add_dataflow_op(Tk2Op::H, [q])?;

let q = res.out_wire(0);

```

Finally, when calling `h.finish_hugr_with_outputs` we need to provide the `ExtensionRegistry` containing our `ExtMBQC` extension:

```

h.finish_hugr_with_outputs([q], registry)

```



## Matching and rewriting



The implementation of the rewrite passes in this project appears in `rewrite.rs`. All of them follow the template below.

```

pub fn push_corrections_and_s_gates(circ: &mut Hugr, reg: &ExtensionRegistry) {

    // Specify the rewrite rules

    let rules = vec![

        // Push corrections

        (xcorr_h(&reg), h_zcorr(&reg)),

        (zcorr_h(&reg), h_xcorr(&reg)),

        ...

        // Push S gates

        (s_cz_0(), cz_s_0()),

        (s_cz_1(), cz_s_1()),

    ]

    // Unwrap all of the above `Result` types into `Hugr`

    .iter()

    .map(|rule| (rule.0.clone().unwrap(), rule.1.clone().unwrap()))

    .collect();

    // Apply them exhaustively

    apply_rules_exhaustively(rules, circ);

}

```

Each rewrite rule where a subcircuit `LHS` is meant to be replace by another subcircuit `RHS` is specified as a pair `(LHS, RHS)`. These subcircuits correspond to the HUGRs built in `patterns.rs`. The lines `.iter()`, `.map(...)` and `collect()` are there just to extract the `Hugr` from each of the `Result` returned by the functions in `patterns.rs`. Alternatively, I could have written each rule as `(s_cz_0().unwrap(), cz_s_0().unwrap())`. Finally, all of the rules collected in the vector `rules` are applied exhaustively on the input circuit by calling `apply_rules_exhaustively(rules, circ);`. The latter function is defined in `utils.rs` and explained below.



### Applying all rewrite rules exhaustively



The workflow of `apply_rules_exhaustively` in `utils.rs` is as follows:

1. Find all matches in the current circuit for the LHS of each of the given rules.

2. Go over each of the matches and replace them with the RHS of the corresponding rule.

3. Since the circuit has changed, there may be new matches we can find, so we repeat the process until no more matches are found.



In practice the implementation is a bit more nuanced, so let's go step by step. The first thing we must do is convert each of the LHS of the input `rules: Vec<(Hugr, Hugr)>` into a `CircuitPattern`.

```

let mut lhs_of_rules = vec![];

for (lhs, _) in rules.iter() {

    lhs_of_rules.push(

        CircuitPattern::try_from_circuit(&lhs).unwrap()

    );

}

```

Here we are collecting the LHS of each rule in a vector `lhs_of_rules`, converting it to a `CircuitPattern` by calling `CircuitPattern::try_from_circuit(&lhs)`. 

Next, we need to create a `PatternMatcher` for this collection of patterns:

```

let matcher = PatternMatcher::from_patterns(lhs_of_rules);

```

which we can immediately use to find all instances of the subcircuits in `lhs_of_rules` in the current circuit:

```

let mut matches = matcher.find_matches(circ);

```

We only need to build the `PatternMatcher` once, and we can reuse it as often as we like on different input circuits. The resulting `matches` is a list of pattern matches `Vec`. Each `PatternMatch` indicates the location in the `circ` HUGR where the match was found (more explicitly, it lists the nodes of the matched subgraph) as well as the `PatternID` identifying the element of `lhs_of_rules` that was matched.



The next step is to convert each `PatternMatch` in `matches` to a `CircuitRewrite`. A `CircuitRewrite` is an object that contains all of the information to apply a particular subgraph replacement on a HUGR, but we are _not applying the rewrite_ yet. The only information that a `PatternMatch` is missing to become a `CircuitRewrite` is the `RHS` that we want to replace the match with. To identify this, we make use of the `PatternID` of the match:

```

let mut rewrites = vec![];

for m in matches {

    let rule_id = m.pattern_id().0;

    let rhs = &rules[rule_id].1;



    let rw = m.to_rewrite(circ, rhs.clone()).unwrap();

    rewrites.push(rw);

};

```

For each match `m` we are extracting its `PatternID`, which identifies the index corresponding to the match in the vector `lhs_of_rules` we passed to `PatternMatcher::from_patterns(lhs_of_rules)`. The type of `PatternID` is `struct PatternID(usize)`, so it's a struct with a single `usize` argument that we can access via `.0` (as if it were a tuple); the variable `rule_id` contains the integer index we were looking for. Since, by construction, the elements of `lhs_of_rules` follow the same order as `rules`, we can retrieve the rule whose LHS was matched as `rules[rule_id]`. Remember each rule is specified as a pair `(LHS, RHS)` so to obtain the RHS we just need to access the second element of the tuple via `.1`. Finally, we can generate a `CircuitRewrite` from our match `m` by calling `m.to_rewrite(circ, rhs.clone())`, where `.clone()` is necessary due to Rust's borrow checker constraints.



Up to this point we have succesfully identified all instances and locations of rewrites that can be applied to `circ`. But we have not yet applied any of the rewrites. To do so, we just need to call `my_rewrite.apply(circ)`, where `circ` is of type `&mut Hugr` (note: if what you have is just a `Hugr` you'll instead need to use `my_rewrite.apply(&mut circ)`). However, there's a caveat: we can't simply iterate over all of the `CircuitRewrite` in our `rewrites` vector and apply them, because there might be some overlaps in their matches. 



A simple example of the problem outlined above is a circuit `CZ (S \otimes S)` where we want to apply both the rewrite rules `CZ (S \otimes I) => (S \otimes I) CZ` and `CZ (I \otimes S) => (I \otimes S) CZ`: both will match, and the matches will overlap on the `CZ` gate. When a `CircuitRewrite` is applied it replaces all of the nodes in the LHS with the nodes in the RHS, and the HUGR does not know that the `CZ` in the LHS is the "same" `CZ` gate in the RHS. Consequently, once we apply the rewrite corresponding to the first rule, if we try to apply the rewrite of the second rule we will get a `SimpleReplacementError`. Fortunately, this has a simple solution: we can match and apply the second rewrite rule during the next iteration of our exhaustive rewrite strategy. All we have to do is avoid the `SimpleReplacementError` by keeping track of the nodes that have been changed by rewrites we've applied, and skipping any rewrite whose LHS contains any of those nodes:

```

fn apply_non_overlapping(

    rewrites: impl IntoIterator,

    circ: &mut Hugr,

) {

    let rewrites = rewrites.into_iter();

    let mut changed_nodes = HashSet::new();

    for rewrite in rewrites {

        if rewrite  // Skip if it changes a node that has already been changed

            .subcircuit()

            .nodes()

            .iter()

            .any(|n| changed_nodes.contains(n))

        {

            continue;

        }

        // Update the set of changed nodes

        changed_nodes.extend(rewrite.subcircuit().nodes().iter().copied());

        // Apply the rewrite

        rewrite

            .apply(circ)

            .expect("Could not perform rewrite in exhaustive strategy");

    }

}

```

The code above is a simplification of the `GreedyRewriteStrategy` in the `tket2` crate, and is included in the `utils.rs` module of this project.



To complete the implementation of `apply_rules_exhaustively` we just need to call `apply_non_overlapping(rewrites, circ)` and wrap both this and the call to `matcher.find_matches(circ)` in a `while matches.len() > 0`.