[Costs] Qubit Counts

qualtran/bloqs/util_bloqs.py

@@ -355,6 +355,9 @@ def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSym
         assert reg.name == 'reg'
         return directional_text_box('alloc', Side.RIGHT)
 
+    def __str__(self):
+        return f'Allocate({self.dtype})'
+
 
 @frozen
 class Free(_BookkeepingBloq):
@@ -402,6 +405,9 @@ def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSym
         assert reg.name == 'reg'
         return directional_text_box('free', Side.LEFT)
 
+    def __str__(self):
+        return f'Free({self.dtype})'
+
 
 @frozen
 class ArbitraryClifford(Bloq):

qualtran/resource_counting/__init__.py

@@ -31,4 +31,6 @@
 
 from ._costing import GeneralizerT, get_cost_value, get_cost_cache, query_costs, CostKey, CostValT
 
+from ._qubit_counts import QubitCount
+
 from . import generalizers

qualtran/resource_counting/_qubit_counts.py

@@ -0,0 +1,132 @@
+#  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, Union
+
+import networkx as nx
+import sympy
+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
+
+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], 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 = 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[int]):
+    """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'], int]) -> int:
+        """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) -> int:
+        """Zero cost is zero qubits."""
+        return 0
+
+    def __str__(self):
+        return 'qubit count'

qualtran/resource_counting/_qubit_counting_test.py → 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.for_testing.interior_alloc import InteriorAlloc
 from qualtran.bloqs.for_testing.with_decomposition import (
     TestIndependentParallelCombo,
     TestSerialCombo,
 )
+from qualtran.bloqs.util_bloqs import Allocate, Free
 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,19 @@ 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")

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
+}