Merge branch 'main' into chore/jira-priorities

Cargo.toml

@@ -69,7 +69,7 @@ harness = false
 
 [workspace]
 
-members = ["pyrs", "compile-matcher", "taso-optimiser"]
+members = ["pyrs", "compile-rewriter", "taso-optimiser"]
 
 [workspace.dependencies]
 

compile-rewriter/src/main.rs

@@ -0,0 +1,75 @@
+use std::fs;
+use std::path::Path;
+use std::process::exit;
+use std::time::Instant;
+
+use clap::Parser;
+
+use tket2::rewrite::ECCRewriter;
+
+/// Program to precompile patterns from files into a PatternMatcher stored as binary file.
+#[derive(Parser, Debug)]
+#[clap(version = "1.0", long_about = None)]
+#[clap(
+    about = "Precompiles ECC sets into a TKET2 Rewriter. The resulting binary files can be loaded into TKET2 for circuit optimisation."
+)]
+struct CmdLineArgs {
+    // TODO: Differentiate between TK1 input and ECC input
+    /// Name of input file/folder
+    #[arg(
+        short,
+        long,
+        value_name = "FILE",
+        help = "Sets the input file to use. It must be a JSON file of ECC sets in the Quartz format."
+    )]
+    input: String,
+    /// Name of output file/folder
+    #[arg(
+        short,
+        long,
+        value_name = "FILE",
+        default_value = ".",
+        help = "Sets the output file or folder. Defaults to \"matcher.rwr\" if no file name is provided. The extension of the file name will always be set or amended to be `.rwr`."
+    )]
+    output: String,
+}
+
+fn main() {
+    let opts = CmdLineArgs::parse();
+
+    let input_path = Path::new(&opts.input);
+    let output_path = Path::new(&opts.output);
+
+    if !input_path.is_file() || input_path.extension().unwrap() != "json" {
+        panic!("Input must be a JSON file");
+    };
+    let start_time = Instant::now();
+    println!("Compiling rewriter...");
+    let Ok(rewriter) = ECCRewriter::try_from_eccs_json_file(input_path) else {
+        eprintln!(
+            "Unable to load ECC file {:?}. Is it a JSON file of Quartz-generated ECCs?",
+            input_path
+        );
+        exit(1);
+    };
+    println!("Saving to file...");
+    let output_file = if output_path.is_dir() {
+        output_path.join("matcher.rwr")
+    } else {
+        output_path.to_path_buf()
+    };
+    let output_file = rewriter.save_binary(output_file).unwrap();
+    println!("Written rewriter to {:?}", output_file);
+
+    // Print the file size of output_file in megabytes
+    if let Ok(metadata) = fs::metadata(&output_file) {
+        let file_size = metadata.len() as f64 / (1024.0 * 1024.0);
+        println!("File size: {:.2} MB", file_size);
+    }
+    let elapsed = start_time.elapsed();
+    println!(
+        "Done in {}.{:03} seconds",
+        elapsed.as_secs(),
+        elapsed.subsec_millis()
+    );
+}

pyrs/src/lib.rs

@@ -1,11 +1,13 @@
 //! Python bindings for TKET2.
 #![warn(missing_docs)]
 use circuit::{add_circuit_module, try_with_hugr};
+use optimiser::add_optimiser_module;
 use pyo3::prelude::*;
 use tket2::{json::TKETDecode, passes::apply_greedy_commutation};
 use tket_json_rs::circuit_json::SerialCircuit;
 
 mod circuit;
+mod optimiser;
 
 #[pyfunction]
 fn greedy_depth_reduce(py_c: PyObject) -> PyResult<(PyObject, u32)> {
@@ -22,6 +24,7 @@ fn pyrs(py: Python, m: &PyModule) -> PyResult<()> {
     add_circuit_module(py, m)?;
     add_pattern_module(py, m)?;
     add_pass_module(py, m)?;
+    add_optimiser_module(py, m)?;
     Ok(())
 }
 

pyrs/src/optimiser.rs

@@ -0,0 +1,67 @@
+//! PyO3 wrapper for the TASO circuit optimiser.
+
+use std::{fs, num::NonZeroUsize, path::PathBuf};
+
+use pyo3::prelude::*;
+use tket2::optimiser::{DefaultTasoOptimiser, TasoLogger};
+
+use crate::circuit::update_hugr;
+
+/// The circuit optimisation module.
+pub fn add_optimiser_module(py: Python, parent: &PyModule) -> PyResult<()> {
+    let m = PyModule::new(py, "optimiser")?;
+    m.add_class::<PyDefaultTasoOptimiser>()?;
+
+    parent.add_submodule(m)
+}
+
+/// Wrapped [`DefaultTasoOptimiser`].
+///
+/// Currently only exposes loading from an ECC file using the constructor
+/// and optimising using default logging settings.
+#[pyclass(name = "TasoOptimiser")]
+pub struct PyDefaultTasoOptimiser(DefaultTasoOptimiser);
+
+#[pymethods]
+impl PyDefaultTasoOptimiser {
+    /// Create a new [`PyDefaultTasoOptimiser`] from a precompiled rewriter.
+    #[staticmethod]
+    pub fn load_precompiled(path: PathBuf) -> Self {
+        Self(DefaultTasoOptimiser::default_with_rewriter_binary(path).unwrap())
+    }
+
+    /// Create a new [`PyDefaultTasoOptimiser`] from ECC sets.
+    ///
+    /// This will compile the rewriter from the provided ECC JSON file.
+    #[staticmethod]
+    pub fn compile_eccs(path: &str) -> Self {
+        Self(DefaultTasoOptimiser::default_with_eccs_json_file(path).unwrap())
+    }
+
+    /// Run the optimiser on a circuit.
+    ///
+    /// Returns an optimised circuit and optionally log the progress to a CSV
+    /// file.
+    pub fn optimise(
+        &self,
+        circ: PyObject,
+        timeout: Option<u64>,
+        n_threads: Option<NonZeroUsize>,
+        log_progress: Option<PathBuf>,
+    ) -> PyResult<PyObject> {
+        let taso_logger = log_progress
+            .map(|file_name| {
+                let log_file = fs::File::create(file_name).unwrap();
+                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()),
+            )
+        })
+    }
+}