feat: Add bindings for tket1-passes from python (#1225)

Adds some pass definitions to `tket.passes`:

- `normalize_guppy(Tk2Circuit, ...) -> Tk2Circuit`
- `clifford_simp(Tk2Circuit, *, allow_swaps = True, traverse_subcircuits
= True) -> Tk2Circuit`
- `squash_phased_rz(Tk2Circuit, *, traverse_subcircuits = True) ->
Tk2Circuit`

The last two use the new bridge to the old tket C++ codebase to run
optimisation on regions of the hugr.

The API here not too ergonomic, we just expose it as a MVP while working
on the API refresh.

To call this from a guppy program output, we'll need to convert the hugr
to the rust representation and back,
```py
from guppylang import guppy
from guppylang.std.angles import angle
from guppylang.std.builtins import result
from guppylang.std.quantum import h, rx, rz, measure, qubit
from hugr import Hugr
from tket.circuit import Tk2Circuit
from tket.passes import normalize_guppy, squash_phasedx_rz

@guppy
def main() -> None:
    q = qubit()

    rz(q, angle(0.1))
    h(q)
    rx(q, angle(-0.1))
    h(q)

    b = measure(q)
    result("b", b)

program = main.compile()
compiler_state = Tk2Circuit.from_bytes(program.to_bytes())
compiler_state = normalize_guppy(compiler_state)
compiler_state = squash_phasedx_rz(compiler_state)
hugr = Hugr.from_str(compiler_state.to_str())
print(hugr.render_dot())
```

drive-by: Add `Tk2Circuit.render_mermaid` method

Author: aborgna-q

Cargo.lock

@@ -27,16 +27,17 @@ tket = { path = "../tket", version = "0.16.0", features = [
 tket-qsystem = { path = "../tket-qsystem", version = "0.22.0" }
 tket1-passes = { path = "../tket1-passes", version = "0.0.0" }
 
-serde = { workspace = true, features = ["derive"] }
-serde_json = { workspace = true }
-tket-json-rs = { workspace = true, features = ["pyo3"] }
-hugr = { workspace = true }
-pyo3 = { workspace = true, features = ["py-clone", "abi3-py310"] }
-num_cpus = { workspace = true }
 derive_more = { workspace = true, features = ["into", "from"] }
+hugr = { workspace = true }
 itertools = { workspace = true }
+num_cpus = { workspace = true }
 portmatching = { workspace = true }
+pyo3 = { workspace = true, features = ["py-clone", "abi3-py310"] }
+rayon = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
 strum = { workspace = true }
+tket-json-rs = { workspace = true, features = ["pyo3"] }
 
 [dev-dependencies]
 rstest = { workspace = true }

tket-py/Cargo.toml

@@ -47,6 +47,7 @@ pub fn module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
         py.get_type::<PyHUGRSerializationError>(),
     )?;
     m.add("TK1EncodeError", py.get_type::<PyTk1EncodeError>())?;
+    m.add("TK1DecodeError", py.get_type::<PyTK1DecodeError>())?;
 
     Ok(m)
 }
@@ -78,13 +79,13 @@ create_py_exception!(
 create_py_exception!(
     tket::serialize::pytket::PytketEncodeError,
     PyTk1EncodeError,
-    "Error type for the conversion between tket and tket1 operations."
+    "Error encoding a HUGR region into a pytket circuit."
 );
 
 create_py_exception!(
     tket::serialize::pytket::PytketDecodeError,
     PyTK1DecodeError,
-    "Error type for the conversion between tket1 and tket operations."
+    "Error decoding a HUGR region from a pytket circuit."
 );
 
 /// Run the validation checks on a circuit.

tket-py/src/circuit.rs

@@ -25,6 +25,7 @@ use derive_more::From;
 use hugr::{Hugr, HugrView, Wire};
 use serde::Serialize;
 use tket::circuit::CircuitHash;
+use tket::modifier::qubit_types_utils::contain_qubit_term;
 use tket::passes::pytket::lower_to_pytket;
 use tket::passes::CircuitChunks;
 use tket::serialize::pytket::{DecodeOptions, EncodeOptions};
@@ -96,11 +97,17 @@ impl Tk2Circuit {
     }
 
     /// Encode the circuit as a HUGR envelope.
-    pub fn to_bytes(&self, config: Bound<'_, PyAny>) -> PyResult<Vec<u8>> {
+    ///
+    /// If no config is given, it defaults to the default binary envelope.
+    #[pyo3(signature = (config = None))]
+    pub fn to_bytes(&self, config: Option<Bound<'_, PyAny>>) -> PyResult<Vec<u8>> {
         fn err(e: impl Display) -> PyErr {
             PyErr::new::<PyAttributeError, _>(format!("Could not encode circuit: {e}"))
         };
-        let config = envelope_config_from_py(config)?;
+        let config = match config {
+            Some(cfg) => envelope_config_from_py(cfg)?,
+            None => EnvelopeConfig::binary(),
+        };
         let mut buf = Vec::new();
         self.circ.store(&mut buf, config).map_err(err)?;
         Ok(buf)
@@ -221,6 +228,7 @@ impl Tk2Circuit {
     pub fn circuit_cost<'py>(&self, cost_fn: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
         let py = cost_fn.py();
         let cost_fn = |op: &OpType| -> PyResult<PyCircuitCost> {
+            // TODO: We should ignore non-tket operations instead.
             let Some(tk2_op) = op.cast::<TketOp>() else {
                 let op_name = op.to_string();
                 return Err(PyErr::new::<PyValueError, _>(format!(
@@ -303,6 +311,10 @@ impl Tk2Circuit {
     fn output_node(&self) -> PyNode {
         self.circ.output_node().into()
     }
+
+    fn render_mermaid(&self) -> String {
+        self.circ.mermaid_string()
+    }
 }
 impl Tk2Circuit {
     /// Tries to extract a Tk2Circuit from a python object.

tket-py/src/circuit/tk2circuit.rs

@@ -1,9 +1,11 @@
 //! Passes for optimising circuits.
 
 pub mod chunks;
+pub mod tket1;
 
 use std::{cmp::min, convert::TryInto, fs, num::NonZeroUsize, path::PathBuf};
 
+use hugr::algorithms::ComposablePass;
 use pyo3::{prelude::*, types::IntoPyDict};
 use tket::optimiser::badger::BadgerOptions;
 use tket::passes;
@@ -24,9 +26,13 @@ pub fn module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
     m.add_function(wrap_pyfunction!(greedy_depth_reduce, &m)?)?;
     m.add_function(wrap_pyfunction!(lower_to_pytket, &m)?)?;
     m.add_function(wrap_pyfunction!(badger_optimise, &m)?)?;
+    m.add_function(wrap_pyfunction!(normalize_guppy, &m)?)?;
     m.add_class::<self::chunks::PyCircuitChunks>()?;
     m.add_function(wrap_pyfunction!(self::chunks::chunks, &m)?)?;
+    m.add_function(wrap_pyfunction!(self::tket1::clifford_simp, &m)?)?;
+    m.add_function(wrap_pyfunction!(self::tket1::squash_phasedx_rz, &m)?)?;
     m.add("PullForwardError", py.get_type::<PyPullForwardError>())?;
+    m.add("TK1PassError", py.get_type::<tket1::PytketPassError>())?;
     Ok(m)
 }
 
@@ -42,6 +48,50 @@ create_py_exception!(
     "Errors that can occur while removing high-level operations from HUGR intended to be encoded as a pytket circuit."
 );
 
+create_py_exception!(
+    tket::passes::guppy::NormalizeGuppyErrors,
+    PyNormalizeGuppyError,
+    "Errors from the Guppy normalization pass."
+);
+
+/// Flatten the structure of a Guppy-generated program to enable additional optimisations.
+///
+/// This should normally be called first before other optimisations.
+///
+/// Parameters:
+/// - simplify_cfgs: Whether to simplify CFG control flow.
+/// - remove_tuple_untuple: Whether to remove tuple/untuple operations.
+/// - constant_folding: Whether to constant fold the program.
+/// - remove_dead_funcs: Whether to remove dead functions.
+/// - inline_dfgs: Whether to inline DFG operations.
+#[pyfunction]
+#[pyo3(signature = (circ, *, simplify_cfgs = true, remove_tuple_untuple = true, constant_folding = false, remove_dead_funcs = true, inline_dfgs = true))]
+fn normalize_guppy<'py>(
+    circ: &Bound<'py, PyAny>,
+    simplify_cfgs: bool,
+    remove_tuple_untuple: bool,
+    constant_folding: bool,
+    remove_dead_funcs: bool,
+    inline_dfgs: bool,
+) -> PyResult<Bound<'py, PyAny>> {
+    let py = circ.py();
+    try_with_circ(circ, |mut circ, typ| {
+        let mut pass = tket::passes::NormalizeGuppy::default();
+
+        pass.simplify_cfgs(simplify_cfgs)
+            .remove_tuple_untuple(remove_tuple_untuple)
+            .constant_folding(constant_folding)
+            .remove_dead_funcs(remove_dead_funcs)
+            .inline_dfgs(inline_dfgs);
+
+        pass.run(circ.hugr_mut()).convert_pyerrs()?;
+
+        let circ = typ.convert(py, circ)?;
+        PyResult::Ok(circ)
+    })
+}
+
+/// Pass which greedily commutes operations forwards in order to reduce depth.
 #[pyfunction]
 fn greedy_depth_reduce<'py>(circ: &Bound<'py, PyAny>) -> PyResult<(Bound<'py, PyAny>, u32)> {
     let py = circ.py();

tket-py/src/passes.rs

@@ -0,0 +1,102 @@
+//! Passes that call to tket1-passes using the tket-c-api.
+
+use rayon::iter::ParallelIterator;
+use std::sync::Arc;
+
+use pyo3::prelude::*;
+use tket::serialize::pytket::{EncodeOptions, EncodedCircuit};
+use tket::Circuit;
+use tket_qsystem::pytket::{qsystem_decoder_config, qsystem_encoder_config};
+
+use crate::circuit::try_with_circ;
+use crate::utils::{create_py_exception, ConvertPyErr};
+
+/// An optimisation pass that applies a number of rewrite rules for simplifying
+/// Clifford gate sequences, similar to Duncan & Fagan
+/// (https://arxiv.org/abs/1901.10114). Produces a circuit comprising TK1 gates
+/// and the two-qubit gate specified as the target.
+///
+/// Parameters:
+/// - allow_swaps: whether the rewriting may introduce implicit wire swaps.
+/// - traverse_subcircuits: Whether to apply the optimisation to nested
+///   subregions in the hugr too, rather than just the top-level region.
+//
+// TODO: We should also expose `target_gate` here, but the most appropriate
+// parameter type [`crate::ops::PyTketOp`] doesn't include `TK2` -.-
+#[pyfunction]
+#[pyo3(signature = (circ, *, allow_swaps = true, traverse_subcircuits = true))]
+pub fn clifford_simp<'py>(
+    circ: &Bound<'py, PyAny>,
+    allow_swaps: bool,
+    traverse_subcircuits: bool,
+) -> PyResult<Bound<'py, PyAny>> {
+    let py = circ.py();
+
+    try_with_circ(circ, |circ, typ| {
+        let circ = run_tket1_pass(circ, traverse_subcircuits, |tk1_circ| {
+            tk1_circ.clifford_simp(tket_json_rs::OpType::CX, allow_swaps)
+        })?;
+
+        let circ = typ.convert(py, circ)?;
+        PyResult::Ok(circ)
+    })
+}
+
+/// Squash single qubit gates into PhasedX and Rz gates. Also remove identity
+/// gates. Commute Rz gates to the back if possible.
+///
+/// Parameters:
+/// - traverse_subcircuits: Whether to apply the optimisation to nested
+///   subregions in the hugr too, rather than just the top-level region.
+#[pyfunction]
+#[pyo3(signature = (circ, *, traverse_subcircuits = true))]
+pub fn squash_phasedx_rz<'py>(
+    circ: &Bound<'py, PyAny>,
+    traverse_subcircuits: bool,
+) -> PyResult<Bound<'py, PyAny>> {
+    let py = circ.py();
+
+    try_with_circ(circ, |circ, typ| {
+        let circ = run_tket1_pass(circ, traverse_subcircuits, |tk1_circ| {
+            tk1_circ.squash_phasedx_rz()
+        })?;
+
+        let circ = typ.convert(py, circ)?;
+        PyResult::Ok(circ)
+    })
+}
+
+fn run_tket1_pass<F>(mut circ: Circuit, traverse_subcircuits: bool, pass: F) -> PyResult<Circuit>
+where
+    F: Fn(&mut tket1_passes::Tket1Circuit) -> Result<(), tket1_passes::PassError> + Send + Sync,
+{
+    let mut encoded_circ = EncodedCircuit::new(
+        &circ,
+        EncodeOptions::new()
+            .with_config(qsystem_encoder_config())
+            .with_subcircuits(traverse_subcircuits),
+    )
+    .convert_pyerrs()?;
+
+    encoded_circ
+        .par_iter_mut()
+        .try_for_each(|(_, circ)| -> Result<(), tket1_passes::PassError> {
+            let mut tk1_circ = tket1_passes::Tket1Circuit::from_serial_circuit(circ)?;
+            pass(&mut tk1_circ)?;
+            *circ = tk1_circ.to_serial_circuit()?;
+            Ok(())
+        })
+        .convert_pyerrs()?;
+
+    encoded_circ
+        .reassemble_inplace(circ.hugr_mut(), Some(Arc::new(qsystem_decoder_config())))
+        .convert_pyerrs()?;
+
+    Ok(circ)
+}
+
+create_py_exception!(
+    tket1_passes::PassError,
+    PytketPassError,
+    "Error from a call to tket-c-api"
+);

tket-py/src/passes/tket1.rs

@@ -1,7 +1,15 @@
 from pytket import Circuit, OpType
 from dataclasses import dataclass
 from typing import Callable, Any
-from tket.passes import badger_pass, greedy_depth_reduce, chunks
+from tket.ops import TketOp
+from tket.passes import (
+    badger_pass,
+    greedy_depth_reduce,
+    chunks,
+    clifford_simp,
+    normalize_guppy,
+    squash_phasedx_rz,
+)
 from tket.circuit import Tk2Circuit
 from tket.pattern import Rule, RuleMatcher
 import hypothesis.strategies as st
@@ -138,3 +146,35 @@ def test_multiple_rules():
 
     out = circ.to_tket1()
     assert out == Circuit(3).CX(0, 1).X(0)
+
+
+def test_clifford_simp():
+    c = Tk2Circuit(Circuit(4).CX(0, 2).CX(1, 2).CX(1, 2))
+
+    c = clifford_simp(c, allow_swaps=False)
+
+    assert c.circuit_cost(lambda op: int(op == TketOp.CX)) == 1
+
+
+def test_squash_phasedx_rz():
+    c = Tk2Circuit(Circuit(1).Rz(0.25, 0).Rz(0.75, 0).Rz(0.25, 0).Rz(-1.25, 0))
+
+    c = squash_phasedx_rz(c)
+
+    # TODO: We cannot use circuit_cost due to a panic on non-tket ops and there
+    # being some parameter loads...
+    assert c.num_operations() == 0
+
+
+def test_normalize_guppy():
+    """Test the normalize_guppy pass.
+
+    This won't actually do anything useful, we just want to check that the pass
+    runs without errors.
+    """
+
+    c = Tk2Circuit(Circuit(4).CX(0, 2).CX(1, 2).CX(1, 2))
+
+    c = normalize_guppy(c)
+
+    assert c.circuit_cost(lambda op: int(op == TketOp.CX)) == 3

tket-py/test/test_pass.py

@@ -71,7 +71,7 @@ class Tk2Circuit:
     def from_hugr_json(json: str) -> Tk2Circuit:
         """Decode a HUGR json string to a Tk2Circuit."""
 
-    def to_bytes(self, config: EnvelopeConfig) -> bytes:
+    def to_bytes(self, config: EnvelopeConfig | None = None) -> bytes:
         """Encode the circuit as a HUGR envelope, according to the given config.
 
         Some envelope formats can be encoded into a string. See :meth:`to_str`.
@@ -150,6 +150,9 @@ class Tk2Circuit:
     def from_tket1_json(json: str) -> Tk2Circuit:
         """Decode a pytket json string to a Tk2Circuit."""
 
+    def render_mermaid(self) -> str:
+        """Render the circuit as a Mermaid graph."""
+
 class Node:
     """Handle to node in HUGR."""