diff --git a/qualtran/_infra/binst_graph_iterators.py b/qualtran/_infra/binst_graph_iterators.py new file mode 100644 index 000000000..37a398961 --- /dev/null +++ b/qualtran/_infra/binst_graph_iterators.py @@ -0,0 +1,76 @@ +# 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 typing import Iterator, TYPE_CHECKING + +import networkx as nx + +if TYPE_CHECKING: + from qualtran import BloqInstance + +_ALLOCATION_PRIORITY: int = int(1e16) +"""A large constant value to ensure that allocations are performed as late as possible +and de-allocations (with -_ALLOCATION_PRIORITY priority) are performed as early as possible. +To determine ordering among allocations, we may add a priority to this base value.""" + + +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 _ALLOCATION_PRIORITY + + if node.bloq_is(Free): + return -_ALLOCATION_PRIORITY + + signature = node.bloq.signature + return total_bits(signature.rights()) - total_bits(signature.lefts()) + + +def greedy_topological_sort(binst_graph: nx.DiGraph) -> Iterator['BloqInstance']: + """Stable greedy topological sorting for the bloq instance graph to minimize qubit counts. + + 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. The method delegates to `networkx.lexicographical_topological_sort` with + the `_priority` function used as a key. + + 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. + """ + yield from nx.lexicographical_topological_sort(binst_graph, key=_priority) 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 b5b28f59b..221c7508a 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 @@ -253,7 +254,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/bloqs/swap_network/swap_with_zero_test.py b/qualtran/bloqs/swap_network/swap_with_zero_test.py index 87fcdfee8..c5a7c2abc 100644 --- a/qualtran/bloqs/swap_network/swap_with_zero_test.py +++ b/qualtran/bloqs/swap_network/swap_with_zero_test.py @@ -102,41 +102,42 @@ def test_swap_with_zero_cirq_gate_diagram(): def test_swap_with_zero_cirq_gate_diagram_multi_dim(): gate = SwapWithZero((2, 1), 2, (3, 2)) gh = cq_testing.GateHelper(gate) + # Bloq -> Cirq conversion preserves insertion ordering when all operations are THRU + # operations cirq.testing.assert_has_diagram( cirq.Circuit(gh.operation, cirq.decompose_once(gh.operation)), """ ┌──────────────────┐ selection0_0: ───────@(r⇋0)────────────────────────────────────────────────────@(approx)─── │ │ -selection0_1: ───────@(r⇋0)──────────────────────────────@(approx)─────────────┼─────────── - │ │ │ -selection1_: ────────@(r⇋0)─────@(approx)───@(approx)────┼────────@(approx)────┼─────────── +selection0_1: ───────@(r⇋0)───────────────────────────────────────@(approx)────┼─────────── + │ │ │ +selection1_: ────────@(r⇋0)─────@(approx)───@(approx)────@(approx)┼────────────┼─────────── │ │ │ │ │ │ -targets[0, 0][0]: ───swap_0_0───×(x)────────┼────────────×(x)─────┼────────────×(x)──────── +targets[0, 0][0]: ───swap_0_0───×(x)────────┼────────────┼────────×(x)─────────×(x)──────── │ │ │ │ │ │ -targets[0, 0][1]: ───swap_0_0───×(x)────────┼────────────×(x)─────┼────────────×(x)──────── +targets[0, 0][1]: ───swap_0_0───×(x)────────┼────────────┼────────×(x)─────────×(x)──────── │ │ │ │ │ │ targets[0, 1][0]: ───swap_0_1───×(y)────────┼────────────┼────────┼────────────┼─────────── │ │ │ │ │ │ targets[0, 1][1]: ───swap_0_1───×(y)────────┼────────────┼────────┼────────────┼─────────── │ │ │ │ │ -targets[1, 0][0]: ───swap_1_0───────────────×(x)─────────×(y)─────┼────────────┼─────────── +targets[1, 0][0]: ───swap_1_0───────────────×(x)─────────┼────────×(y)─────────┼─────────── │ │ │ │ │ -targets[1, 0][1]: ───swap_1_0───────────────×(x)─────────×(y)─────┼────────────┼─────────── - │ │ │ │ -targets[1, 1][0]: ───swap_1_1───────────────×(y)──────────────────┼────────────┼─────────── - │ │ │ │ -targets[1, 1][1]: ───swap_1_1───────────────×(y)──────────────────┼────────────┼─────────── - │ │ │ -targets[2, 0][0]: ───swap_2_0─────────────────────────────────────×(x)─────────×(y)──────── - │ │ │ -targets[2, 0][1]: ───swap_2_0─────────────────────────────────────×(x)─────────×(y)──────── - │ │ -targets[2, 1][0]: ───swap_2_1─────────────────────────────────────×(y)───────────────────── - │ │ -targets[2, 1][1]: ───swap_2_1─────────────────────────────────────×(y)───────────────────── - └──────────────────┘ -""", +targets[1, 0][1]: ───swap_1_0───────────────×(x)─────────┼────────×(y)─────────┼─────────── + │ │ │ │ +targets[1, 1][0]: ───swap_1_1───────────────×(y)─────────┼─────────────────────┼─────────── + │ │ │ │ +targets[1, 1][1]: ───swap_1_1───────────────×(y)─────────┼─────────────────────┼─────────── + │ │ │ +targets[2, 0][0]: ───swap_2_0────────────────────────────×(x)──────────────────×(y)──────── + │ │ │ +targets[2, 0][1]: ───swap_2_0────────────────────────────×(x)──────────────────×(y)──────── + │ │ +targets[2, 1][0]: ───swap_2_1────────────────────────────×(y)────────────────────────────── + │ │ +targets[2, 1][1]: ───swap_2_1────────────────────────────×(y)────────────────────────────── + └──────────────────┘""", ) @@ -144,9 +145,11 @@ def test_swap_with_zero_classically(): data = np.array([131, 255, 92, 2]) swz = SwapWithZero(selection_bitsizes=2, target_bitsize=8, n_target_registers=4) - for sel in range(2**2): - sel, out_data = swz.call_classically(selection=sel, targets=data) # type: ignore[assignment] - print(sel, out_data) + for sel_in in range(2**2): + sel_out, out_data = swz.call_classically(selection=sel_in, targets=data) # type: ignore[assignment] + assert sel_in == sel_out + assert isinstance(out_data, np.ndarray) + assert out_data[0] == data[sel_in] @pytest.mark.parametrize( 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 c149bdfb2..1e4b37e61 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 @@ -147,6 +148,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)} diff --git a/qualtran/cirq_interop/_cirq_to_bloq_test.py b/qualtran/cirq_interop/_cirq_to_bloq_test.py index 64a7c3087..c4b62fd31 100644 --- a/qualtran/cirq_interop/_cirq_to_bloq_test.py +++ b/qualtran/cirq_interop/_cirq_to_bloq_test.py @@ -162,9 +162,11 @@ def signature(self) -> Signature: cbloq = cirq_optree_to_cbloq(circuit) assert cbloq.signature == qualtran.Signature([qualtran.Register('qubits', QBit(), shape=(28,))]) bloq_instances = [binst for binst, _, _ in cbloq.iter_bloqnections()] - assert all(bloq_instances[i].bloq == Join(QAny(2)) for i in range(14)) - assert bloq_instances[14].bloq == CirqGateWithRegisters(reg1) - assert bloq_instances[14].bloq.signature == qualtran.Signature( + # Greedy iteration of iter_bloqnections first joins only qubits needed + # for the first gate. + assert all(bloq_instances[i].bloq == Join(QAny(2)) for i in range(12)) + assert bloq_instances[12].bloq == CirqGateWithRegisters(reg1) + assert bloq_instances[12].bloq.signature == qualtran.Signature( [qualtran.Register('x', QAny(bitsize=2), shape=(3, 4))] ) assert bloq_instances[15].bloq == CirqGateWithRegisters(anc_reg) diff --git a/qualtran/drawing/qpic_diagram_test.py b/qualtran/drawing/qpic_diagram_test.py index 9e9c42d7e..456f375f1 100644 --- a/qualtran/drawing/qpic_diagram_test.py +++ b/qualtran/drawing/qpic_diagram_test.py @@ -39,46 +39,3 @@ def test_qpic_data_for_reflect_using_prepare(): selection G:width=17:shape=box \textrm{\scalebox{0.8}{R\_L}} """, ) - - _assert_bloq_has_qpic_diagram( - bloq.decompose_bloq(), - r""" -DEFINE off color=white -DEFINE on color=black -_empty_wire W off -selection W \textrm{\scalebox{0.8}{selection}} -selection / \textrm{\scalebox{0.5}{BoundedQUInt(2, 4)}} -LABEL length=10 -reg W off -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: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}} -reg_1 / \textrm{\scalebox{0.5}{QAny(2)}} -reg_2:on G:width=25:shape=box \textrm{\scalebox{0.8}{alloc}} -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)}} -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: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}}""", - )