diff --git a/qualtran/resource_counting/__init__.py b/qualtran/resource_counting/__init__.py index 8288931a6c..9f8503627a 100644 --- a/qualtran/resource_counting/__init__.py +++ b/qualtran/resource_counting/__init__.py @@ -32,5 +32,6 @@ from ._costing import GeneralizerT, get_cost_value, get_cost_cache, query_costs, CostKey, CostValT from ._success_prob import SuccessProb +from ._qubit_counts import QubitCount from . import generalizers diff --git a/qualtran/resource_counting/_qubit_counting.py b/qualtran/resource_counting/_qubit_counting.py deleted file mode 100644 index eb3e5ebf32..0000000000 --- a/qualtran/resource_counting/_qubit_counting.py +++ /dev/null @@ -1,61 +0,0 @@ -# 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. - -from typing import Callable, Set, Union - -import networkx as nx -import sympy - -from qualtran import Bloq, Connection, DanglingT -from qualtran._infra.composite_bloq import _binst_to_cxns - - -def _cbloq_max_width( - binst_graph: nx.DiGraph, _bloq_max_width: Callable[[Bloq], int] = lambda b: 0 -) -> Union[int, sympy.Expr]: - """Get the maximum width of a composite bloq. - - Specifically, we treat each binst in series. The width at each inter-bloq time point - is the sum of the bitsizes of all the connections that are "in play". The width at each - during-a-binst time point is the sum of the binst width (which is provided by the - `_bloq_max_width` callable) and the bystander connections that are "in play". The max - width is the maximum over all the time points. - - If the dataflow graph has more than one connected component, we treat each component - independently. - """ - max_width: Union[int, sympy.Expr] = 0 - in_play: Set[Connection] = set() - - for cc in nx.weakly_connected_components(binst_graph): - for binst in nx.topological_sort(binst_graph.subgraph(cc)): - pred_cxns, succ_cxns = _binst_to_cxns(binst, binst_graph=binst_graph) - - # Remove inbound connections from those that are 'in play'. - for cxn in pred_cxns: - in_play.remove(cxn) - - if not isinstance(binst, DanglingT): - # During the application of the binst, we have "observer" connections that have - # width as well as the width from the binst itself. We consider the case where - # the bloq may have a max_width greater than the max of its left/right registers. - during_size = _bloq_max_width(binst.bloq) + sum(s.shape for s in in_play) - max_width = sympy.Max(max_width, during_size) - - # After the binst, its successor connections are 'in play'. - in_play.update(succ_cxns) - after_size = sum(s.shape for s in in_play) - max_width = sympy.Max(max_width, after_size) - - return max_width diff --git a/qualtran/resource_counting/_qubit_counts.py b/qualtran/resource_counting/_qubit_counts.py new file mode 100644 index 0000000000..c39881c059 --- /dev/null +++ b/qualtran/resource_counting/_qubit_counts.py @@ -0,0 +1,133 @@ +# 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 logging +from typing import Callable, Set + +import networkx as nx +from attrs import frozen + +from qualtran import Bloq, Connection, DanglingT, DecomposeNotImplementedError, DecomposeTypeError +from qualtran._infra.composite_bloq import _binst_to_cxns +from qualtran.symbolics import smax, SymbolicInt + +from ._call_graph import get_bloq_callee_counts +from ._costing import CostKey + +logger = logging.getLogger(__name__) + + +def _cbloq_max_width( + binst_graph: nx.DiGraph, _bloq_max_width: Callable[[Bloq], SymbolicInt] = lambda b: 0 +) -> SymbolicInt: + """Get the maximum width of a composite bloq. + + Specifically, we treat each binst in series. The width at each inter-bloq time point + is the sum of the bitsizes of all the connections that are "in play". The width at each + during-a-binst time point is the sum of the binst width (which is provided by the + `_bloq_max_width` callable) and the bystander connections that are "in play". The max + width is the maximum over all the time points. + + If the dataflow graph has more than one connected component, we treat each component + independently. + """ + max_width: SymbolicInt = 0 + in_play: Set[Connection] = set() + + for cc in nx.weakly_connected_components(binst_graph): + for binst in nx.topological_sort(binst_graph.subgraph(cc)): + pred_cxns, succ_cxns = _binst_to_cxns(binst, binst_graph=binst_graph) + + # Remove inbound connections from those that are 'in play'. + for cxn in pred_cxns: + in_play.remove(cxn) + + if not isinstance(binst, DanglingT): + # During the application of the binst, we have "observer" connections that have + # width as well as the width from the binst itself. We consider the case where + # the bloq may have a max_width greater than the max of its left/right registers. + during_size = _bloq_max_width(binst.bloq) + sum(s.shape for s in in_play) + max_width = smax(max_width, during_size) + + # After the binst, its successor connections are 'in play'. + in_play.update(succ_cxns) + after_size = sum(s.shape for s in in_play) + max_width = smax(max_width, after_size) + + return max_width + + +@frozen +class QubitCount(CostKey[SymbolicInt]): + """A cost estimating the number of qubits required to implement a bloq. + + The number of qubits is bounded from below by the number of qubits implied by the signature. + If a bloq has no callees, the size implied by the signature will be returned. Otherwise, + this CostKey will try to compute the number of qubits by inspecting the decomposition. + + In the decomposition, each (sub)bloq is considered to be executed sequentially. The "width" + of the circuit (i.e. the number of qubits) at each sequence point is the number of qubits + required by the subbloq (computed recursively) plus any "bystander" idling wires. + + This is an estimate for the number of qubits required by an algorithm. Specifically: + - Bloqs are assumed to be executed sequentially, minimizing the number of qubits potentially + at the expense of greater circuit depth or execution time. + - We do not consider "tetris-ing" subbloqs. In a decomposition, each subbloq is assumed + to be using all of its qubits for the duration of its execution. This could potentially + overestimate the total number of qubits. + + This Min-Max style estimate can provide a good balance between accuracy and scalability + of the accounting. To fully account for each qubit and manage space-vs-time trade-offs, + you must comprehensively decompose your algorithm to a `cirq.Circuit` of basic gates and + use a `cirq.QubitManager` to manage trade-offs. This may be computationally expensive for + large algorithms. + """ + + def compute( + self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], SymbolicInt] + ) -> SymbolicInt: + """Compute an estimate of the number of qubits used by `bloq`. + + See the class docstring for more information. + """ + # Base case: No callees; use the signature + min_bloq_size = bloq.signature.n_qubits() + callees = get_bloq_callee_counts(bloq) + if len(callees) == 0: + logger.info("Computing %s for %s from signature", self, bloq) + return min_bloq_size + + # Compute the number of qubits ("width") from the bloq's decomposition. We forward + # the `get_callee_cost` function so this can recurse into subbloqs. + try: + cbloq = bloq.decompose_bloq() + logger.info("Computing %s for %s from its decomposition", self, bloq) + return _cbloq_max_width(cbloq._binst_graph, get_callee_cost) + except (DecomposeNotImplementedError, DecomposeTypeError): + pass + + # No decomposition specified, but callees present. Take the simple maximum of + # all the callees' sizes. This is likely an under-estimate. + tot: int = min_bloq_size + logger.info("Computing %s for %s from %d callee(s)", self, bloq, len(callees)) + for callee, n in callees: + tot = smax(tot, get_callee_cost(callee)) + return tot + + def zero(self) -> SymbolicInt: + """Zero cost is zero qubits.""" + return 0 + + def __str__(self): + return 'qubit count' diff --git a/qualtran/resource_counting/_qubit_counting_test.py b/qualtran/resource_counting/_qubit_counts_test.py similarity index 68% rename from qualtran/resource_counting/_qubit_counting_test.py rename to qualtran/resource_counting/_qubit_counts_test.py index 47c9a83011..ded0fee783 100644 --- a/qualtran/resource_counting/_qubit_counting_test.py +++ b/qualtran/resource_counting/_qubit_counts_test.py @@ -12,15 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import sympy +import qualtran.testing as qlt_testing +from qualtran import QAny +from qualtran.bloqs.basic_gates import Swap, TwoBitSwap +from qualtran.bloqs.bookkeeping import Allocate, Free from qualtran.bloqs.for_testing.interior_alloc import InteriorAlloc from qualtran.bloqs.for_testing.with_decomposition import ( TestIndependentParallelCombo, TestSerialCombo, ) from qualtran.drawing import show_bloq -from qualtran.resource_counting._qubit_counting import _cbloq_max_width +from qualtran.resource_counting import get_cost_cache, QubitCount +from qualtran.resource_counting._qubit_counts import _cbloq_max_width +from qualtran.resource_counting.generalizers import ignore_split_join def test_max_width_interior_alloc_symb(): @@ -54,3 +61,20 @@ def test_max_width_simple(): show_bloq(TestSerialCombo().decompose_bloq()) max_width = _cbloq_max_width(TestSerialCombo().decompose_bloq()._binst_graph) assert max_width == 1 + + +def test_qubit_count_cost(): + bloq = InteriorAlloc(n=10) + qubit_counts = get_cost_cache(bloq, QubitCount(), generalizer=ignore_split_join) + assert qubit_counts == { + InteriorAlloc(n=10): 30, + Allocate(QAny(10)): 10, + Free(QAny(10)): 10, + Swap(10): 20, + TwoBitSwap(): 2, + } + + +@pytest.mark.notebook +def test_notebook(): + qlt_testing.execute_notebook("qubit_counts") diff --git a/qualtran/resource_counting/qubit_counts.ipynb b/qualtran/resource_counting/qubit_counts.ipynb new file mode 100644 index 0000000000..a7d0bee5c2 --- /dev/null +++ b/qualtran/resource_counting/qubit_counts.ipynb @@ -0,0 +1,117 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "c3a2e0b1-d7f3-4761-aaf7-9fe1b4202c63", + "metadata": {}, + "source": [ + "# Qubit Counts\n", + "\n", + "The number of qubits is an important cost for running a quantum algorithm. The provided `QubitCounts()` cost key can efficiently estimate the qubit count of even large-scale algorithms by exploiting the hierarchical structure of bloq decomposition.\n", + "\n", + "\n", + "The number of qubits is bounded from below by the number of qubits implied by the signature.\n", + "If a bloq has no callees, the size implied by the signature will be returned. Otherwise,\n", + "`QubitCounts()` will try to compute the number of qubits by inspecting the decomposition.\n", + "\n", + "In the decomposition, each (sub)bloq is considered to be executed sequentially. The \"width\"\n", + "of the circuit (i.e. the number of qubits) at each sequence point is the number of qubits\n", + "required by the subbloq (computed recursively) plus any \"bystander\" idling wires.\n", + "\n", + "This is an estimate for the number of qubits required by an algorithm. Specifically:\n", + " - Bloqs are assumed to be executed sequentially, minimizing the number of qubits potentially\n", + " at the expense of greater circuit depth or execution time.\n", + " - We do not consider \"tetris-ing\" subbloqs. In a decomposition, each subbloq is assumed\n", + " to be using all of its qubits for the duration of its execution. This could potentially\n", + " overestimate the total number of qubits.\n", + "\n", + "This Min-Max style estimate can provide a good balance between accuracy and scalability\n", + "of the accounting. To fully account for each qubit and manage space-vs-time trade-offs,\n", + "you must comprehensively decompose your algorithm to a `cirq.Circuit` of basic gates and\n", + "use a `cirq.QubitManager` to manage trade-offs. This may be computationally expensive for\n", + "large algorithms." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ddb82393-8bb0-42d1-9b84-6b3e22455f7e", + "metadata": {}, + "outputs": [], + "source": [ + "import sympy\n", + "\n", + "from qualtran.drawing import show_bloq\n", + "\n", + "from qualtran.bloqs.for_testing.interior_alloc import InteriorAlloc\n", + "from qualtran.resource_counting import get_cost_value, query_costs, QubitCount" + ] + }, + { + "cell_type": "markdown", + "id": "58f0823f-f76f-4adb-8f2a-a25d1e2ee070", + "metadata": {}, + "source": [ + "For illustrative purposes, we use a bloq that has two $n$ bit registers, but allocates an additional $n$ bit register as part of its decomposition. Looking purely at the signature, you would conclude that the bloq uses $2n$ qubits; but by looking at the decomposition we can see that at its maximum circuit width it uses $3n$ qubits. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2a5768b5-7f2e-4851-bc15-3be8a66df4f5", + "metadata": {}, + "outputs": [], + "source": [ + "n = sympy.Symbol('n', positive=True, integer=True)\n", + "bloq = InteriorAlloc(n=n)\n", + "show_bloq(bloq)\n", + "show_bloq(bloq.decompose_bloq(), 'musical_score')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "329c02d8-35b0-4ed8-881b-0bb5e6856813", + "metadata": {}, + "outputs": [], + "source": [ + "get_cost_value(bloq, QubitCount())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21a84a95-1ed6-4ce1-8c43-512621efdbdf", + "metadata": {}, + "outputs": [], + "source": [ + "from qualtran.drawing import GraphvizCallGraph\n", + "\n", + "g, _ = bloq.call_graph()\n", + "costs = query_costs(bloq, [QubitCount()])\n", + "GraphvizCallGraph(g, costs).get_svg()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}