Profiling large composite bloq construction and qubit counting timings (

quantumlib#1303)

* Large bloq for testing

* aggressive profile stuff

* rollback aggressive changes, keep good ones

* lint and type

qualtran/_infra/quantum_graph.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 """Plumbing for bloq-to-bloq `Connection`s."""
-
+from functools import cached_property
 from typing import Tuple, TYPE_CHECKING, Union
 
 from attrs import field, frozen
@@ -143,7 +143,7 @@ class Connection:
     left: Soquet
     right: Soquet
 
-    @property
+    @cached_property
     def shape(self) -> int:
         ls = self.left.reg.bitsize
         rs = self.right.reg.bitsize

qualtran/bloqs/for_testing/large_bloq.py

@@ -0,0 +1,62 @@
+#  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 Dict, List
+
+import numpy as np
+from attrs import frozen
+
+from qualtran import Bloq, BloqBuilder, QUInt, Signature, Soquet, SoquetT
+from qualtran.bloqs.for_testing.atom import TestAtom
+from qualtran.bloqs.mcmt import And
+
+
+@frozen
+class LargeBloq(Bloq):
+    n_select: int
+    n_ops: int
+
+    @property
+    def signature(self) -> 'Signature':
+        return Signature.build(select=self.n_select, target=1)
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', select, target) -> Dict[str, 'SoquetT']:
+        sel = bb.split(select)
+        ancs: List[Soquet] = [None] * self.n_select  # type: ignore
+        ancs[0] = sel[0]
+
+        cvs = QUInt(self.n_select - 1).to_bits_array(np.arange(self.n_ops) % (self.n_select - 1))
+        assert cvs.shape == (self.n_ops, self.n_select - 1)
+
+        for op_i in range(self.n_ops):
+
+            # Ladder of ands
+            for i in range(self.n_select - 1):
+                and_op = And(cv1=cvs[op_i, i], cv2=1)
+                [ancs[i], sel[i + 1]], ancs[i + 1] = bb.add(
+                    and_op, ctrl=np.array([ancs[i], sel[i + 1]])
+                )
+
+            # The placeholder op
+            ancs[-1], target = bb.add(TestAtom().controlled(), ctrl=ancs[-1], q=target)
+
+            # Un-ladder of ands
+            for i in range(self.n_select - 1 - 1, 0 - 1, -1):
+                and_op = And(cv1=cvs[op_i, i], cv2=1)
+                [ancs[i], sel[i + 1]] = bb.add(
+                    and_op.adjoint(), ctrl=np.array([ancs[i], sel[i + 1]]), target=ancs[i + 1]
+                )
+
+        sel[0] = ancs[0]
+        return {'select': bb.join(sel), 'target': target}

qualtran/resource_counting/_qubit_counts.py

@@ -101,13 +101,7 @@ def compute(
 
         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
-
+        # 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.
         try:
@@ -117,8 +111,12 @@ def compute(
         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.
+        # Fallback:
+        # Use the simple maximum of callees and of this bloq's signature. If there
+        # are no callees, this will be the number of qubits implied by the signature.
+        # In any case, this strategy is likely an under-estimate of the qubit count.
+        min_bloq_size = bloq.signature.n_qubits()
+        callees = get_bloq_callee_counts(bloq)
         tot: int = min_bloq_size
         logger.info("Computing %s for %s from %d callee(s)", self, bloq, len(callees))
         for callee, n in callees: