Move GreedyQubitManager from Cirq-FT to Cirq-Core

cirq-core/cirq/__init__.py

@@ -224,6 +224,7 @@
     givens,
     GlobalPhaseGate,
     global_phase_operation,
+    GreedyQubitManager,
     H,
     HPowGate,
     I,
@@ -301,6 +302,7 @@
     ry,
     rz,
     S,
+    SimpleQubitManager,
     SingleQubitCliffordGate,
     SingleQubitPauliStringGateOperation,
     SQRT_ISWAP,

cirq-core/cirq/ops/__init__.py

@@ -122,6 +122,8 @@
 
 from cirq.ops.qubit_manager import BorrowableQubit, CleanQubit, QubitManager, SimpleQubitManager
 
+from cirq.ops.greedy_qubit_manager import GreedyQubitManager
+
 from cirq.ops.qubit_order import QubitOrder
 
 from cirq.ops.qubit_order_or_list import QubitOrderOrList

cirq-core/cirq/ops/greedy_qubit_manager.py

@@ -0,0 +1,86 @@
+# Copyright 2023 The Cirq Developers
+#
+# 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.
+
+from typing import Iterable, List, Set, TYPE_CHECKING
+
+from cirq.ops import named_qubit, qid_util, qubit_manager
+
+if TYPE_CHECKING:
+    import cirq
+
+
+class GreedyQubitManager(qubit_manager.QubitManager):
+    """Greedy allocator that maximizes/minimizes qubit reuse based on a configurable parameter.
+
+    GreedyQubitManager can be configured, using `maximize_reuse` flag, to work in one of two modes:
+    - Minimize qubit reuse (maximize_reuse=False): For a fixed width, this mode uses a FIFO (First
+             in First out) strategy s.t. next allocated qubit is one which was freed the earliest.
+    - Maximize qubit reuse (maximize_reuse=True): For a fixed width, this mode uses a LIFO (Last in
+            First out) strategy s.t. the next allocated qubit is one which was freed the latest.
+
+    If the requested qubits are more than the set of free qubits, the qubit manager automatically
+    resizes the size of the managed qubit pool and adds new free qubits, that have their last
+    freed time to be -infinity.
+
+    For borrowing qubits, the qubit manager simply delegates borrow requests to `self.qalloc`, thus
+    always allocating new clean qubits.
+    """
+
+    def __init__(self, prefix: str, *, size: int = 0, maximize_reuse: bool = False):
+        """Initializes `GreedyQubitManager`
+
+        Args:
+            prefix: The prefix to use for naming new clean ancillas allocated by the qubit manager.
+                    The i'th allocated qubit is of the type `cirq.NamedQubit(f'{prefix}_{i}')`.
+            size: The initial size of the pool of ancilla qubits managed by the qubit manager. The
+                    qubit manager can automatically resize itself when the allocation request
+                    exceeds the number of available qubits.
+            maximize_reuse: Flag to control a FIFO vs LIFO strategy, defaults to False (FIFO).
+        """
+        self._prefix = prefix
+        self._used_qubits: Set['cirq.Qid'] = set()
+        self._free_qubits: List['cirq.Qid'] = []
+        self._size = 0
+        self.maximize_reuse = maximize_reuse
+        self.resize(size)
+
+    def _allocate_qid(self, name: str, dim: int) -> 'cirq.Qid':
+        return qid_util.q(name) if dim == 2 else named_qubit.NamedQid(name, dimension=dim)
+
+    def resize(self, new_size: int, dim: int = 2) -> None:
+        if new_size <= self._size:
+            return
+        new_qubits: List['cirq.Qid'] = [
+            self._allocate_qid(f'{self._prefix}_{s}', dim) for s in range(self._size, new_size)
+        ]
+        self._free_qubits = new_qubits + self._free_qubits
+        self._size = new_size
+
+    def qalloc(self, n: int, dim: int = 2) -> List['cirq.Qid']:
+        if not n:
+            return []
+        self.resize(self._size + n - len(self._free_qubits), dim=dim)
+        ret_qubits = self._free_qubits[-n:] if self.maximize_reuse else self._free_qubits[:n]
+        self._free_qubits = self._free_qubits[:-n] if self.maximize_reuse else self._free_qubits[n:]
+        self._used_qubits.update(ret_qubits)
+        return ret_qubits
+
+    def qfree(self, qubits: Iterable['cirq.Qid']) -> None:
+        qs = list(dict(zip(qubits, qubits)).keys())
+        assert self._used_qubits.issuperset(qs), "Only managed qubits currently in-use can be freed"
+        self._used_qubits = self._used_qubits.difference(qs)
+        self._free_qubits.extend(qs)
+
+    def qborrow(self, n: int, dim: int = 2) -> List['cirq.Qid']:
+        return self.qalloc(n, dim)

cirq-core/cirq/ops/greedy_qubit_manager_test.py

@@ -0,0 +1,101 @@
+# Copyright 2023 The Cirq Developers
+#
+# 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
+
+
+class GateAllocInDecompose(cirq.Gate):
+    def __init__(self, num_alloc: int = 1):
+        self.num_alloc = num_alloc
+
+    def _num_qubits_(self) -> int:
+        return 1
+
+    def _decompose_with_context_(self, qubits, context):
+        assert context is not None
+        qm = context.qubit_manager
+        for q in qm.qalloc(self.num_alloc):
+            yield cirq.CNOT(qubits[0], q)
+            qm.qfree([q])
+
+    def __str__(self):
+        return 'TestGateAlloc'
+
+
+def test_greedy_qubit_manager():
+    def make_circuit(qm: cirq.QubitManager):
+        q = cirq.LineQubit.range(2)
+        g = GateAllocInDecompose(1)
+        context = cirq.DecompositionContext(qubit_manager=qm)
+        circuit = cirq.Circuit(
+            cirq.decompose_once(g.on(q[0]), context=context),
+            cirq.decompose_once(g.on(q[1]), context=context),
+        )
+        return circuit
+
+    qm = cirq.GreedyQubitManager(prefix="ancilla", size=1)
+    # Qubit manager with only 1 managed qubit. Will always repeat the same qubit.
+    circuit = make_circuit(qm)
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+0: ───────────@───────
+              │
+1: ───────────┼───@───
+              │   │
+ancilla_0: ───X───X───
+        """,
+    )
+
+    qm = cirq.GreedyQubitManager(prefix="ancilla", size=2)
+    # Qubit manager with 2 managed qubits and maximize_reuse=False, tries to minimize adding
+    # additional data dependencies by minimizing qubit reuse.
+    circuit = make_circuit(qm)
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+              ┌──┐
+0: ────────────@─────
+               │
+1: ────────────┼@────
+               ││
+ancilla_0: ────X┼────
+                │
+ancilla_1: ─────X────
+              └──┘
+        """,
+    )
+
+    qm = cirq.GreedyQubitManager(prefix="ancilla", size=2, maximize_reuse=True)
+    # Qubit manager with 2 managed qubits and maximize_reuse=True, tries to maximize reuse by
+    # potentially adding new data dependencies.
+    circuit = make_circuit(qm)
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+0: ───────────@───────
+              │
+1: ───────────┼───@───
+              │   │
+ancilla_1: ───X───X───
+     """,
+    )
+
+
+def test_greedy_qubit_manager_preserves_order():
+    qm = cirq.GreedyQubitManager(prefix="anc")
+    ancillae = [cirq.q(f"anc_{i}") for i in range(100)]
+    assert qm.qalloc(100) == ancillae
+    qm.qfree(ancillae)
+    assert qm.qalloc(100) == ancillae

cirq-core/cirq/protocols/json_test_data/spec.py

@@ -89,6 +89,9 @@
         'LineInitialMapper',
         'MappingManager',
         'RouteCQC',
+        # Qubit Managers,
+        'SimpleQubitManager',
+        'GreedyQubitManager',
         # global objects
         'CONTROL_TAG',
         'PAULI_BASIS',

cirq-ft/cirq_ft/__init__.py

@@ -45,7 +45,6 @@
 )
 from cirq_ft.infra import (
     GateWithRegisters,
-    GreedyQubitManager,
     Register,
     Signature,
     SelectionRegister,

cirq-ft/cirq_ft/algos/arithmetic_gates_test.py

@@ -18,7 +18,7 @@
 import cirq_ft
 import numpy as np
 import pytest
-from cirq_ft.infra import bit_tools, GreedyQubitManager
+from cirq_ft.infra import bit_tools
 
 
 def identity_map(n: int):
@@ -174,7 +174,7 @@ def test_add(a: int, b: int, num_bits: int):
     num_anc = num_bits - 1
     gate = cirq_ft.AdditionGate(num_bits)
     qubits = cirq.LineQubit.range(2 * num_bits)
-    greedy_mm = cirq_ft.GreedyQubitManager(prefix="_a", maximize_reuse=True)
+    greedy_mm = cirq.GreedyQubitManager(prefix="_a", maximize_reuse=True)
     context = cirq.DecompositionContext(greedy_mm)
     circuit = cirq.Circuit(cirq.decompose_once(gate.on(*qubits), context=context))
     ancillas = sorted(circuit.all_qubits())[-num_anc:]
@@ -258,7 +258,7 @@ def test_add_truncated():
     num_anc = num_bits - 1
     gate = cirq_ft.AdditionGate(num_bits)
     qubits = cirq.LineQubit.range(2 * num_bits)
-    greedy_mm = cirq_ft.GreedyQubitManager(prefix="_a", maximize_reuse=True)
+    greedy_mm = cirq.GreedyQubitManager(prefix="_a", maximize_reuse=True)
     context = cirq.DecompositionContext(greedy_mm)
     circuit = cirq.Circuit(cirq.decompose_once(gate.on(*qubits), context=context))
     ancillas = sorted(circuit.all_qubits() - frozenset(qubits))
@@ -272,7 +272,7 @@ def test_add_truncated():
     num_anc = num_bits - 1
     gate = cirq_ft.AdditionGate(num_bits)
     qubits = cirq.LineQubit.range(2 * num_bits)
-    greedy_mm = cirq_ft.GreedyQubitManager(prefix="_a", maximize_reuse=True)
+    greedy_mm = cirq.GreedyQubitManager(prefix="_a", maximize_reuse=True)
     context = cirq.DecompositionContext(greedy_mm)
     circuit = cirq.Circuit(cirq.decompose_once(gate.on(*qubits), context=context))
     ancillas = sorted(circuit.all_qubits() - frozenset(qubits))
@@ -290,7 +290,7 @@ def test_subtract(a, b, num_bits):
     num_anc = num_bits - 1
     gate = cirq_ft.AdditionGate(num_bits)
     qubits = cirq.LineQubit.range(2 * num_bits)
-    greedy_mm = cirq_ft.GreedyQubitManager(prefix="_a", maximize_reuse=True)
+    greedy_mm = cirq.GreedyQubitManager(prefix="_a", maximize_reuse=True)
     context = cirq.DecompositionContext(greedy_mm)
     circuit = cirq.Circuit(cirq.decompose_once(gate.on(*qubits), context=context))
     ancillas = sorted(circuit.all_qubits())[-num_anc:]
@@ -340,7 +340,7 @@ def test_decompose_less_than_equal_gate(P: int, n: int, Q: int, m: int):
     circuit = cirq.Circuit(
         cirq.decompose_once(
             cirq_ft.LessThanEqualGate(n, m).on(*cirq.LineQubit.range(n + m + 1)),
-            context=cirq.DecompositionContext(GreedyQubitManager(prefix='_c')),
+            context=cirq.DecompositionContext(cirq.GreedyQubitManager(prefix='_c')),
         )
     )
     qubit_order = tuple(sorted(circuit.all_qubits()))

cirq-ft/cirq_ft/algos/programmable_rotation_gate_array_test.py

@@ -63,7 +63,7 @@ def construct_prga_with_identity(*args, **kwargs) -> cirq_ft.ProgrammableRotatio
 def test_programmable_rotation_gate_array(angles, kappa, constructor):
     rotation_gate = cirq.X
     programmable_rotation_gate = constructor(*angles, kappa=kappa, rotation_gate=rotation_gate)
-    greedy_mm = cirq_ft.GreedyQubitManager(prefix="_a")
+    greedy_mm = cirq.GreedyQubitManager(prefix="_a")
     g = cirq_ft.testing.GateHelper(
         programmable_rotation_gate, context=cirq.DecompositionContext(greedy_mm)
     )

cirq-ft/cirq_ft/algos/qrom_test.py

@@ -29,7 +29,7 @@
 @pytest.mark.parametrize("num_controls", [0, 1, 2])
 def test_qrom_1d(data, num_controls):
     qrom = cirq_ft.QROM.build(*data, num_controls=num_controls)
-    greedy_mm = cirq_ft.GreedyQubitManager('a', maximize_reuse=True)
+    greedy_mm = cirq.GreedyQubitManager('a', maximize_reuse=True)
     g = cirq_ft.testing.GateHelper(qrom, context=cirq.DecompositionContext(greedy_mm))
     decomposed_circuit = cirq.Circuit(cirq.decompose(g.operation, context=g.context))
     inverse = cirq.Circuit(cirq.decompose(g.operation**-1, context=g.context))
@@ -121,7 +121,7 @@ def test_t_complexity(data):
 def _assert_qrom_has_diagram(qrom: cirq_ft.QROM, expected_diagram: str):
     gh = cirq_ft.testing.GateHelper(qrom)
     op = gh.operation
-    context = cirq.DecompositionContext(qubit_manager=cirq_ft.GreedyQubitManager(prefix="anc"))
+    context = cirq.DecompositionContext(qubit_manager=cirq.GreedyQubitManager(prefix="anc"))
     circuit = cirq.Circuit(cirq.decompose_once(op, context=context))
     selection = [
         *itertools.chain.from_iterable(gh.quregs[reg.name] for reg in qrom.selection_registers)
@@ -208,7 +208,7 @@ def test_qrom_multi_dim(data, num_controls):
         target_bitsizes=target_bitsizes,
         num_controls=num_controls,
     )
-    greedy_mm = cirq_ft.GreedyQubitManager('a', maximize_reuse=True)
+    greedy_mm = cirq.GreedyQubitManager('a', maximize_reuse=True)
     g = cirq_ft.testing.GateHelper(qrom, context=cirq.DecompositionContext(greedy_mm))
     decomposed_circuit = cirq.Circuit(cirq.decompose(g.operation, context=g.context))
     inverse = cirq.Circuit(cirq.decompose(g.operation**-1, context=g.context))

cirq-ft/cirq_ft/algos/qubitization_walk_operator_test.py

@@ -63,7 +63,7 @@ def test_qubitization_walk_operator(num_sites: int, eps: float):
     L_state = np.zeros(2 ** len(g.quregs['selection']))
     L_state[: len(ham_coeff)] = np.sqrt(ham_coeff / qubitization_lambda)
 
-    greedy_mm = cirq_ft.GreedyQubitManager('ancilla', maximize_reuse=True)
+    greedy_mm = cirq.GreedyQubitManager('ancilla', maximize_reuse=True)
     walk_circuit = cirq_ft.map_clean_and_borrowable_qubits(walk_circuit, qm=greedy_mm)
     assert len(walk_circuit.all_qubits()) < 23
     qubit_order = cirq.QubitOrder.explicit(

cirq-ft/cirq_ft/algos/reflection_using_prepare_test.py

@@ -44,7 +44,7 @@ def keep(op: cirq.Operation):
 
 
 def greedily_allocate_ancilla(circuit: cirq.AbstractCircuit) -> cirq.Circuit:
-    greedy_mm = cirq_ft.GreedyQubitManager(prefix="ancilla", maximize_reuse=True)
+    greedy_mm = cirq.GreedyQubitManager(prefix="ancilla", maximize_reuse=True)
     circuit = cirq_ft.map_clean_and_borrowable_qubits(circuit, qm=greedy_mm)
     assert len(circuit.all_qubits()) < 30
     return circuit

cirq-ft/cirq_ft/algos/select_swap_qrom_test.py

@@ -29,7 +29,7 @@ def test_select_swap_qrom(data, block_size):
     selection_q, selection_r = selection[: qrom.selection_q], selection[qrom.selection_q :]
     targets = [qubit_regs[f"target{i}"] for i in range(len(data))]
 
-    greedy_mm = cirq_ft.GreedyQubitManager(prefix="_a", maximize_reuse=True)
+    greedy_mm = cirq.GreedyQubitManager(prefix="_a", maximize_reuse=True)
     context = cirq.DecompositionContext(greedy_mm)
     qrom_circuit = cirq.Circuit(cirq.decompose(qrom.on_registers(**qubit_regs), context=context))
 

cirq-ft/cirq_ft/algos/selected_majorana_fermion_test.py

@@ -100,7 +100,7 @@ def test_selected_majorana_fermion_gate_decomposed_diagram():
         cirq_ft.SelectionRegister('selection', selection_bitsize, target_bitsize),
         target_gate=cirq.X,
     )
-    greedy_mm = cirq_ft.GreedyQubitManager(prefix="_a", maximize_reuse=True)
+    greedy_mm = cirq.GreedyQubitManager(prefix="_a", maximize_reuse=True)
     g = cirq_ft.testing.GateHelper(gate)
     context = cirq.DecompositionContext(greedy_mm)
     circuit = cirq.Circuit(cirq.decompose_once(g.operation, context=context))

cirq-ft/cirq_ft/algos/unary_iteration_gate.py

@@ -182,13 +182,13 @@ def unary_iteration(
     Users can write multi-dimensional coherent for loops as follows:
 
     >>> import cirq
-    >>> from cirq_ft import unary_iteration, GreedyQubitManager
+    >>> from cirq_ft import unary_iteration
     >>> N, M = 5, 7
     >>> target = [[cirq.q(f't({i}, {j})') for j in range(M)] for i in range(N)]
     >>> selection = [[cirq.q(f's({i}, {j})') for j in range(3)] for i in range(3)]
     >>> circuit = cirq.Circuit()
     >>> i_ops = []
-    >>> qm = GreedyQubitManager("ancilla", maximize_reuse=True)
+    >>> qm = cirq.GreedyQubitManager("ancilla", maximize_reuse=True)
     >>> for i_optree, i_ctrl, i in unary_iteration(0, N, i_ops, [], selection[0], qm):
     ...     circuit.append(i_optree)
     ...     j_ops = []