Use greedy_topological_sort in qubit counting (#1641)

Fixes #1636

Author: mpharrigan

qualtran/bloqs/for_testing/qubit_count_many_alloc.py

@@ -0,0 +1,104 @@
+#  Copyright 2025 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 Dict
+
+import attrs
+
+from qualtran import Bloq, BloqBuilder, Signature, Soquet, SoquetT
+from qualtran.bloqs.basic_gates import CNOT, TGate
+
+
+@attrs.frozen
+class TestManyAllocOnce(Bloq):
+    """Allocate an ancilla once, and re-use it explicitly
+
+    See qualtran.resource.counting._qubit_counts_test:test_many_alloc
+    """
+
+    n: int
+
+    @property
+    def signature(self) -> Signature:
+        return Signature.build(x=self.n)
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'Soquet') -> dict[str, 'SoquetT']:
+        x = bb.split(x)
+        anc = bb.allocate()
+        for i in range(self.n):
+            x[i], anc = bb.add(CNOT(), ctrl=x[i], target=anc)
+            anc = bb.add(TGate(), q=anc)
+            x[i], anc = bb.add(CNOT(), ctrl=x[i], target=anc)
+        bb.free(anc)
+        return {'x': bb.join(x)}
+
+
+@attrs.frozen
+class TestManyAllocMany(Bloq):
+    """Allocate a new ancilla for each pseudo-subcall.
+
+    See qualtran.resource.counting._qubit_counts_test:test_many_alloc
+    """
+
+    n: int
+
+    @property
+    def signature(self) -> Signature:
+        return Signature.build(x=self.n)
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'Soquet') -> dict[str, 'SoquetT']:
+        x = bb.split(x)
+        for i in range(self.n):
+            anc = bb.allocate()
+            x[i], anc = bb.add(CNOT(), ctrl=x[i], target=anc)
+            anc = bb.add(TGate(), q=anc)
+            x[i], anc = bb.add(CNOT(), ctrl=x[i], target=anc)
+            bb.free(anc)
+        return {'x': bb.join(x)}
+
+
+@attrs.frozen
+class _Inner(Bloq):
+    """Inner part, used by `TestManyAllocAbstracted`"""
+
+    @property
+    def signature(self) -> Signature:
+        return Signature.build(x=1)
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'Soquet') -> Dict[str, 'SoquetT']:
+        anc = bb.allocate()
+        x, anc = bb.add(CNOT(), ctrl=x, target=anc)
+        anc = bb.add(TGate(), q=anc)
+        x, anc = bb.add(CNOT(), ctrl=x, target=anc)
+        bb.free(anc)
+        return {'x': x}
+
+
+@attrs.frozen
+class TestManyAllocAbstracted(Bloq):
+    """Factor allocation into subbloq
+
+    See qualtran.resource.counting._qubit_counts_test:test_many_alloc
+    """
+
+    n: int
+
+    @property
+    def signature(self) -> Signature:
+        return Signature.build(x=self.n)
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'Soquet') -> dict[str, 'SoquetT']:
+        x = bb.split(x)
+        for i in range(self.n):
+            x[i] = bb.add(_Inner(), x=x[i])
+        return {'x': bb.join(x)}

qualtran/bloqs/for_testing/qubit_count_many_alloc_test.py

@@ -0,0 +1,26 @@
+#  Copyright 2025 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 qualtran.testing as qlt_testing
+from qualtran.bloqs.for_testing.qubit_count_many_alloc import (
+    TestManyAllocAbstracted,
+    TestManyAllocMany,
+    TestManyAllocOnce,
+)
+
+
+def test_many_alloc_validity():
+    qlt_testing.assert_valid_bloq_decomposition(TestManyAllocMany(10))
+    qlt_testing.assert_valid_bloq_decomposition(TestManyAllocOnce(10))
+    qlt_testing.assert_valid_bloq_decomposition(TestManyAllocAbstracted(10))
+    qlt_testing.assert_valid_cbloq(TestManyAllocAbstracted(10).as_composite_bloq().flatten())

qualtran/resource_counting/_qubit_counts.py

@@ -18,8 +18,16 @@
 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, CompositeBloq
+from qualtran import (
+    Bloq,
+    CompositeBloq,
+    Connection,
+    DanglingT,
+    DecomposeNotImplementedError,
+    DecomposeTypeError,
+)
+from qualtran._infra.binst_graph_iterators import greedy_topological_sort
+from qualtran._infra.composite_bloq import _binst_to_cxns
 from qualtran.symbolics import smax, SymbolicInt
 
 from ._call_graph import get_bloq_callee_counts
@@ -46,7 +54,7 @@ def _cbloq_max_width(
     in_play: Set[Connection] = set()
 
     for cc in nx.weakly_connected_components(binst_graph):
-        for binst in nx.topological_sort(binst_graph.subgraph(cc)):
+        for binst in greedy_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'.

qualtran/resource_counting/_qubit_counts_test.py

@@ -81,6 +81,20 @@ def test_on_cbloq():
     assert n_qubits == 3 * n
 
 
+def test_many_alloc():
+    from qualtran.bloqs.for_testing.qubit_count_many_alloc import (
+        TestManyAllocAbstracted,
+        TestManyAllocMany,
+        TestManyAllocOnce,
+    )
+
+    n = 10
+    for bloq in [TestManyAllocMany(n), TestManyAllocOnce(n), TestManyAllocAbstracted(n)]:
+        # These should all give n+1 despite their unique constructions
+        # https://github.com/quantumlib/Qualtran/issues/1636
+        assert get_cost_value(bloq, QubitCount()) == 11
+
+
 @pytest.mark.notebook
 def test_notebook():
     qlt_testing.execute_notebook("qubit_counts")