|
| 1 | +# Copyright 2023 The Cirq Developers |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# https://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +from typing import Iterable, List, Set, TYPE_CHECKING |
| 16 | + |
| 17 | +from cirq.ops import named_qubit, qid_util, qubit_manager |
| 18 | + |
| 19 | +if TYPE_CHECKING: |
| 20 | + import cirq |
| 21 | + |
| 22 | + |
| 23 | +class GreedyQubitManager(qubit_manager.QubitManager): |
| 24 | + """Greedy allocator that maximizes/minimizes qubit reuse based on a configurable parameter. |
| 25 | +
|
| 26 | + GreedyQubitManager can be configured, using `maximize_reuse` flag, to work in one of two modes: |
| 27 | + - Minimize qubit reuse (maximize_reuse=False): For a fixed width, this mode uses a FIFO (First |
| 28 | + in First out) strategy s.t. next allocated qubit is one which was freed the earliest. |
| 29 | + - Maximize qubit reuse (maximize_reuse=True): For a fixed width, this mode uses a LIFO (Last in |
| 30 | + First out) strategy s.t. the next allocated qubit is one which was freed the latest. |
| 31 | +
|
| 32 | + If the requested qubits are more than the set of free qubits, the qubit manager automatically |
| 33 | + resizes the size of the managed qubit pool and adds new free qubits, that have their last |
| 34 | + freed time to be -infinity. |
| 35 | +
|
| 36 | + For borrowing qubits, the qubit manager simply delegates borrow requests to `self.qalloc`, thus |
| 37 | + always allocating new clean qubits. |
| 38 | + """ |
| 39 | + |
| 40 | + def __init__(self, prefix: str, *, size: int = 0, maximize_reuse: bool = False): |
| 41 | + """Initializes `GreedyQubitManager` |
| 42 | +
|
| 43 | + Args: |
| 44 | + prefix: The prefix to use for naming new clean ancillas allocated by the qubit manager. |
| 45 | + The i'th allocated qubit is of the type `cirq.NamedQubit(f'{prefix}_{i}')`. |
| 46 | + size: The initial size of the pool of ancilla qubits managed by the qubit manager. The |
| 47 | + qubit manager can automatically resize itself when the allocation request |
| 48 | + exceeds the number of available qubits. |
| 49 | + maximize_reuse: Flag to control a FIFO vs LIFO strategy, defaults to False (FIFO). |
| 50 | + """ |
| 51 | + self._prefix = prefix |
| 52 | + self._used_qubits: Set['cirq.Qid'] = set() |
| 53 | + self._free_qubits: List['cirq.Qid'] = [] |
| 54 | + self._size = 0 |
| 55 | + self.maximize_reuse = maximize_reuse |
| 56 | + self.resize(size) |
| 57 | + |
| 58 | + def _allocate_qid(self, name: str, dim: int) -> 'cirq.Qid': |
| 59 | + return qid_util.q(name) if dim == 2 else named_qubit.NamedQid(name, dimension=dim) |
| 60 | + |
| 61 | + def resize(self, new_size: int, dim: int = 2) -> None: |
| 62 | + if new_size <= self._size: |
| 63 | + return |
| 64 | + new_qubits: List['cirq.Qid'] = [ |
| 65 | + self._allocate_qid(f'{self._prefix}_{s}', dim) for s in range(self._size, new_size) |
| 66 | + ] |
| 67 | + self._free_qubits = new_qubits + self._free_qubits |
| 68 | + self._size = new_size |
| 69 | + |
| 70 | + def qalloc(self, n: int, dim: int = 2) -> List['cirq.Qid']: |
| 71 | + if not n: |
| 72 | + return [] |
| 73 | + self.resize(self._size + n - len(self._free_qubits), dim=dim) |
| 74 | + ret_qubits = self._free_qubits[-n:] if self.maximize_reuse else self._free_qubits[:n] |
| 75 | + self._free_qubits = self._free_qubits[:-n] if self.maximize_reuse else self._free_qubits[n:] |
| 76 | + self._used_qubits.update(ret_qubits) |
| 77 | + return ret_qubits |
| 78 | + |
| 79 | + def qfree(self, qubits: Iterable['cirq.Qid']) -> None: |
| 80 | + qs = list(dict(zip(qubits, qubits)).keys()) |
| 81 | + assert self._used_qubits.issuperset(qs), "Only managed qubits currently in-use can be freed" |
| 82 | + self._used_qubits = self._used_qubits.difference(qs) |
| 83 | + self._free_qubits.extend(qs) |
| 84 | + |
| 85 | + def qborrow(self, n: int, dim: int = 2) -> List['cirq.Qid']: |
| 86 | + return self.qalloc(n, dim) |
0 commit comments