diff --git a/test_files/guppy_examples/flat_quantum.hugr b/test_files/guppy_examples/flat_quantum.hugr new file mode 100644 index 000000000..1241e3149 Binary files /dev/null and b/test_files/guppy_examples/flat_quantum.hugr differ diff --git a/test_files/guppy_examples/flat_quantum.py b/test_files/guppy_examples/flat_quantum.py new file mode 100644 index 000000000..cd862b087 --- /dev/null +++ b/test_files/guppy_examples/flat_quantum.py @@ -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()) diff --git a/tket-py/src/passes.rs b/tket-py/src/passes.rs index c69de7947..65f601a6b 100644 --- a/tket-py/src/passes.rs +++ b/tket-py/src/passes.rs @@ -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> { 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::())?; 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) -> 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, +) -> 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(()) +} diff --git a/tket-py/test/test_pass.py b/tket-py/test/test_pass.py index 4c8957955..a3c4d5431 100644 --- a/tket-py/test/test_pass.py +++ b/tket-py/test/test_pass.py @@ -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 diff --git a/tket-py/tket/_state/__init__.py b/tket-py/tket/_state/__init__.py index 6212cbadd..d4b043d6c 100644 --- a/tket-py/tket/_state/__init__.py +++ b/tket-py/tket/_state/__init__.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 diff --git a/tket-py/tket/_tket/passes.pyi b/tket-py/tket/_tket/passes.pyi index 699a3beea..6fdbdbfb5 100644 --- a/tket-py/tket/_tket/passes.pyi +++ b/tket-py/tket/_tket/passes.pyi @@ -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. + """ diff --git a/tket-py/tket/passes/__init__.py b/tket-py/tket/passes/__init__.py index e21f78e71..259e2c5e7 100644 --- a/tket-py/tket/passes/__init__.py +++ b/tket-py/tket/passes/__init__.py @@ -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 diff --git a/tket-qsystem/src/lib.rs b/tket-qsystem/src/lib.rs index 7ba4c4f4f..faf83971b 100644 --- a/tket-qsystem/src/lib.rs +++ b/tket-qsystem/src/lib.rs @@ -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();