Support CompositeBloq in QubitCount() (#1379)

Author: mpharrigan

qualtran/_infra/composite_bloq.py

@@ -232,7 +232,10 @@ def as_composite_bloq(self) -> 'CompositeBloq':
         return self
 
     def decompose_bloq(self) -> 'CompositeBloq':
-        raise DecomposeTypeError("CompositeBloq cannot be decomposed.")
+        raise ValueError(
+            "Calling `decompose_bloq` on a CompositeBloq is ill-defined. "
+            "Consider using the composite bloq directly or using `.flatten()`."
+        )
 
     def build_call_graph(self, ssa: Optional['SympySymbolAllocator']) -> Set['BloqCountT']:
         """Return the bloq counts by counting up all the subbloqs."""

qualtran/resource_counting/_bloq_counts_test.py

@@ -18,7 +18,7 @@
 from qualtran.bloqs.basic_gates import Hadamard, TGate, Toffoli
 from qualtran.bloqs.basic_gates._shims import Measure
 from qualtran.bloqs.for_testing.costing import make_example_costing_bloqs
-from qualtran.bloqs.mcmt import MultiTargetCNOT
+from qualtran.bloqs.mcmt import MultiAnd, MultiTargetCNOT
 from qualtran.cirq_interop.t_complexity_protocol import TComplexity
 from qualtran.resource_counting import BloqCount, GateCounts, get_cost_value, QECGatesCost
 
@@ -68,6 +68,12 @@ def test_qec_gates_cost():
     assert gc == GateCounts(toffoli=100, t=2 * 2 * 10, clifford=2 * 10)
 
 
+def test_qec_gates_cost_cbloq():
+    bloq = MultiAnd(cvs=(1,) * 5)
+    cbloq = bloq.decompose_bloq()
+    assert get_cost_value(bloq, QECGatesCost()) == get_cost_value(cbloq, QECGatesCost())
+
+
 @pytest.mark.parametrize(
     ['bloq', 'counts'],
     [

qualtran/resource_counting/_qubit_counts.py

@@ -19,7 +19,7 @@
 from attrs import frozen
 
 from qualtran import Bloq, Connection, DanglingT, DecomposeNotImplementedError, DecomposeTypeError
-from qualtran._infra.composite_bloq import _binst_to_cxns
+from qualtran._infra.composite_bloq import _binst_to_cxns, CompositeBloq
 from qualtran.symbolics import smax, SymbolicInt
 
 from ._call_graph import get_bloq_callee_counts
@@ -104,6 +104,9 @@ def compute(
         # Most accurate:
         # Compute the number of qubits ("width") from the bloq's decomposition. We forward
         # the `get_callee_cost` function so this can recurse into subbloqs.
+        if isinstance(bloq, CompositeBloq):
+            logger.info("Computing %s by the passed-in CompositeBloq", self)
+            return _cbloq_max_width(bloq._binst_graph, get_callee_cost)
         try:
             cbloq = bloq.decompose_bloq()
             logger.info("Computing %s for %s from its decomposition", self, bloq)

qualtran/resource_counting/_qubit_counts_test.py

@@ -25,15 +25,14 @@
     TestSerialCombo,
 )
 from qualtran.drawing import show_bloq
-from qualtran.resource_counting import get_cost_cache, QubitCount
+from qualtran.resource_counting import get_cost_cache, get_cost_value, 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():
     n = sympy.Symbol('n', positive=True)
     bloq = InteriorAlloc(n=n)
-    show_bloq(bloq.decompose_bloq())
 
     binst_graph = bloq.decompose_bloq()._binst_graph
     max_width = _cbloq_max_width(binst_graph)
@@ -43,7 +42,6 @@ def test_max_width_interior_alloc_symb():
 def test_max_width_interior_alloc_nums():
     n = 10
     bloq = InteriorAlloc(n=n)
-    show_bloq(bloq.decompose_bloq())
 
     binst_graph = bloq.decompose_bloq()._binst_graph
     max_width = _cbloq_max_width(binst_graph)
@@ -75,6 +73,14 @@ def test_qubit_count_cost():
     }
 
 
+def test_on_cbloq():
+    n = sympy.Symbol('n', positive=True, integer=True)
+    bloq = InteriorAlloc(n=n)
+    cbloq = bloq.decompose_bloq()
+    n_qubits = get_cost_value(cbloq, QubitCount())
+    assert n_qubits == 3 * n
+
+
 @pytest.mark.notebook
 def test_notebook():
     qlt_testing.execute_notebook("qubit_counts")

qualtran/resource_counting/qubit_counts.ipynb

@@ -52,6 +52,8 @@
    "id": "58f0823f-f76f-4adb-8f2a-a25d1e2ee070",
    "metadata": {},
    "source": [
+    "## Example: a bloq with an interior allocation\n",
+    "\n",
     "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. "
    ]
   },
@@ -91,6 +93,25 @@
     "costs = query_costs(bloq, [QubitCount()])\n",
     "GraphvizCallGraph(g, costs).get_svg()"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e421621b-cbb4-4603-acab-2b65f8cbf08c",
+   "metadata": {},
+   "source": [
+    "You can call `get_cost_value` on a composite bloq directly."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a7666cc8-4023-4c15-a946-acbaae3443d0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cbloq = bloq.decompose_bloq()\n",
+    "get_cost_value(cbloq, QubitCount())"
+   ]
   }
  ],
  "metadata": {