From cfda75d5df32a6579784775a785a97a683861bf5 Mon Sep 17 00:00:00 2001 From: Tanuj Khattar Date: Sun, 30 Jun 2024 15:21:43 -0700 Subject: [PATCH 1/7] Greedy topological sort of the binst graph to minimize qubit allocations / deallocations --- qualtran/_infra/binst_graph_iterators.py | 98 +++++++++++++++++++ qualtran/_infra/binst_graph_iterators_test.py | 63 ++++++++++++ qualtran/_infra/composite_bloq.py | 3 +- qualtran/cirq_interop/_bloq_to_cirq.py | 32 +++--- qualtran/cirq_interop/_bloq_to_cirq_test.py | 22 +++++ 5 files changed, 199 insertions(+), 19 deletions(-) create mode 100644 qualtran/_infra/binst_graph_iterators.py create mode 100644 qualtran/_infra/binst_graph_iterators_test.py diff --git a/qualtran/_infra/binst_graph_iterators.py b/qualtran/_infra/binst_graph_iterators.py new file mode 100644 index 000000000..83eeab9c8 --- /dev/null +++ b/qualtran/_infra/binst_graph_iterators.py @@ -0,0 +1,98 @@ +# Copyright 2024 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 heapq +from collections import Counter +from typing import Dict, Iterator, List, TYPE_CHECKING + +import attrs +import networkx as nx + +if TYPE_CHECKING: + from qualtran import BloqInstance + +_INFINITY: int = int(1e18) + + +def _priority(node: 'BloqInstance') -> int: + from qualtran._infra.gate_with_registers import total_bits + from qualtran._infra.quantum_graph import DanglingT + from qualtran.bloqs.bookkeeping import Allocate, Free + + if isinstance(node, DanglingT): + return 0 + + if node.bloq_is(Allocate): + return _INFINITY + + if node.bloq_is(Free): + return -_INFINITY + + signature = node.bloq.signature + return total_bits(signature.rights()) - total_bits(signature.lefts()) + + +@attrs.frozen(order=True) +class PrioritizedItem: + item: 'BloqInstance' = attrs.field(eq=_priority, order=_priority) + priority: int + + +def greedy_topological_sort(binst_graph: nx.DiGraph) -> Iterator['BloqInstance']: + """Stable greedy topological sorting for the bloq instance graph. + + Topological sorting for the Bloq Instances graph which maintains a priority queue + instead of a queue. Priority for each bloq is a tuple of the form + (_priority(bloq), insertion_index); where each term corresponds to + + ### Priority of a bloq + - 0: For Left / Right Dangling bloqs. + - +Infinity / -Infinity: For `Allocate` / `Free` bloqs. + - total_bits(right registers) - total_bits(left registers): For all other bloqs. + + ### Insertion Index + `insertion_index` is a unique ID used to breaks ties between bloqs that have the + same priority and follows from the order of nodes inserted in the networkx Graph. + + The stability condition guarantees that two networkx graphs constructed with + identical ordering of Graph.nodes and Graph.edges will have the same topological + sorting. + + Args: + binst_graph: A networkx DiGraph with `BloqInstances` as nodes. Usually obtained + from `cbloq._binst_graph` where `cbloq` is a `CompositeBloq`. + + Yields: + Nodes from the input graph returned in a greedy topological sorted order with the + goal to minimize qubit allocations and deallocations by pushing allocations to the + right and de-allocations to the left. + """ + heap: List[PrioritizedItem] = [] + idx: int = 0 + in_degree: Dict[BloqInstance, int] = Counter() + + for x in binst_graph.nodes(): + in_degree[x] = binst_graph.in_degree(x) + if not in_degree[x]: + heapq.heappush(heap, PrioritizedItem(x, idx)) + idx = idx + 1 + + while heap: + x = heapq.heappop(heap).item + yield x + for y in binst_graph.neighbors(x): + in_degree[y] -= 1 + if not in_degree[y]: + heapq.heappush(heap, PrioritizedItem(y, idx)) + idx = idx + 1 diff --git a/qualtran/_infra/binst_graph_iterators_test.py b/qualtran/_infra/binst_graph_iterators_test.py new file mode 100644 index 000000000..58dbb13c8 --- /dev/null +++ b/qualtran/_infra/binst_graph_iterators_test.py @@ -0,0 +1,63 @@ +# Copyright 2024 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. + +from attrs import frozen + +from qualtran import ( + Bloq, + BloqBuilder, + BloqInstance, + LeftDangle, + QAny, + QBit, + RightDangle, + Signature, + SoquetT, +) +from qualtran._infra.binst_graph_iterators import greedy_topological_sort +from qualtran.bloqs.basic_gates import CNOT +from qualtran.bloqs.bookkeeping import Allocate, Free + + +@frozen +class MultiAlloc(Bloq): + rounds: int = 2 + + @property + def signature(self) -> Signature: + return Signature.build(q=1) + + def build_composite_bloq(self, bb: BloqBuilder, q: SoquetT) -> dict[str, SoquetT]: + for _ in range(self.rounds): + a = bb.allocate(1) + a, q = bb.add(CNOT(), ctrl=a, target=q) + bb.free(a) + + return {"q": q} + + +def test_greedy_topological_sort(): + bloq = MultiAlloc() + binst_graph = bloq.decompose_bloq()._binst_graph + greedy_toposort = [*greedy_topological_sort(binst_graph)] + assert greedy_toposort == [ + LeftDangle, + BloqInstance(bloq=Allocate(dtype=QAny(bitsize=1)), i=0), + BloqInstance(bloq=CNOT(), i=1), + BloqInstance(bloq=Free(dtype=QBit()), i=2), + BloqInstance(bloq=Allocate(dtype=QAny(bitsize=1)), i=3), + BloqInstance(bloq=CNOT(), i=4), + BloqInstance(bloq=Free(dtype=QBit()), i=5), + RightDangle, + ] diff --git a/qualtran/_infra/composite_bloq.py b/qualtran/_infra/composite_bloq.py index e483774ee..245e5fac4 100644 --- a/qualtran/_infra/composite_bloq.py +++ b/qualtran/_infra/composite_bloq.py @@ -40,6 +40,7 @@ import sympy from numpy.typing import NDArray +from .binst_graph_iterators import greedy_topological_sort from .bloq import Bloq, DecomposeNotImplementedError, DecomposeTypeError from .data_types import check_dtypes_consistent, QAny, QBit, QDType from .quantum_graph import BloqInstance, Connection, DanglingT, LeftDangle, RightDangle, Soquet @@ -248,7 +249,7 @@ def iter_bloqnections( a predecessor and again as a successor. """ g = self._binst_graph - for binst in nx.topological_sort(g): + for binst in greedy_topological_sort(g): if isinstance(binst, DanglingT): continue pred_cxns, succ_cxns = _binst_to_cxns(binst, binst_graph=g) diff --git a/qualtran/cirq_interop/_bloq_to_cirq.py b/qualtran/cirq_interop/_bloq_to_cirq.py index 0face5128..9ab2fd19e 100644 --- a/qualtran/cirq_interop/_bloq_to_cirq.py +++ b/qualtran/cirq_interop/_bloq_to_cirq.py @@ -34,6 +34,7 @@ Signature, Soquet, ) +from qualtran._infra.binst_graph_iterators import greedy_topological_sort from qualtran._infra.composite_bloq import _binst_to_cxns from qualtran._infra.gate_with_registers import ( _get_all_and_output_quregs_from_input, @@ -266,23 +267,18 @@ def _cbloq_to_cirq_circuit( for reg in signature.lefts() for idx in reg.all_idxs() } - moments: List[cirq.Moment] = [] - for binsts in nx.topological_generations(binst_graph): - moment: List[cirq.Operation] = [] - - for binst in binsts: - if binst is LeftDangle: - continue - pred_cxns, succ_cxns = _binst_to_cxns(binst, binst_graph=binst_graph) - if binst is RightDangle: - _track_soq_name_changes(pred_cxns, qvar_to_qreg) - continue - - op = _bloq_to_cirq_op(binst.bloq, pred_cxns, succ_cxns, qvar_to_qreg, qubit_manager) - if op is not None: - moment.append(op) - if moment: - moments.append(cirq.Moment(moment)) + ops: List[cirq.Operation] = [] + for binst in greedy_topological_sort(binst_graph): + if binst is LeftDangle: + continue + pred_cxns, succ_cxns = _binst_to_cxns(binst, binst_graph=binst_graph) + if binst is RightDangle: + _track_soq_name_changes(pred_cxns, qvar_to_qreg) + continue + + op = _bloq_to_cirq_op(binst.bloq, pred_cxns, succ_cxns, qvar_to_qreg, qubit_manager) + if op is not None: + ops.append(op) # Find output Cirq quregs using `qvar_to_qreg` mapping for registers in `signature.rights()`. def _f_quregs(reg: Register) -> CirqQuregT: @@ -294,7 +290,7 @@ def _f_quregs(reg: Register) -> CirqQuregT: out_quregs = {reg.name: _f_quregs(reg) for reg in signature.rights()} - return cirq.FrozenCircuit(moments), out_quregs + return cirq.FrozenCircuit(ops), out_quregs def _wire_symbol_to_cirq_diagram_info( diff --git a/qualtran/cirq_interop/_bloq_to_cirq_test.py b/qualtran/cirq_interop/_bloq_to_cirq_test.py index 9d9ab0f2e..2e9772b65 100644 --- a/qualtran/cirq_interop/_bloq_to_cirq_test.py +++ b/qualtran/cirq_interop/_bloq_to_cirq_test.py @@ -23,6 +23,7 @@ from qualtran.bloqs.basic_gates import Toffoli, XGate from qualtran.bloqs.factoring import ModExp from qualtran.bloqs.mcmt.and_bloq import And, MultiAnd +from qualtran.bloqs.state_preparation import PrepareUniformSuperposition from qualtran.cirq_interop._bloq_to_cirq import BloqAsCirqGate, CirqQuregT from qualtran.cirq_interop.t_complexity_protocol import t_complexity from qualtran.testing import execute_notebook @@ -148,6 +149,27 @@ def test_multi_and_allocates(): assert sorted(out_quregs.keys()) == ['ctrl', 'junk', 'target'] +def test_flat_cbloq_to_cirq_circuit_minimizes_qubit_allocation(): + bloq = PrepareUniformSuperposition(n=3, cvs=(1,)) + qm = cirq.GreedyQubitManager(prefix='anc', maximize_reuse=True) + cbloq = bloq.as_composite_bloq() + assert len(cbloq.to_cirq_circuit(qubit_manager=qm).all_qubits()) == 3 + cbloq = bloq.decompose_bloq() + assert len(cbloq.to_cirq_circuit(qubit_manager=qm).all_qubits()) == 5 + cbloq = bloq.decompose_bloq().flatten_once() + assert len(cbloq.to_cirq_circuit(qubit_manager=qm).all_qubits()) == 7 + qm = cirq.GreedyQubitManager(prefix='anc', maximize_reuse=True) + # Note: This should also be 7 but to work correctly, it relies on + # `greedy_topological_sort` iterating on allocation nodes in insertion order. + # `cbloq.flatten()` preserves this now because cbloq.iter_bloqnections is also + # updated to use `greedy_topological_sort` instead of `nx.topological_sort`. + # In general, we should have a more stable way to preserve this property, + # potentially by maintaing a sorted order in `binst.i`; + # xref: https://github.com/quantumlib/Qualtran/issues/1098 + cbloq = bloq.decompose_bloq().flatten() + assert len(cbloq.to_cirq_circuit(qubit_manager=qm).all_qubits()) == 7 + + def test_contruct_op_from_gate(): and_gate = And() in_quregs = {'ctrl': np.array([*cirq.LineQubit.range(2)]).reshape(2, 1)} From b38b230f245eaa260b030b093fb56df93bb0a507 Mon Sep 17 00:00:00 2001 From: Tanuj Khattar Date: Sun, 30 Jun 2024 15:26:04 -0700 Subject: [PATCH 2/7] Docstring for _PrioritizedItem --- qualtran/_infra/binst_graph_iterators.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/qualtran/_infra/binst_graph_iterators.py b/qualtran/_infra/binst_graph_iterators.py index 83eeab9c8..dba247abc 100644 --- a/qualtran/_infra/binst_graph_iterators.py +++ b/qualtran/_infra/binst_graph_iterators.py @@ -44,7 +44,9 @@ def _priority(node: 'BloqInstance') -> int: @attrs.frozen(order=True) -class PrioritizedItem: +class _PrioritizedItem: + """Helper dataclass to insert items in a heap as part of `greedy_topological_sort`.""" + item: 'BloqInstance' = attrs.field(eq=_priority, order=_priority) priority: int @@ -78,14 +80,14 @@ def greedy_topological_sort(binst_graph: nx.DiGraph) -> Iterator['BloqInstance'] goal to minimize qubit allocations and deallocations by pushing allocations to the right and de-allocations to the left. """ - heap: List[PrioritizedItem] = [] + heap: List[_PrioritizedItem] = [] idx: int = 0 in_degree: Dict[BloqInstance, int] = Counter() for x in binst_graph.nodes(): in_degree[x] = binst_graph.in_degree(x) if not in_degree[x]: - heapq.heappush(heap, PrioritizedItem(x, idx)) + heapq.heappush(heap, _PrioritizedItem(x, idx)) idx = idx + 1 while heap: @@ -94,5 +96,5 @@ def greedy_topological_sort(binst_graph: nx.DiGraph) -> Iterator['BloqInstance'] for y in binst_graph.neighbors(x): in_degree[y] -= 1 if not in_degree[y]: - heapq.heappush(heap, PrioritizedItem(y, idx)) + heapq.heappush(heap, _PrioritizedItem(y, idx)) idx = idx + 1 From ae3fd59a9b140ca814cdcbc7e51d74267cebf7fb Mon Sep 17 00:00:00 2001 From: Tanuj Khattar Date: Mon, 1 Jul 2024 00:37:00 -0700 Subject: [PATCH 3/7] Fix failing CI --- qualtran/bloqs/chemistry/thc/thc.ipynb | 2 +- qualtran/drawing/qpic_diagram_test.py | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/qualtran/bloqs/chemistry/thc/thc.ipynb b/qualtran/bloqs/chemistry/thc/thc.ipynb index bebdac055..c5e42144d 100644 --- a/qualtran/bloqs/chemistry/thc/thc.ipynb +++ b/qualtran/bloqs/chemistry/thc/thc.ipynb @@ -433,7 +433,7 @@ " print(f\"{k+':':20s} qualtran = {binned_counts[k]:5d} vs paper cost = {v:5d}.\")\n", "\n", "print(f\"Total cost = {sum(v for v in binned_counts.values())}\")\n", - "assert binned_counts['data_loading'] == 304" + "assert binned_counts['data_loading'] == 296" ] }, { diff --git a/qualtran/drawing/qpic_diagram_test.py b/qualtran/drawing/qpic_diagram_test.py index 9e9c42d7e..6dfbd6fc8 100644 --- a/qualtran/drawing/qpic_diagram_test.py +++ b/qualtran/drawing/qpic_diagram_test.py @@ -53,10 +53,10 @@ def test_qpic_data_for_reflect_using_prepare(): reg_1 W off reg_2 W off reg_3 W off -reg_4 W off reg[0] W off reg[1] W off -reg_5 W off +reg_4 W off +_empty_wire G:width=65:shape=8 GPhase((-0-1j)) reg:on G:width=25:shape=box \textrm{\scalebox{0.8}{alloc}} reg / \textrm{\scalebox{0.5}{QAny(2)}} reg_1:on G:width=25:shape=box \textrm{\scalebox{0.8}{alloc}} @@ -65,20 +65,20 @@ def test_qpic_data_for_reflect_using_prepare(): reg_2 / \textrm{\scalebox{0.5}{QAny(2)}} reg_3:on G:width=25:shape=box \textrm{\scalebox{0.8}{alloc}} reg_3 / \textrm{\scalebox{0.5}{QAny(1)}} +selection G:width=121:shape=box \textrm{\scalebox{0.8}{StatePreparationAliasSampling}} reg G:width=37:shape=box \textrm{\scalebox{0.8}{sigma\_mu}} reg_1 G:width=17:shape=box \textrm{\scalebox{0.8}{alt}} reg_2 G:width=21:shape=box \textrm{\scalebox{0.8}{keep}} reg_3 G:width=65:shape=box \textrm{\scalebox{0.8}{less\_than\_equal}} +selection:off G:width=5:shape=> \textrm{\scalebox{0.8}{}} reg[0]:on G:width=17:shape=box \textrm{\scalebox{0.8}{[0]}} reg[1]:on G:width=17:shape=box \textrm{\scalebox{0.8}{[1]}} reg_4:on G:width=25:shape=box \textrm{\scalebox{0.8}{alloc}} reg_4 / \textrm{\scalebox{0.5}{QAny(1)}} -_empty_wire G:width=65:shape=8 GPhase((-0-1j)) -selection G:width=121:shape=box \textrm{\scalebox{0.8}{StatePreparationAliasSampling}} reg G:width=37:shape=box \textrm{\scalebox{0.8}{sigma\_mu}} reg_1 G:width=17:shape=box \textrm{\scalebox{0.8}{alt}} reg_2 G:width=21:shape=box \textrm{\scalebox{0.8}{keep}} reg_3 G:width=65:shape=box \textrm{\scalebox{0.8}{less\_than\_equal}} +reg_4 -selection:off G:width=5:shape=> \textrm{\scalebox{0.8}{}} reg[0]:on G:width=17:shape=box \textrm{\scalebox{0.8}{[0]}} reg[1]:on G:width=17:shape=box \textrm{\scalebox{0.8}{[1]}} reg_4 G:width=9:shape=box \textrm{\scalebox{0.8}{Z}} -reg[0] -reg[1] +reg_4 -reg[0]:off G:width=17:shape=box \textrm{\scalebox{0.8}{[0]}} reg[1]:off G:width=17:shape=box \textrm{\scalebox{0.8}{[1]}} reg_5:on G:width=5:shape=< \textrm{\scalebox{0.8}{}} -reg_5 / \textrm{\scalebox{0.5}{BoundedQUInt(2, 4)}} reg_4:off G:width=21:shape=box \textrm{\scalebox{0.8}{free}} -reg_5 G:width=121:shape=box \textrm{\scalebox{0.8}{StatePreparationAliasSampling}} reg G:width=37:shape=box \textrm{\scalebox{0.8}{sigma\_mu}} reg_1 G:width=17:shape=box \textrm{\scalebox{0.8}{alt}} reg_2 G:width=21:shape=box \textrm{\scalebox{0.8}{keep}} reg_3 G:width=65:shape=box \textrm{\scalebox{0.8}{less\_than\_equal}} +reg[0]:off G:width=17:shape=box \textrm{\scalebox{0.8}{[0]}} reg[1]:off G:width=17:shape=box \textrm{\scalebox{0.8}{[1]}} reg_4:on G:width=5:shape=< \textrm{\scalebox{0.8}{}} +reg_4 / \textrm{\scalebox{0.5}{BoundedQUInt(2, 4)}} +reg_4 G:width=121:shape=box \textrm{\scalebox{0.8}{StatePreparationAliasSampling}} reg G:width=37:shape=box \textrm{\scalebox{0.8}{sigma\_mu}} reg_1 G:width=17:shape=box \textrm{\scalebox{0.8}{alt}} reg_2 G:width=21:shape=box \textrm{\scalebox{0.8}{keep}} reg_3 G:width=65:shape=box \textrm{\scalebox{0.8}{less\_than\_equal}} reg:off G:width=21:shape=box \textrm{\scalebox{0.8}{free}} reg_1:off G:width=21:shape=box \textrm{\scalebox{0.8}{free}} reg_2:off G:width=21:shape=box \textrm{\scalebox{0.8}{free}} -reg_3:off G:width=21:shape=box \textrm{\scalebox{0.8}{free}}""", +reg_3:off G:width=21:shape=box \textrm{\scalebox{0.8}{free}} +""", ) From b3a4a92e24c0922a006ad5cfae5eac98e8ad2dd0 Mon Sep 17 00:00:00 2001 From: Tanuj Khattar Date: Mon, 1 Jul 2024 00:47:25 -0700 Subject: [PATCH 4/7] Undo thc.ipynb change --- qualtran/bloqs/chemistry/thc/thc.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qualtran/bloqs/chemistry/thc/thc.ipynb b/qualtran/bloqs/chemistry/thc/thc.ipynb index c5e42144d..bebdac055 100644 --- a/qualtran/bloqs/chemistry/thc/thc.ipynb +++ b/qualtran/bloqs/chemistry/thc/thc.ipynb @@ -433,7 +433,7 @@ " print(f\"{k+':':20s} qualtran = {binned_counts[k]:5d} vs paper cost = {v:5d}.\")\n", "\n", "print(f\"Total cost = {sum(v for v in binned_counts.values())}\")\n", - "assert binned_counts['data_loading'] == 296" + "assert binned_counts['data_loading'] == 304" ] }, { From 308a60f1201f7f2685ba46ed3309f886f330aac5 Mon Sep 17 00:00:00 2001 From: Tanuj Khattar Date: Mon, 1 Jul 2024 00:48:15 -0700 Subject: [PATCH 5/7] Another attempt at cirq interop bugfix --- qualtran/_infra/adjoint.py | 8 ------- qualtran/_infra/gate_with_registers_test.py | 6 ++--- qualtran/bloqs/arithmetic/comparison.py | 3 +++ qualtran/bloqs/bookkeeping/allocate.py | 14 ++++++++++- qualtran/bloqs/bookkeeping/free.py | 10 +++++++- .../reflections/reflection_using_prepare.py | 4 ++-- .../reflection_using_prepare_test.py | 12 ++++++++-- .../prepare_uniform_superposition_test.py | 13 ++++++++++ qualtran/cirq_interop/_bloq_to_cirq.py | 1 - qualtran/cirq_interop/_bloq_to_cirq_test.py | 2 +- .../cirq_interop/_interop_qubit_manager.py | 24 +++++++++++++++++-- 11 files changed, 76 insertions(+), 21 deletions(-) diff --git a/qualtran/_infra/adjoint.py b/qualtran/_infra/adjoint.py index 16b9d2d3d..11766c244 100644 --- a/qualtran/_infra/adjoint.py +++ b/qualtran/_infra/adjoint.py @@ -17,7 +17,6 @@ import cirq from attrs import frozen -from numpy.typing import NDArray from .composite_bloq import _binst_to_cxns, _cxn_to_soq_dict, _map_soqs, _reg_to_soq, BloqBuilder from .gate_with_registers import GateWithRegisters @@ -142,13 +141,6 @@ def decompose_bloq(self) -> 'CompositeBloq': """The decomposition is the adjoint of `subbloq`'s decomposition.""" return self.subbloq.decompose_bloq().adjoint() - def decompose_from_registers( - self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid] # type: ignore[type-var] - ) -> cirq.OP_TREE: - if isinstance(self.subbloq, GateWithRegisters): - return cirq.inverse(self.subbloq.decompose_from_registers(context=context, **quregs)) - return super().decompose_from_registers(context=context, **quregs) - def _circuit_diagram_info_( self, args: 'cirq.CircuitDiagramInfoArgs' ) -> cirq.CircuitDiagramInfo: diff --git a/qualtran/_infra/gate_with_registers_test.py b/qualtran/_infra/gate_with_registers_test.py index 329c20f36..a3735a889 100644 --- a/qualtran/_infra/gate_with_registers_test.py +++ b/qualtran/_infra/gate_with_registers_test.py @@ -141,11 +141,11 @@ def test_gate_with_registers_decompose_from_context_auto_generated(): cirq.testing.assert_has_diagram( circuit, """ -l: ───BloqWithDecompose───X───────free─── +l: ───BloqWithDecompose───X─── │ -r: ───r───────────────────alloc───Z────── +r: ───r───────────────────Z─── │ -t: ───t───────────────────Y────────────── +t: ───t───────────────────Y─── """, ) diff --git a/qualtran/bloqs/arithmetic/comparison.py b/qualtran/bloqs/arithmetic/comparison.py index a5a1e5987..5cd7599e5 100644 --- a/qualtran/bloqs/arithmetic/comparison.py +++ b/qualtran/bloqs/arithmetic/comparison.py @@ -610,6 +610,9 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: def _has_unitary_(self): return True + def adjoint(self) -> 'Bloq': + return self + @bloq_example def _leq_symb() -> LessThanEqual: diff --git a/qualtran/bloqs/bookkeeping/allocate.py b/qualtran/bloqs/bookkeeping/allocate.py index 953807739..cd0bef7b5 100644 --- a/qualtran/bloqs/bookkeeping/allocate.py +++ b/qualtran/bloqs/bookkeeping/allocate.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Any, Dict, Tuple, TYPE_CHECKING +from typing import Any, Dict, Tuple, TYPE_CHECKING, Union import numpy as np import sympy @@ -35,8 +35,11 @@ from qualtran.drawing import directional_text_box, Text, WireSymbol if TYPE_CHECKING: + import cirq import quimb.tensor as qtn + from qualtran.cirq_interop import CirqQuregT + @frozen class Allocate(_BookkeepingBloq): @@ -86,6 +89,15 @@ def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSym assert reg.name == 'reg' return directional_text_box('alloc', Side.RIGHT) + def as_cirq_op( + self, qubit_manager: 'cirq.QubitManager' + ) -> Tuple[Union['cirq.Operation', None], Dict[str, 'CirqQuregT']]: + shape = (*self.signature[0].shape, self.signature[0].bitsize) + return ( + None, + {'reg': np.array(qubit_manager.qalloc(self.signature.n_qubits())).reshape(shape)}, + ) + @bloq_example def _alloc() -> Allocate: diff --git a/qualtran/bloqs/bookkeeping/free.py b/qualtran/bloqs/bookkeeping/free.py index 2aea1c11b..0b15dd2e1 100644 --- a/qualtran/bloqs/bookkeeping/free.py +++ b/qualtran/bloqs/bookkeeping/free.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Any, Dict, Tuple, TYPE_CHECKING +from typing import Any, Dict, Tuple, TYPE_CHECKING, Union import numpy as np import sympy @@ -36,8 +36,10 @@ from qualtran.drawing import directional_text_box, Text, WireSymbol if TYPE_CHECKING: + import cirq import quimb.tensor as qtn + from qualtran.cirq_interop import CirqQuregT from qualtran.simulation.classical_sim import ClassicalValT @@ -95,6 +97,12 @@ def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSym assert reg.name == 'reg' return directional_text_box('free', Side.LEFT) + def as_cirq_op( + self, qubit_manager: 'cirq.QubitManager', reg: 'CirqQuregT' + ) -> Tuple[Union['cirq.Operation', None], Dict[str, 'CirqQuregT']]: + qubit_manager.qfree(reg.flatten().tolist()) + return (None, {}) + @bloq_example def _free() -> Free: diff --git a/qualtran/bloqs/reflections/reflection_using_prepare.py b/qualtran/bloqs/reflections/reflection_using_prepare.py index 88a531596..9dabf41f5 100644 --- a/qualtran/bloqs/reflections/reflection_using_prepare.py +++ b/qualtran/bloqs/reflections/reflection_using_prepare.py @@ -29,12 +29,12 @@ from qualtran.bloqs.basic_gates.global_phase import GlobalPhase from qualtran.bloqs.basic_gates.rotation import ZPowGate from qualtran.bloqs.basic_gates.x_basis import XGate -from qualtran.bloqs.block_encoding.lcu_select_and_prepare import PrepareOracle from qualtran.bloqs.mcmt.multi_control_multi_target_pauli import MultiControlPauli from qualtran.resource_counting.generalizers import ignore_split_join from qualtran.symbolics.types import SymbolicInt if TYPE_CHECKING: + from qualtran.bloqs.block_encoding.lcu_select_and_prepare import PrepareOracle from qualtran.resource_counting import BloqCountT, SympySymbolAllocator @@ -77,7 +77,7 @@ class ReflectionUsingPrepare(SpecializedSingleQubitControlledGate): Babbush et. al. (2018). Figure 1. """ - prepare_gate: PrepareOracle + prepare_gate: 'PrepareOracle' control_val: Optional[int] = None global_phase: complex = 1 eps: float = 1e-11 diff --git a/qualtran/bloqs/reflections/reflection_using_prepare_test.py b/qualtran/bloqs/reflections/reflection_using_prepare_test.py index 6f2cb9677..9aac402c4 100644 --- a/qualtran/bloqs/reflections/reflection_using_prepare_test.py +++ b/qualtran/bloqs/reflections/reflection_using_prepare_test.py @@ -20,7 +20,7 @@ import pytest from numpy.typing import NDArray -from qualtran import Bloq +from qualtran import Adjoint, Bloq from qualtran._infra.gate_with_registers import get_named_qubits from qualtran.bloqs.arithmetic import LessThanConstant, LessThanEqual from qualtran.bloqs.basic_gates import ZPowGate @@ -30,6 +30,7 @@ from qualtran.bloqs.reflections.prepare_identity import PrepareIdentity from qualtran.bloqs.reflections.reflection_using_prepare import ReflectionUsingPrepare from qualtran.bloqs.state_preparation import StatePreparationAliasSampling +from qualtran.cirq_interop import BloqAsCirqGate from qualtran.cirq_interop.testing import GateHelper from qualtran.resource_counting.generalizers import ( ignore_alloc_free, @@ -58,6 +59,13 @@ def keep(op: cirq.Operation): ret = op in gateset_to_keep if op.gate is not None and isinstance(op.gate, cirq.ops.raw_types._InverseCompositeGate): ret |= op.gate._original in gateset_to_keep + if op.gate is not None and isinstance(op.gate, Adjoint): + subgate = ( + op.gate.subbloq + if isinstance(op.gate.subbloq, cirq.Gate) + else BloqAsCirqGate(op.gate.subbloq) + ) + ret |= subgate in gateset_to_keep return ret @@ -73,7 +81,7 @@ def construct_gate_helper_and_qubit_order(gate, decompose_once: bool = False): ) ordered_input = list(itertools.chain(*g.quregs.values())) qubit_order = cirq.QubitOrder.explicit(ordered_input, fallback=cirq.QubitOrder.DEFAULT) - assert len(circuit.all_qubits()) < 30 + assert len(circuit.all_qubits()) < 24 return g, qubit_order, circuit diff --git a/qualtran/bloqs/state_preparation/prepare_uniform_superposition_test.py b/qualtran/bloqs/state_preparation/prepare_uniform_superposition_test.py index a239cd4dd..7409588e6 100644 --- a/qualtran/bloqs/state_preparation/prepare_uniform_superposition_test.py +++ b/qualtran/bloqs/state_preparation/prepare_uniform_superposition_test.py @@ -104,3 +104,16 @@ def test_prepare_uniform_superposition_consistent_protocols(): PrepareUniformSuperposition(5, cvs=()), PrepareUniformSuperposition(5, cvs=[]), ) + + +def test_prepare_uniform_superposition_adjoint(): + n = 3 + target = cirq.NamedQubit.range((n - 1).bit_length(), prefix='target') + control = [cirq.NamedQubit('control')] + op = PrepareUniformSuperposition(n, cvs=(0,)).on_registers(ctrl=control, target=target) + gqm = cirq.GreedyQubitManager(prefix="_ancilla", maximize_reuse=True) + context = cirq.DecompositionContext(gqm) + circuit = cirq.Circuit(op, cirq.decompose(cirq.inverse(op), context=context)) + identity = cirq.Circuit(cirq.identity_each(*circuit.all_qubits())).final_state_vector() + result = cirq.Simulator(dtype=np.complex128).simulate(circuit) + np.testing.assert_allclose(result.final_state_vector, identity, atol=1e-8) diff --git a/qualtran/cirq_interop/_bloq_to_cirq.py b/qualtran/cirq_interop/_bloq_to_cirq.py index 9ab2fd19e..a20a81b15 100644 --- a/qualtran/cirq_interop/_bloq_to_cirq.py +++ b/qualtran/cirq_interop/_bloq_to_cirq.py @@ -230,7 +230,6 @@ def _bloq_to_cirq_op( del qvar_to_qreg[soq] op, out_quregs = bloq.as_cirq_op(qubit_manager=qubit_manager, **in_quregs) - # 2. Update the mappings based on output soquets and `out_quregs`. for cxn in succ_cxns: soq = cxn.left diff --git a/qualtran/cirq_interop/_bloq_to_cirq_test.py b/qualtran/cirq_interop/_bloq_to_cirq_test.py index 2e9772b65..fd6cd2d76 100644 --- a/qualtran/cirq_interop/_bloq_to_cirq_test.py +++ b/qualtran/cirq_interop/_bloq_to_cirq_test.py @@ -209,7 +209,7 @@ def test_bloq_as_cirq_gate_left_register(): bb.free(q) cbloq = bb.finalize() circuit = cbloq.to_cirq_circuit() - cirq.testing.assert_has_diagram(circuit, """_c(0): ───alloc───X───free───""") + cirq.testing.assert_has_diagram(circuit, """_c(0): ───X───""") def test_bloq_as_cirq_gate_for_mod_exp(): diff --git a/qualtran/cirq_interop/_interop_qubit_manager.py b/qualtran/cirq_interop/_interop_qubit_manager.py index 0cb2cccd1..da3e8d2ee 100644 --- a/qualtran/cirq_interop/_interop_qubit_manager.py +++ b/qualtran/cirq_interop/_interop_qubit_manager.py @@ -28,10 +28,30 @@ def __init__(self, qm: Optional[cirq.QubitManager] = None): self._managed_qubits: Set[cirq.Qid] = set() def qalloc(self, n: int, dim: int = 2) -> List['cirq.Qid']: - return self._qm.qalloc(n, dim) + ret: List['cirq.Qid'] = [] + qubits_to_free: List['cirq.Qid'] = [] + while len(ret) < n: + new_alloc = self._qm.qalloc(n - len(ret), dim) + for q in new_alloc: + if q in self._managed_qubits: + qubits_to_free.append(q) + else: + ret.append(q) + self._qm.qfree(qubits_to_free) + return ret def qborrow(self, n: int, dim: int = 2) -> List['cirq.Qid']: - return self._qm.qborrow(n, dim) + ret: List['cirq.Qid'] = [] + qubits_to_free: List['cirq.Qid'] = [] + while len(ret) < n: + new_alloc = self._qm.qborrow(n - len(ret), dim) + for q in new_alloc: + if q in self._managed_qubits: + qubits_to_free.append(q) + else: + ret.append(q) + self._qm.qfree(qubits_to_free) + return ret def manage_qubits(self, qubits: Iterable[cirq.Qid]): self._managed_qubits |= set(qubits) From a107562a1321234c2e63b50ea45eed246a47d21c Mon Sep 17 00:00:00 2001 From: Tanuj Khattar Date: Sat, 6 Jul 2024 10:43:28 -0700 Subject: [PATCH 6/7] Fix formatting --- qualtran/bloqs/bookkeeping/allocate.py | 4 ++-- qualtran/bloqs/bookkeeping/free.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/qualtran/bloqs/bookkeeping/allocate.py b/qualtran/bloqs/bookkeeping/allocate.py index efab646e1..7ab3fb383 100644 --- a/qualtran/bloqs/bookkeeping/allocate.py +++ b/qualtran/bloqs/bookkeeping/allocate.py @@ -12,11 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import cached_property -from typing import Dict, List, Tuple, Union, TYPE_CHECKING +from typing import Dict, List, Tuple, TYPE_CHECKING, Union +import numpy as np import sympy from attrs import frozen -import numpy as np from qualtran import ( Bloq, diff --git a/qualtran/bloqs/bookkeeping/free.py b/qualtran/bloqs/bookkeeping/free.py index 40b7867dd..f8a2a0700 100644 --- a/qualtran/bloqs/bookkeeping/free.py +++ b/qualtran/bloqs/bookkeeping/free.py @@ -13,7 +13,7 @@ # limitations under the License. from functools import cached_property -from typing import Dict, List, Tuple, Union, TYPE_CHECKING +from typing import Dict, List, Tuple, TYPE_CHECKING, Union import sympy from attrs import frozen From 3a1d0140eaaf5d1762bb250fb1af8f885b771995 Mon Sep 17 00:00:00 2001 From: Tanuj Khattar Date: Sat, 6 Jul 2024 12:01:57 -0700 Subject: [PATCH 7/7] Fix as_cirq_op for SGate(is_adjoint=True) --- qualtran/bloqs/basic_gates/s_gate.py | 3 ++- qualtran/bloqs/basic_gates/s_gate_test.py | 3 ++- .../bloqs/data_loading/select_swap_qrom_test.py | 6 +++--- qualtran/bloqs/swap_network/cswap_approx_test.py | 14 +++++++++++++- qualtran/bloqs/swap_network/swap_with_zero.py | 2 +- 5 files changed, 21 insertions(+), 7 deletions(-) diff --git a/qualtran/bloqs/basic_gates/s_gate.py b/qualtran/bloqs/basic_gates/s_gate.py index a716d06fe..98449242d 100644 --- a/qualtran/bloqs/basic_gates/s_gate.py +++ b/qualtran/bloqs/basic_gates/s_gate.py @@ -73,7 +73,8 @@ def as_cirq_op( import cirq (q,) = q - return cirq.S(q), {'q': np.array([q])} + p = -1 if self.is_adjoint else 1 + return cirq.S(q) ** p, {'q': np.array([q])} def pretty_name(self) -> str: maybe_dag = '†' if self.is_adjoint else '' diff --git a/qualtran/bloqs/basic_gates/s_gate_test.py b/qualtran/bloqs/basic_gates/s_gate_test.py index 0fc6fa8c1..0af74b219 100644 --- a/qualtran/bloqs/basic_gates/s_gate_test.py +++ b/qualtran/bloqs/basic_gates/s_gate_test.py @@ -32,9 +32,10 @@ def test_to_cirq(): bb = BloqBuilder() q = bb.add(PlusState()) q = bb.add(SGate(), q=q) + q = bb.add(SGate().adjoint(), q=q) cbloq = bb.finalize(q=q) circuit = cbloq.to_cirq_circuit() - cirq.testing.assert_has_diagram(circuit, "_c(0): ───H───S───") + cirq.testing.assert_has_diagram(circuit, "_c(0): ───H───S───S^-1───") def test_tensors(): diff --git a/qualtran/bloqs/data_loading/select_swap_qrom_test.py b/qualtran/bloqs/data_loading/select_swap_qrom_test.py index 83287702c..11dd47ef6 100644 --- a/qualtran/bloqs/data_loading/select_swap_qrom_test.py +++ b/qualtran/bloqs/data_loading/select_swap_qrom_test.py @@ -62,9 +62,9 @@ def test_select_swap_qrom(data, block_size): cirq.decompose_once(qrom.on_registers(**qubit_regs), context=context) ) - dirty_target_ancilla = [ - q for q in qrom_circuit.all_qubits() if isinstance(q, cirq.ops.BorrowableQubit) - ] + dirty_target_ancilla = sorted( + qrom_circuit.all_qubits() - set(q for qs in qubit_regs.values() for q in qs.flatten()) + ) circuit = cirq.Circuit( # Prepare dirty ancillas in an arbitrary state. diff --git a/qualtran/bloqs/swap_network/cswap_approx_test.py b/qualtran/bloqs/swap_network/cswap_approx_test.py index 7ed1700fd..665e1cb12 100644 --- a/qualtran/bloqs/swap_network/cswap_approx_test.py +++ b/qualtran/bloqs/swap_network/cswap_approx_test.py @@ -11,10 +11,11 @@ # 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 random from typing import Dict, Tuple, Union +import cirq +import numpy as np import pytest import sympy @@ -38,6 +39,17 @@ def test_cswap_approx_decomp(): assert_valid_bloq_decomposition(csa) +def test_cswap_approx_decomposition(): + csa = CSwapApprox(4) + circuit = ( + csa.as_composite_bloq().to_cirq_circuit() + + csa.adjoint().as_composite_bloq().to_cirq_circuit() + ) + initial_state = cirq.testing.random_superposition(2**9, random_state=1234) + result = cirq.Simulator(dtype=np.complex128).simulate(circuit, initial_state=initial_state) + np.testing.assert_allclose(result.final_state_vector, initial_state) + + @pytest.mark.parametrize('n', [5, 32]) def test_approx_cswap_t_count(n): cswap = CSwapApprox(bitsize=n) diff --git a/qualtran/bloqs/swap_network/swap_with_zero.py b/qualtran/bloqs/swap_network/swap_with_zero.py index be301aa5a..112bd6e91 100644 --- a/qualtran/bloqs/swap_network/swap_with_zero.py +++ b/qualtran/bloqs/swap_network/swap_with_zero.py @@ -153,7 +153,7 @@ def build_composite_bloq( def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: num_swaps = prod(x for x in self.n_target_registers) - 1 - return {(CSwapApprox(self.target_bitsize), num_swaps)} + return {(self.cswap_n, num_swaps)} def _circuit_diagram_info_(self, args) -> cirq.CircuitDiagramInfo: from qualtran.cirq_interop._bloq_to_cirq import _wire_symbol_to_cirq_diagram_info