feat: expose the QSystemPass to Python (#1556)

closes #1553 

Driveby: now ncludes a fix to how we resolve extensions in Python. See
[076c176](076c176).
Prior to this fix the transformed HUGR contained Custom ops rather than
QSystem ops as we would expect.

Another driveby: There was a Boolean flag called `hide_funcs` which did
not have a corresponding `with_*` method in Rust. I assume this was an
oversight so I added a method in
[739870d](https://github.com/Quantinuum/tket2/pull/1556/commits/739870dc8167dcb5ca911346a044d51c0e354b5a)[739870d](https://github.com/Quantinuum/tket2/pull/1556/commits/739870dc8167dcb5ca911346a044d51c0e354b5a).

---------

Co-authored-by: Agustín Borgna <agustin.borgna@quantinuum.com>

Author: CalMacCQ

test_files/guppy_examples/flat_quantum.hugr

@@ -0,0 +1,23 @@
+# /// script
+# requires-python = ">=3.13"
+# dependencies = [
+#     "guppylang ==0.21.13",
+# ]
+# ///
+from pathlib import Path
+from sys import argv
+
+from guppylang import guppy
+from guppylang.std.quantum import qubit, h, cx, rz
+from guppylang.std.angles import pi
+
+
+@guppy
+def flat_quantum_func(q0: qubit, q1: qubit) -> None:
+    h(q0)
+    rz(q0, 3 * pi / 4)
+    cx(q0, q1)
+
+
+program = flat_quantum_func.compile_function()
+Path(argv[0]).with_suffix(".hugr").write_bytes(program.to_bytes())

test_files/guppy_examples/flat_quantum.py

@@ -11,10 +11,11 @@ use std::{cmp::min, convert::TryInto, fs, num::NonZeroUsize, path::PathBuf};
 
 use pyo3::prelude::*;
 use tket::optimiser::badger::BadgerOptions;
-use tket::passes;
 use tket::passes::composable::{ComposablePass, WithScope};
 use tket::{Circuit, TketOp, op_matches};
 
+use tket::passes;
+
 use crate::optimiser::PyBadgerOptimiser;
 use crate::state::CompilationState;
 use crate::utils::{ConvertPyErr, create_py_exception};
@@ -32,6 +33,7 @@ pub fn module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
     m.add_function(wrap_pyfunction!(self::chunks::chunks, &m)?)?;
     m.add_function(wrap_pyfunction!(self::tket1::tket1_pass, &m)?)?;
     m.add_function(wrap_pyfunction!(resolve_modifiers, &m)?)?;
+    m.add_function(wrap_pyfunction!(qsystem_rebase_pass, &m)?)?;
     m.add("PullForwardError", py.get_type::<PyPullForwardError>())?;
     m.add(
         "InlineFunctionsError",
@@ -64,6 +66,13 @@ create_py_exception!(
     PyInlineFunctionsError,
     "Errors from the function inlining pass."
 );
+
+create_py_exception!(
+    tket_qsystem::QSystemPassError,
+    PyQSystemPassError,
+    "Errors from the QSystem rebase pass."
+);
+
 /// Flatten the structure of a Guppy-generated program to enable additional optimisations.
 ///
 /// This should normally be called first before other optimisations.
@@ -203,3 +212,26 @@ fn resolve_modifiers(circ: &mut CompilationState, scope: Option<PyPassScope>) ->
     pass.run(&mut circ.hugr).convert_pyerrs()?;
     Ok(())
 }
+
+#[pyfunction]
+#[pyo3(signature=(circ, *, constant_fold = true, monomorphize = true, force_order = true, lazify = true, hide_funcs = true, scope = None))]
+fn qsystem_rebase_pass(
+    circ: &mut CompilationState,
+    constant_fold: bool,
+    monomorphize: bool,
+    force_order: bool,
+    lazify: bool,
+    hide_funcs: bool,
+    scope: Option<PyPassScope>,
+) -> PyResult<()> {
+    let py_scope = scope.unwrap_or_default();
+    let qsystem_pass = tket_qsystem::QSystemPass::default_with_scope(py_scope.scope)
+        .with_constant_fold(constant_fold)
+        .with_monomorphize(monomorphize)
+        .with_force_order(force_order)
+        .with_lazify(lazify)
+        .with_hide_funcs(hide_funcs);
+
+    qsystem_pass.run(&mut circ.hugr).convert_pyerrs()?;
+    Ok(())
+}

tket-py/src/passes.rs

@@ -21,7 +21,8 @@
 from hypothesis.strategies._internal import SearchStrategy
 from hypothesis import given, settings
 
-from tket.passes import PytketHugrPass
+
+from tket.passes import PytketHugrPass, QSystemPass
 from pytket.passes import (
     CliffordSimp,
     SquashRzPhasedX,
@@ -323,3 +324,13 @@ def test_issue_1516() -> None:
     opt_hugr = opt(hugr, inplace=False)
 
     CompilationState.from_python(opt_hugr).validate()
+
+
+def test_python_qsystem_pass() -> None:
+    normalize = NormalizeGuppy()
+    hugr = normalize(_hugr_from_path("test_files/guppy_examples/flat_quantum.hugr"))
+    qsystem_pass = QSystemPass()
+    qsystem_hugr = qsystem_pass(hugr)
+    assert _count_ops(qsystem_hugr, "ZZPhase") == 1
+    assert _count_ops(qsystem_hugr, "Custom") == 0
+    assert _count_ops(qsystem_hugr, "tket.quantum") == 0

tket-py/test/test_pass.py

@@ -6,6 +6,7 @@
 
 from hugr.envelope import EnvelopeConfig
 from hugr.ext import ExtensionRegistry
+from tket_exts import tket_registry
 from .._tket import state as _state
 from .build import CircBuild, Command
 
@@ -97,11 +98,10 @@ def to_python(self) -> Package:
         """Convert this CompilationState back to a python Hugr package."""
         # Convert the inner hugr to bytes and load it in Python.
         hugr_bytes = self._inner.to_bytes()
-        package = Package.from_bytes(hugr_bytes, self._py_extensions)
+        package = Package.from_bytes(hugr_bytes, tket_registry())
         if self._py_extensions is not None:
             # Resolve the extensions in the loaded package using the python registry, if needed.
-            # TODO: Use the `package.resolve_extensions` for clarity once it's been released in `hugr-py 0.16.0`.
-            package.used_extensions(self._py_extensions)
+            package.resolve_extensions(self._py_extensions)
         return package
 
     @staticmethod

tket-py/tket/_state/__init__.py

@@ -99,3 +99,21 @@ def resolve_modifiers(
     :param circ: The input program as a CompilationState.
     :param scope: A scope to control how the pass is applied to HUGR regions.
     """
+
+def qsystem_rebase_pass(
+    circ: CompilationState,
+    constant_fold: bool = True,
+    monomorphize: bool = True,
+    force_order: bool = True,
+    lazify: bool = True,
+    hide_funcs: bool = True,
+    scope: PassScope | None = None,
+) -> None:
+    """Runs a rust backed pass to convert quantum ops to qsystem ops.
+
+    :param constant_fold: Whether to perform constant folding.
+    :param monomorphize: Whether to monomorphize generic functions.
+    :param force_order: Whether to enforce total ordering of all HUGR operations.
+    :param lazify: Whether to replace measurements with lazy measurements.
+    :param hide_funcs: Make all HUGR functions private.
+    """

tket-py/tket/_tket/passes.pyi

@@ -29,6 +29,7 @@
     "InlineFunctions",
     "NormalizeGuppy",
     "ModifierResolverPass",
+    "QSystemPass",
 ]
 
 
@@ -294,3 +295,59 @@ def _run_tk(self, program: _state.CompilationState) -> _state.CompilationState:
             scope=self._scope,
         )
         return program
+
+
+@dataclass(kw_only=True)
+class QSystemPass(ComposablePass):
+    """A pass to convert quantum ops to qsystem ops.
+
+     Parameters:
+    - constant_fold: Whether to perform constant folding.
+    - monomorphize: Whether to monomorphize generic functions.
+    - force_order: Whether to enforce total ordering of all HUGR operations.
+    - lazify: Whether to replace measurements with lazy measurements.
+    - hide_funcs: Whether to mark all functions as private.
+    """
+
+    constant_fold: bool = True
+    monomorphize: bool = True
+    force_order: bool = True
+    lazify: bool = True
+    hide_funcs: bool = True
+    _scope: PassScope = GlobalScope.PRESERVE_PUBLIC
+
+    def run(self, hugr: Hugr, *, inplace: bool = True) -> PassResult:
+        return implement_pass_run(
+            self,
+            hugr=hugr,
+            inplace=inplace,
+            copy_call=lambda h: self._qsystem_rebase(h, inplace),
+        )
+
+    def with_scope(self, scope: PassScope) -> QSystemPass:
+        """Set the scope of this pass and return self."""
+        self._scope = scope
+        return self
+
+    def _qsystem_rebase(self, hugr: Hugr, inplace: bool) -> PassResult:
+        tk_program = _state.CompilationState.from_python(hugr)
+
+        self._run_tk(tk_program)
+
+        package = tk_program.to_python()
+        return PassResult.for_pass(
+            self, hugr=package.modules[0], inplace=inplace, result=None
+        )
+
+    def _run_tk(self, program: _state.CompilationState) -> _state.CompilationState:
+        """Run the pass in the CompilationState"""
+        _passes.qsystem_rebase_pass(
+            program._inner,
+            constant_fold=self.constant_fold,
+            monomorphize=self.monomorphize,
+            force_order=self.force_order,
+            lazify=self.lazify,
+            hide_funcs=self.hide_funcs,
+            scope=self._scope,
+        )
+        return program

tket-py/tket/passes/__init__.py

@@ -140,6 +140,16 @@ impl QSystemPass {
         self
     }
 
+    /// Makes all functions private.
+    ///
+    /// On by default
+    ///
+    /// When enabled all functions are marked as private. This enables LLVM to drop functions which are not called.
+    pub fn with_hide_funcs(mut self, hide_funcs: bool) -> Self {
+        self.hide_funcs = hide_funcs;
+        self
+    }
+
     /// Add order edges in the HUGR regions to force qubit frees to be as early
     /// as possible, quantum ops to be as early as possible, and Future::Reads
     /// to be as late as possible.
@@ -444,9 +454,9 @@ mod test {
         // Run again without hiding...
         let mut hugr_public = orig;
         QSystemPass {
-            hide_funcs: false,
             ..Default::default()
         }
+        .with_hide_funcs(false)
         .run(&mut hugr_public)
         .unwrap();