feat: Pauli Graph Resynthesis (#2020)

Implements a tket2 pass that takes flat HUGR regions, converts them to
pauli graphs, applies the greedy resynth pass, and converts the
resulting pauli graph back into a HUGR.

Note: currently the result of greedy resynth does not simplify single
qubit gate sequences so it might be useful to add a simple single qubit
gate merging pass after resynthesis in the future (`RotationMergingPass`
is only for pauli graphs in canonical form)

---------

Co-authored-by: Silas Dilkes <36165522+sjdilkes@users.noreply.github.com>
Co-authored-by: Alec Edgington <54802828+cqc-alec@users.noreply.github.com>

Author: 3 people

Cargo.lock

@@ -121,6 +121,11 @@ cool_asserts = "2.0.4"
 zstd = "0.14.0"
 anyhow = "1.0.102"
 num-rational = "0.4.2"
+pg-core = { package = "tk-pg-core", path = "pg-libs/tk-pg-core", version = "0.1.0" }
+pg-canonical-form = { package = "tk-pg-canonical-form", path = "pg-libs/pg-passes/tk-pg-canonical-form", version = "0.1.0" }
+pg-greedy-synth = { package = "tk-pg-greedy-synth", path = "pg-libs/pg-passes/tk-pg-greedy-synth", version = "0.1.0" }
+pg-optimise = { package = "tk-pg-optimize", path = "pg-libs/pg-passes/tk-pg-optimize", version = "0.1.0" }
+pg-rebase = { package = "tk-pg-rebase", path = "pg-libs/pg-passes/tk-pg-rebase", version = "0.1.0" }
 
 [profile.release.package.tket-py]
 # Some configurations to reduce the size of tket wheels

Cargo.toml

@@ -28,6 +28,7 @@ tket-qsystem = { path = "../tket-qsystem", version = "0.29.0" }
 tket1-passes = { path = "../tket1-passes", version = "0.0.0" }
 
 derive_more = { workspace = true, features = ["into", "from"] }
+pg-greedy-synth = { workspace = true }
 hugr = { workspace = true }
 itertools = { workspace = true }
 num_cpus = { workspace = true }

tket/tket-py/Cargo.toml

@@ -34,6 +34,7 @@ pub fn module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
     m.add_class::<self::chunks::PyCircuitChunks>()?;
     m.add_function(wrap_pyfunction!(self::chunks::chunks, &m)?)?;
     m.add_function(wrap_pyfunction!(self::tket1::tket1_pass, &m)?)?;
+    m.add_function(wrap_pyfunction!(pauli_graph_resynthesis, &m)?)?;
     m.add_function(wrap_pyfunction!(resolve_modifiers, &m)?)?;
     m.add_function(wrap_pyfunction!(qsystem::qsystem_rebase_pass, &m)?)?;
     m.add_function(wrap_pyfunction!(qsystem::qsystem_llvm_pass, &m)?)?;
@@ -43,6 +44,10 @@ pub fn module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
         py.get_type::<PyInlineFunctionsError>(),
     )?;
     m.add("TK1PassError", py.get_type::<tket1::PytketPassError>())?;
+    m.add(
+        "PauliGraphResynthesisError",
+        py.get_type::<PauliGraphResynthesisError>(),
+    )?;
     Ok(m)
 }
 
@@ -64,6 +69,12 @@ create_py_exception!(
     "Errors from the modifer resolver pass."
 );
 
+create_py_exception!(
+    tket::passes::pauli_graph_resynthesis::PauliGraphResynthesisErrors,
+    PauliGraphResynthesisError,
+    "Errors from the Pauli graph resynthesis pass."
+);
+
 create_py_exception!(
     tket::passes::inline_funcs::InlineFuncsError,
     PyInlineFunctionsError,
@@ -230,3 +241,61 @@ fn resolve_modifiers(circ: &mut CompilationState, scope: Option<PyPassScope>) ->
     pass.run(&mut circ.hugr).convert_pyerrs()?;
     Ok(())
 }
+
+struct PyParallelMode(pg_greedy_synth::ParallelMode);
+
+impl<'a, 'py> FromPyObject<'a, 'py> for PyParallelMode {
+    type Error = PyErr;
+
+    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
+        use pg_greedy_synth::ParallelMode;
+
+        let enum_type = ob.py().import("tket.passes")?.getattr("ParallelMode")?;
+
+        for (name, mode) in [
+            ("Auto", ParallelMode::Auto),
+            ("On", ParallelMode::On),
+            ("Off", ParallelMode::Off),
+        ] {
+            if ob.is(&enum_type.getattr(name)?) {
+                return Ok(Self(mode));
+            }
+        }
+
+        Err(pyo3::exceptions::PyTypeError::new_err(
+            "parallel_mode must be an instance of the ParallelMode enum",
+        ))
+    }
+}
+
+#[pyfunction]
+#[pyo3(signature = (circ, scope = None, window_size=None, pool_size=None, top_up_size=None, seed=None, parallel_mode=None))]
+fn pauli_graph_resynthesis(
+    circ: &mut CompilationState,
+    scope: Option<PyPassScope>,
+    window_size: Option<usize>,
+    pool_size: Option<usize>,
+    top_up_size: Option<usize>,
+    seed: Option<usize>,
+    parallel_mode: Option<PyParallelMode>,
+) -> PyResult<()> {
+    let py_scope = scope.unwrap_or_default();
+    let mut pass = tket::passes::PauliGraphResynthesis::default_with_scope(py_scope.scope);
+    if let Some(ws) = window_size {
+        pass = pass.with_window_size(ws);
+    }
+    if let Some(ps) = pool_size {
+        pass = pass.with_pool_size(ps);
+    }
+    if let Some(tus) = top_up_size {
+        pass = pass.with_top_up_size(tus);
+    }
+    if let Some(s) = seed {
+        pass = pass.with_seed(s as u64);
+    }
+    let parallel_mode = parallel_mode.map_or(pg_greedy_synth::ParallelMode::Auto, |mode| mode.0);
+    pass = pass.with_parallel_mode(parallel_mode);
+
+    pass.run(&mut circ.hugr).convert_pyerrs()?;
+    Ok(())
+}

tket/tket-py/src/passes.rs

@@ -4,6 +4,7 @@
 from collections.abc import Callable
 from pathlib import Path
 from typing import Any
+from unittest.mock import ANY, patch
 
 import hypothesis.strategies as st
 import numpy as np
@@ -17,12 +18,16 @@
 from tket._ops import TketOp
 from tket._pattern import Rule, RuleMatcher
 from tket._state import CompilationState
+from tket._state.build import H, from_coms
+from tket._tket import passes as rust_passes
 from tket.passes import (
     GlobalScope,
     InlineFunctions,
     ModifierResolverPass,
     Normalize,
     NormalizeGuppy,
+    ParallelMode,
+    PauliGraphResynthesis,
     PlatformTarget,
     PytketHugrPass,
     QSystemRebasePass,
@@ -507,3 +512,94 @@ def test_python_qsystem_pass_with_modifiers() -> None:
         except Exception as exc:  # noqa: BLE001
             failures.append(f"{hugr_path}: {exc}")
     assert not failures, "QSystem pass failures:\n" + "\n".join(failures)
+
+
+def _count_hadamards(hugr: Hugr) -> int:
+    return sum(
+        data.op.name().rsplit(".", maxsplit=1)[-1] == "H" for _, data in hugr.nodes()
+    )
+
+
+def test_resynthesis_with_default_options() -> None:
+    hugr = from_coms(H(0), H(0)).to_python().modules[0]
+    optimisation = PauliGraphResynthesis()
+
+    result = optimisation.run(hugr, inplace=False)
+
+    assert optimisation.parallel_mode is ParallelMode.Auto
+    assert optimisation.window_size is None
+    assert optimisation.pool_size is None
+    assert optimisation.top_up_size is None
+    assert optimisation.seed is None
+    assert result.results == [("PauliGraphResynthesis", None)]
+    assert _count_hadamards(result.hugr) == 0
+    assert _count_hadamards(hugr) == 2
+    CompilationState.from_python(result.hugr).validate()
+
+
+def test_custom_options_and_scope() -> None:
+    hugr = from_coms(H(0), H(0)).to_python().modules[0]
+    optimisation = PauliGraphResynthesis(
+        window_size=16,
+        pool_size=32,
+        top_up_size=4,
+        seed=7,
+        parallel_mode=ParallelMode.On,
+    )
+    assert optimisation.with_scope(GlobalScope.PRESERVE_ALL) is optimisation
+
+    with patch.object(
+        rust_passes,
+        "pauli_graph_resynthesis",
+        wraps=rust_passes.pauli_graph_resynthesis,
+    ) as resynthesis:
+        result = optimisation.run(hugr, inplace=False)
+
+    resynthesis.assert_called_once_with(
+        ANY,
+        scope=GlobalScope.PRESERVE_ALL,
+        window_size=16,
+        pool_size=32,
+        top_up_size=4,
+        seed=7,
+        parallel_mode=ParallelMode.On,
+    )
+    assert _count_hadamards(result.hugr) == 0
+
+
+@pytest.mark.parametrize(
+    ("parameter", "value", "message"),
+    [
+        ("window_size", -1, "window_size must be positive"),
+        ("window_size", 0, "window_size must be positive"),
+        ("pool_size", -1, "pool_size must be positive"),
+        ("pool_size", 0, "pool_size must be positive"),
+        ("top_up_size", -1, "top_up_size must be positive"),
+        ("top_up_size", 0, "top_up_size must be positive"),
+        ("seed", -1, "seed must be non-negative"),
+    ],
+)
+def test_resynthesis_rejects_invalid_parameters(
+    parameter: str, value: int, message: str
+) -> None:
+    hugr = from_coms(H(0), H(0)).to_python().modules[0]
+    options: dict[str, Any] = {parameter: value}
+
+    with pytest.raises(ValueError, match=f"^{message}$"):
+        PauliGraphResynthesis(**options)
+
+    optimisation = PauliGraphResynthesis()
+    setattr(optimisation, parameter, value)
+
+    with pytest.raises(ValueError, match=f"^{message}$"):
+        optimisation.run(hugr, inplace=True)
+
+    assert _count_hadamards(hugr) == 2
+
+
+def test_wrapper_requires_enum() -> None:
+    invalid_mode: Any = "on"
+    with pytest.raises(
+        TypeError, match="parallel_mode must be an instance of the ParallelMode enum"
+    ):
+        PauliGraphResynthesis(parallel_mode=invalid_mode)

tket/tket-py/test/test_pass.py

@@ -3,7 +3,7 @@ from typing import Literal
 
 from hugr.passes.scope import GlobalScope, PassScope
 
-from ..passes import inline_funcs
+from ..passes import ParallelMode, inline_funcs
 from .optimiser import BadgerOptimiser
 from .state import CompilationState
 
@@ -20,6 +20,9 @@ class CircuitChunks:
 class PullForwardError(Exception):
     """Error from a `PullForward` operation."""
 
+class PauliGraphResynthesisError(Exception):
+    """Error from the Pauli graph resynthesis pass."""
+
 # ruff: ignore[B008]
 def normalize_guppy(
     circ: CompilationState,
@@ -102,6 +105,36 @@ def tket1_pass(
       Defaults to the platform-agnostic ``"tket"`` target.
     """
 
+def pauli_graph_resynthesis(
+    circ: CompilationState,
+    scope: PassScope = GlobalScope.PRESERVE_PUBLIC,
+    window_size: int | None = None,
+    pool_size: int | None = None,
+    top_up_size: int | None = None,
+    seed: int | None = None,
+    parallel_mode: ParallelMode | None = None,
+) -> None:
+    """
+    Resynthesizes a Clifford + Rz circuit by converting it to a Pauli graph and applying various
+    optimisation techniques such as:
+    - phase folding
+    - a synthesis algorithm from pauli graph to Clifford + Rz aimed at reducing the number of 2
+    qubit gates
+
+    Rotation angles must be numeric as symbolic angles are not supported currently.
+
+    Parameters:
+    - window_size: Sets the size of the sliding window used for lookahead during synthesis.
+    - pool_size: Sets the number of candidate gates to maintain in the pool.
+    - top_up_size: Sets the number of candidate gates to add after each TQE gate.
+    - seed: Sets the random seed used to sample candidate gates.
+    - parallel_mode: ParallelMode.Auto, ParallelMode.On, or ParallelMode.Off.
+      Omitted or None selects ParallelMode.Auto.
+
+    Raises:
+        PauliGraphResynthesisError: If the resynthesis pass fails.
+    """
+
 def resolve_modifiers(
     circ: CompilationState, scope: PassScope = GlobalScope.PRESERVE_PUBLIC
 ) -> None:

tket/tket-py/tket/_tket/passes.pyi

@@ -35,14 +35,24 @@
     "ModifierResolverPass",
     "Normalize",
     "NormalizeGuppy",
+    "ParallelMode",
     "PassResult",
+    "PauliGraphResynthesis",
     "PlatformTarget",
     "PytketHugrPass",
     "QSystemRebasePass",
     "_QSystemLLVMPass",
 ]
 
 
+class ParallelMode(Enum):
+    """Parallel processing mode for Pauli graph resynthesis."""
+
+    Auto = "Auto"  # Let synthesis choose when to use parallel processing.
+    On = "On"  # Enable parallel processing.
+    Off = "Off"  # Disable parallel processing.
+
+
 class PlatformTarget(Enum):
     """A hardware platform that passes can target.
 
@@ -505,6 +515,79 @@ def _run_tk(self, program: _state.CompilationState) -> _state.CompilationState:
         return program
 
 
+@dataclass
+class PauliGraphResynthesis(ComposablePass):
+    """
+    An optimisation pass that resynthesizes a Clifford + Rz circuit by converting it to a Pauli Graph
+    and applying various optimisation techniques such as:
+    - phase folding
+    - a synthesis algorithm from Pauli Graph to Clifford + Rz aimed at reducing the number of 2
+    qubit gates
+
+    Rotation angles must be numeric as symbolic angles are not supported currently.
+
+    Parameters:
+    - window_size: Sets the size of the sliding window used for lookahead during synthesis. Must be positive.
+    - pool_size: Sets the number of candidate gates to maintain in the pool. Must be positive.
+    - top_up_size: Sets the number of candidate gates to add after each TQE gate. Must be positive.
+    - seed: Sets the random seed used to sample candidate gates. Must be non-negative.
+    - parallel_mode: A :class:`ParallelMode` for candidate synthesis.
+      Defaults to :attr:`ParallelMode.Auto`.
+    """
+
+    window_size: int | None = None
+    pool_size: int | None = None
+    top_up_size: int | None = None
+    seed: int | None = None
+    parallel_mode: ParallelMode = ParallelMode.Auto
+    _scope: PassScope = GlobalScope.PRESERVE_PUBLIC
+
+    def __post_init__(self) -> None:
+        self._validate_parameters()
+
+    def _validate_parameters(self) -> None:
+        for parameter in ("window_size", "pool_size", "top_up_size"):
+            value = getattr(self, parameter)
+            if value is not None and value <= 0:
+                raise ValueError(f"{parameter} must be positive")
+        if self.seed is not None and self.seed < 0:
+            raise ValueError("seed must be non-negative")
+        if not isinstance(self.parallel_mode, ParallelMode):
+            raise TypeError(
+                "parallel_mode must be an instance of the ParallelMode enum"
+            )
+
+    def with_scope(self, scope: PassScope) -> PauliGraphResynthesis:
+        """Set the scope of this pass and return self."""
+        self._scope = scope
+        return self
+
+    def run(self, hugr: Hugr, *, inplace: bool = True) -> PassResult:
+        return implement_pass_run(
+            self,
+            hugr=hugr,
+            inplace=inplace,
+            copy_call=lambda h: self._pauli_graph_resynthesis(h, inplace),
+        )
+
+    def _pauli_graph_resynthesis(self, hugr: Hugr, inplace: bool) -> PassResult:
+        self._validate_parameters()
+        program = _state.CompilationState.from_python(hugr)
+        _passes.pauli_graph_resynthesis(
+            program._inner,
+            scope=self._scope,
+            window_size=self.window_size,
+            pool_size=self.pool_size,
+            top_up_size=self.top_up_size,
+            seed=self.seed,
+            parallel_mode=self.parallel_mode,
+        )
+        package = program.to_python()
+        return PassResult.for_pass(
+            self, hugr=package.modules[0], inplace=inplace, result=None
+        )
+
+
 @dataclass(kw_only=True)
 class _QSystemLLVMPass(ComposablePass):
     """Prepare a QSystem program for LLVM lowering.