From 00ad7f20589dc51feb95e265fb706c98c1f84702 Mon Sep 17 00:00:00 2001 From: Tanuj Khattar Date: Mon, 9 Oct 2023 13:07:12 -0700 Subject: [PATCH] Add `InteropQubitManager` to use in Cirq/Bloqs interop code. (#385) * Add InteropQubitManager * Fix typo and use cirq.ops.SimpleQubitManager since dependence is not up to date * Add more tests --- qualtran/cirq_interop/_cirq_to_bloq.py | 16 ++++++--- qualtran/cirq_interop/_cirq_to_bloq_test.py | 19 ++++++++++ .../cirq_interop/_interop_qubit_manager.py | 36 +++++++++++++++++++ .../_interop_qubit_manager_test.py | 30 ++++++++++++++++ 4 files changed, 97 insertions(+), 4 deletions(-) create mode 100644 qualtran/cirq_interop/_interop_qubit_manager.py create mode 100644 qualtran/cirq_interop/_interop_qubit_manager_test.py diff --git a/qualtran/cirq_interop/_cirq_to_bloq.py b/qualtran/cirq_interop/_cirq_to_bloq.py index 956e0088e..79c6cd417 100644 --- a/qualtran/cirq_interop/_cirq_to_bloq.py +++ b/qualtran/cirq_interop/_cirq_to_bloq.py @@ -26,6 +26,7 @@ from numpy.typing import NDArray from qualtran import Bloq, BloqBuilder, CompositeBloq, Register, Side, Signature, Soquet, SoquetT +from qualtran.cirq_interop._interop_qubit_manager import InteropQubitManager if TYPE_CHECKING: from qualtran.drawing import WireSymbol @@ -34,6 +35,12 @@ CirqQuregInT = Union[NDArray[cirq.Qid], Sequence[cirq.Qid]] +def get_cirq_quregs(signature: Signature, qm: InteropQubitManager): + ret = signature.get_cirq_quregs() + qm.manage_qubits(itertools.chain.from_iterable(qreg.flatten() for qreg in ret.values())) + return ret + + @frozen class CirqGateAsBloq(Bloq): """A Bloq wrapper around a `cirq.Gate`, preserving signature if gate is a `GateWithRegisters`.""" @@ -66,8 +73,8 @@ def cirq_registers(self) -> cirq_ft.Signature: ) def decompose_bloq(self) -> 'CompositeBloq': - in_quregs = self.signature.get_cirq_quregs() - qubit_manager = cirq.ops.SimpleQubitManager() + qubit_manager = InteropQubitManager() + in_quregs = get_cirq_quregs(self.signature, qubit_manager) cirq_op, out_quregs = self.as_cirq_op(qubit_manager, **in_quregs) context = cirq.DecompositionContext(qubit_manager=qubit_manager) decomposed_optree = cirq.decompose_once(cirq_op, context=context, default=None) @@ -358,8 +365,9 @@ def decompose_from_cirq_op(bloq: 'Bloq') -> 'CompositeBloq': ): raise NotImplementedError(f"{bloq} does not support decomposition.") - in_quregs = bloq.signature.get_cirq_quregs() - cirq_op, out_quregs = bloq.as_cirq_op(cirq.ops.SimpleQubitManager(), **in_quregs) + qubit_manager = InteropQubitManager() + in_quregs = get_cirq_quregs(bloq.signature, qubit_manager) + cirq_op, out_quregs = bloq.as_cirq_op(qubit_manager, **in_quregs) from qualtran.cirq_interop._bloq_to_cirq import BloqAsCirqGate if cirq_op is None or ( diff --git a/qualtran/cirq_interop/_cirq_to_bloq_test.py b/qualtran/cirq_interop/_cirq_to_bloq_test.py index 52c1eae58..405d9efa3 100644 --- a/qualtran/cirq_interop/_cirq_to_bloq_test.py +++ b/qualtran/cirq_interop/_cirq_to_bloq_test.py @@ -201,3 +201,22 @@ def signature(self) -> cirq_ft.Signature: assert bloqs_list.count(Join(3)) == 6 assert bloqs_list.count(Allocate(2)) == 2 assert bloqs_list.count(Free(2)) == 2 + + +def test_cirq_gate_as_bloq_for_left_only_gates(): + class LeftOnlyGate(cirq_ft.GateWithRegisters): + @property + def signature(self): + return cirq_ft.Signature([cirq_ft.Register('junk', 2, side=cirq_ft.infra.Side.LEFT)]) + + def decompose_from_registers(self, *, context, junk) -> cirq.OP_TREE: + yield cirq.CNOT(*junk) + yield cirq.reset_each(*junk) + + # Using InteropQubitManager enables support for LeftOnlyGate's in CirqGateAsBloq. + cbloq = CirqGateAsBloq(gate=LeftOnlyGate()).decompose_bloq() + bloqs_list = [binst.bloq for binst in cbloq.bloq_instances] + assert bloqs_list.count(Split(2)) == 1 + assert bloqs_list.count(Free(1)) == 2 + assert bloqs_list.count(CirqGateAsBloq(cirq.CNOT)) == 1 + assert bloqs_list.count(CirqGateAsBloq(cirq.ResetChannel())) == 2 diff --git a/qualtran/cirq_interop/_interop_qubit_manager.py b/qualtran/cirq_interop/_interop_qubit_manager.py new file mode 100644 index 000000000..78def1bda --- /dev/null +++ b/qualtran/cirq_interop/_interop_qubit_manager.py @@ -0,0 +1,36 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qubit Manager to use when converting Cirq gates to/from Bloqs.""" +from typing import Iterable + +import cirq + + +class InteropQubitManager(cirq.ops.SimpleQubitManager): + """Qubit Manager to use to facilitate interop of Cirq gates and Bloqs.""" + + def __init__(self): + super().__init__() + self._managed_qubits = set() + + def manage_qubits(self, qubits: Iterable[cirq.Qid]): + self._managed_qubits |= set(qubits) + + def qfree(self, qubits: Iterable[cirq.Qid]): + qs = set(qubits) + managed_qs = qs & self._managed_qubits + qs -= managed_qs + self._managed_qubits -= managed_qs + super().qfree(qs) diff --git a/qualtran/cirq_interop/_interop_qubit_manager_test.py b/qualtran/cirq_interop/_interop_qubit_manager_test.py new file mode 100644 index 000000000..03c939f52 --- /dev/null +++ b/qualtran/cirq_interop/_interop_qubit_manager_test.py @@ -0,0 +1,30 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import cirq +import pytest + +from qualtran.cirq_interop._interop_qubit_manager import InteropQubitManager + + +def test_interop_qubit_manager(): + qm = InteropQubitManager() + q = cirq.q('junk') + with pytest.raises(ValueError, match='not allocated'): + qm.qfree([q]) + # You can delegate qubits to be "managed" by the InteropQubitManager. + qm.manage_qubits([q]) + qm.qfree([q]) + # q was already deallocated. + with pytest.raises(ValueError, match='not allocated'): + qm.qfree([q])