feat: Add Python TASO pass

pyrs/Cargo.toml

@@ -18,3 +18,4 @@ tket-json-rs = { git = "https://github.com/CQCL/tket-json-rs", rev="619db15d3",
 quantinuum-hugr = { workspace = true }
 portgraph = { workspace = true, features = ["pyo3"] }
 pyo3 = { workspace = true, features = ["extension-module"] }
+num_cpus = "1.16.0"

pyrs/src/lib.rs

@@ -1,19 +1,20 @@
 //! Python bindings for TKET2.
 #![warn(missing_docs)]
-use circuit::{add_circuit_module, to_hugr, try_with_hugr, T2Circuit};
-use hugr::Hugr;
+
+mod circuit;
+mod optimiser;
+mod pass;
+
+use circuit::{add_circuit_module, to_hugr, T2Circuit};
 use optimiser::add_optimiser_module;
+use pass::add_pass_module;
+
+use hugr::Hugr;
 use pyo3::prelude::*;
 use tket2::{
-    json::TKETDecode,
-    passes::apply_greedy_commutation,
     portmatching::{pyo3::PyPatternMatch, CircuitPattern, PatternMatcher},
     rewrite::CircuitRewrite,
 };
-use tket_json_rs::circuit_json::SerialCircuit;
-
-mod circuit;
-mod optimiser;
 
 #[derive(Clone)]
 #[pyclass]
@@ -63,15 +64,6 @@ impl RuleMatcher {
     }
 }
 
-#[pyfunction]
-fn greedy_depth_reduce(py_c: PyObject) -> PyResult<(PyObject, u32)> {
-    try_with_hugr(py_c, |mut h| {
-        let n_moves = apply_greedy_commutation(&mut h)?;
-        let py_c = SerialCircuit::encode(&h)?.to_tket1()?;
-        PyResult::Ok((py_c, n_moves))
-    })
-}
-
 /// The Python bindings to TKET2.
 #[pymodule]
 fn pyrs(py: Python, m: &PyModule) -> PyResult<()> {
@@ -102,15 +94,3 @@ fn add_pattern_module(py: Python, parent: &PyModule) -> PyResult<()> {
 
     parent.add_submodule(m)
 }
-
-fn add_pass_module(py: Python, parent: &PyModule) -> PyResult<()> {
-    let m = PyModule::new(py, "passes")?;
-    m.add_function(wrap_pyfunction!(greedy_depth_reduce, m)?)?;
-    m.add_class::<tket2::T2Op>()?;
-    m.add(
-        "PullForwardError",
-        py.get_type::<tket2::passes::PyPullForwardError>(),
-    )?;
-    parent.add_submodule(m)?;
-    Ok(())
-}

pyrs/src/optimiser.rs

@@ -3,6 +3,7 @@
 use std::io::BufWriter;
 use std::{fs, num::NonZeroUsize, path::PathBuf};
 
+use hugr::Hugr;
 use pyo3::prelude::*;
 use tket2::optimiser::{DefaultTasoOptimiser, TasoLogger};
 
@@ -56,7 +57,8 @@ impl PyDefaultTasoOptimiser {
     ///   `n_threads` chunks and optimise each on a separate thread.
     /// * `log_progress`: The path to a CSV file to log progress to.
     ///
-    pub fn optimise(
+    #[pyo3(name = "optimise")]
+    pub fn py_optimise(
         &self,
         circ: PyObject,
         timeout: Option<u64>,
@@ -65,22 +67,44 @@ impl PyDefaultTasoOptimiser {
         log_progress: Option<PathBuf>,
         queue_size: Option<usize>,
     ) -> PyResult<PyObject> {
+        update_hugr(circ, |circ| {
+            self.optimise(
+                circ,
+                timeout,
+                n_threads,
+                split_circ,
+                log_progress,
+                queue_size,
+            )
+        })
+    }
+}
+
+impl PyDefaultTasoOptimiser {
+    /// The Python optimise method, but on Hugrs.
+    pub(super) fn optimise(
+        &self,
+        circ: Hugr,
+        timeout: Option<u64>,
+        n_threads: Option<NonZeroUsize>,
+        split_circ: Option<bool>,
+        log_progress: Option<PathBuf>,
+        queue_size: Option<usize>,
+    ) -> Hugr {
         let taso_logger = log_progress
             .map(|file_name| {
                 let log_file = fs::File::create(file_name).unwrap();
                 let log_file = BufWriter::new(log_file);
                 TasoLogger::new(log_file)
             })
             .unwrap_or_default();
-        update_hugr(circ, |circ| {
-            self.0.optimise_with_log(
-                &circ,
-                taso_logger,
-                timeout,
-                n_threads.unwrap_or(NonZeroUsize::new(1).unwrap()),
-                split_circ.unwrap_or(false),
-                queue_size.unwrap_or(100),
-            )
-        })
+        self.0.optimise_with_log(
+            &circ,
+            taso_logger,
+            timeout,
+            n_threads.unwrap_or(NonZeroUsize::new(1).unwrap()),
+            split_circ.unwrap_or(false),
+            queue_size.unwrap_or(100),
+        )
     }
 }

pyrs/src/pass.rs

@@ -0,0 +1,147 @@
+use std::{cmp::min, convert::TryInto, fs, num::NonZeroUsize, path::PathBuf};
+
+use pyo3::{prelude::*, types::IntoPyDict};
+use tket2::{json::TKETDecode, op_matches, passes::apply_greedy_commutation, Circuit, T2Op};
+use tket_json_rs::circuit_json::SerialCircuit;
+
+use crate::{
+    circuit::{try_update_hugr, try_with_hugr},
+    optimiser::PyDefaultTasoOptimiser,
+};
+
+#[pyfunction]
+fn greedy_depth_reduce(py_c: PyObject) -> PyResult<(PyObject, u32)> {
+    try_with_hugr(py_c, |mut h| {
+        let n_moves = apply_greedy_commutation(&mut h)?;
+        let py_c = SerialCircuit::encode(&h)?.to_tket1()?;
+        PyResult::Ok((py_c, n_moves))
+    })
+}
+
+/// TASO optimisation pass.
+///
+/// HyperTKET's best attempt at optimising a circuit using circuit rewriting
+/// and TASO.
+///
+/// Input can be in any gate set and will be rebased to Nam, i.e. CX + Rz + H.
+///
+/// Will use at most `max_threads` threads (plus a constant) and take at most
+/// `timeout` seconds (plus a constant). Default to the number of cpus and
+/// 15min respectively.
+///
+/// Log files will be written to the directory `log_dir` if specified.
+///
+/// This requires a `nam_6_3.rwr` file in the current directory. The location
+/// can alternatively be specified using the `rewriter_dir` argument.
+#[pyfunction]
+fn taso_optimise(
+    circ: PyObject,
+    max_threads: Option<NonZeroUsize>,
+    rewriter_dir: Option<PathBuf>,
+    timeout: Option<u64>,
+    log_dir: Option<PathBuf>,
+) -> PyResult<PyObject> {
+    // Runs the following code:
+    // ```python
+    // from pytket.passes.auto_rebase import auto_rebase_pass
+    // from pytket import OpType
+    // auto_rebase_pass({OpType.CX, OpType.Rz, OpType.H}).apply(circ)"
+    // ```
+    Python::with_gil(|py| {
+        let auto_rebase = py
+            .import("pytket.passes.auto_rebase")?
+            .getattr("auto_rebase_pass")?;
+        let optype = py.import("pytket")?.getattr("OpType")?;
+        let locals = [("OpType", &optype)].into_py_dict(py);
+        let op_set = py.eval("{OpType.CX, OpType.Rz, OpType.H}", None, Some(locals))?;
+        let rebase_pass = auto_rebase.call1((op_set,))?.getattr("apply")?;
+        rebase_pass.call1((&circ,)).map(|_| ())
+    })?;
+    nam_taso_optimise(circ, max_threads, rewriter_dir, timeout, log_dir)
+}
+
+/// TASO optimisation pass.
+///
+/// HyperTKET's best attempt at optimising a circuit using circuit rewriting
+/// and TASO.
+///
+/// Input must be in the Nam gate set, i.e. CX + Rz + H.
+///
+/// Will use at most `max_threads` threads (plus a constant) and take at most
+/// `timeout` seconds (plus a constant). Default to the number of cpus and
+/// 30s respectively.
+///
+/// Log files will be written to the directory `log_dir` if specified.
+///
+/// This requires a `nam_6_3.rwr` file in the current directory. The location
+/// can alternatively be specified using the `rewriter_dir` argument.
+#[pyfunction]
+fn nam_taso_optimise(
+    circ: PyObject,
+    max_threads: Option<NonZeroUsize>,
+    rewriter_dir: Option<PathBuf>,
+    timeout: Option<u64>,
+    log_dir: Option<PathBuf>,
+) -> PyResult<PyObject> {
+    let max_threads = max_threads.unwrap_or(num_cpus::get().try_into().unwrap());
+    let rewrite_dir = rewriter_dir.unwrap_or(PathBuf::from("."));
+    let timeout = timeout.unwrap_or(30);
+    if let Some(log_dir) = log_dir.as_ref() {
+        fs::create_dir_all(log_dir)?;
+    }
+    let taso_splits = |n_threads: NonZeroUsize| match n_threads.get() {
+        n if n >= 7 => (
+            vec![n, 3, 1],
+            vec![timeout / 2, timeout / 10 * 3, timeout / 10 * 2],
+        ),
+        n if n >= 4 => (
+            vec![n, 2, 1],
+            vec![timeout / 2, timeout / 10 * 3, timeout / 10 * 2],
+        ),
+        n if n > 1 => (vec![n, 1], vec![timeout / 2, timeout / 2]),
+        1 => (vec![1], vec![timeout]),
+        _ => unreachable!(),
+    };
+    let optimiser = PyDefaultTasoOptimiser::load_precompiled(rewrite_dir.join("nam_6_3.rwr"));
+    try_update_hugr(circ, |mut circ| {
+        let n_cx = circ
+            .commands()
+            .filter(|c| op_matches(c.optype(), T2Op::CX))
+            .count();
+        let n_threads = min(
+            (n_cx / 50).try_into().unwrap_or(1.try_into().unwrap()),
+            max_threads,
+        );
+        let (split_threads, split_timeouts) = taso_splits(n_threads);
+        for (i, (n_threads, timeout)) in split_threads.into_iter().zip(split_timeouts).enumerate() {
+            let log_file = log_dir.as_ref().map(|log_dir| {
+                let mut log_file = log_dir.clone();
+                log_file.push(format!("cycle-{i}.log"));
+                log_file
+            });
+            circ = optimiser.optimise(
+                circ,
+                Some(timeout),
+                Some(n_threads.try_into().unwrap()),
+                Some(true),
+                log_file,
+                None,
+            );
+        }
+        PyResult::Ok(circ)
+    })
+}
+
+pub(crate) fn add_pass_module(py: Python, parent: &PyModule) -> PyResult<()> {
+    let m = PyModule::new(py, "passes")?;
+    m.add_function(wrap_pyfunction!(greedy_depth_reduce, m)?)?;
+    m.add_function(wrap_pyfunction!(nam_taso_optimise, m)?)?;
+    m.add_function(wrap_pyfunction!(taso_optimise, m)?)?;
+    m.add_class::<tket2::T2Op>()?;
+    m.add(
+        "PullForwardError",
+        py.get_type::<tket2::passes::PyPullForwardError>(),
+    )?;
+    parent.add_submodule(m)?;
+    Ok(())
+}

pyrs/test/test_optimiser.py

@@ -0,0 +1,14 @@
+from pytket import Circuit
+from pyrs.pyrs import optimiser
+
+
+def test_simple_optimiser():
+    """ a simple circuit matching test """
+    c = Circuit(3).CX(0, 1).CX(0, 1).CX(1, 2)
+    opt = optimiser.TasoOptimiser.compile_eccs("../test_files/cx_cx_eccs.json")
+
+    # optimise for 1s
+    cc = opt.optimise(c, 3)
+    exp_c = Circuit(3).CX(1, 2)
+
+    assert cc == exp_c

pyrs/test/test_pass.py

@@ -0,0 +1,9 @@
+from pytket import Circuit, OpType
+from pyrs.pyrs import passes
+
+
+def test_simple_taso_pass_no_opt():
+    c = Circuit(3).CCX(0, 1, 2)
+    c = passes.taso_optimise(c, max_threads = 1, timeout = 0)
+    print(c)
+    assert c.n_gates_of_type(OpType.CX) == 6

src/lib.rs

@@ -21,4 +21,4 @@ pub mod portmatching;
 mod utils;
 
 pub use circuit::Circuit;
-pub use ops::{symbolic_constant_op, Pauli, T2Op};
+pub use ops::{op_matches, symbolic_constant_op, Pauli, T2Op};

src/ops.rs

@@ -72,7 +72,7 @@ pub enum T2Op {
 }
 
 /// Whether an op is a given T2Op.
-pub(crate) fn op_matches(op: &OpType, t2op: T2Op) -> bool {
+pub fn op_matches(op: &OpType, t2op: T2Op) -> bool {
     let Ok(op) = T2Op::try_from(op) else {
         return false;
     };