Add InteropQubitManager

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 (

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.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)

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 add 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])