[Costs] get_bloq_callee_counts and call graph organization (#900)

get_bloq_callee_counts and call graph organization

Author: mpharrigan

qualtran/_infra/bloq.py

@@ -322,7 +322,7 @@ def call_graph(
                 according to `keep` and `max_depth` (if provided) or if a bloq cannot be
                 decomposed.
         """
-        from qualtran.resource_counting.bloq_counts import get_bloq_call_graph
+        from qualtran.resource_counting import get_bloq_call_graph
 
         return get_bloq_call_graph(self, generalizer=generalizer, keep=keep, max_depth=max_depth)
 

qualtran/cirq_interop/t_complexity_protocol.py

@@ -150,16 +150,16 @@ def _from_iterable(it: Any) -> Optional[TComplexity]:
 
 def _from_bloq_build_call_graph(stc: Any) -> Optional[TComplexity]:
     # Uses the depth 1 call graph of Bloq `stc` to recursively compute the complexity.
+    from qualtran.resource_counting import get_bloq_callee_counts
     from qualtran.resource_counting.generalizers import cirq_to_bloqs
 
     if not isinstance(stc, Bloq):
         return None
-    _, sigma = stc.call_graph(max_depth=1, generalizer=cirq_to_bloqs)
-    if sigma == {stc: 1}:
-        # No decomposition found.
+    callee_counts = get_bloq_callee_counts(bloq=stc, generalizer=cirq_to_bloqs)
+    if len(callee_counts) == 0:
         return None
     ret = TComplexity()
-    for bloq, n in sigma.items():
+    for bloq, n in callee_counts:
         r = t_complexity(bloq)
         if r is None:
             return None

qualtran/cirq_interop/t_complexity_protocol_test.py

@@ -15,6 +15,7 @@
 
 import cirq
 import pytest
+from attrs import frozen
 
 from qualtran import Bloq, GateWithRegisters, Signature
 from qualtran._infra.gate_with_registers import get_named_qubits
@@ -36,6 +37,7 @@ class DoesNotSupportTComplexity:
     ...
 
 
+@frozen
 class SupportsTComplexityGateWithRegisters(GateWithRegisters):
     @property
     def signature(self) -> Signature:
@@ -64,6 +66,7 @@ def signature(self) -> 'Signature':
         return Signature.build(q=1)
 
 
+@frozen
 class SupportsTComplexityBloqViaBuildCallGraph(Bloq):
     @property
     def signature(self) -> 'Signature':
@@ -75,6 +78,8 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
 
 def test_t_complexity_for_bloq_via_build_call_graph():
     bloq = SupportsTComplexityBloqViaBuildCallGraph()
+    _, sigma = bloq.call_graph(max_depth=1)
+    assert sigma != {}
     assert t_complexity(bloq) == TComplexity(t=5, clifford=10)
 
 

qualtran/drawing/bloq_counts_graph_test.py

@@ -44,11 +44,12 @@ def test_format_counts_graph_markdown():
     ret = format_counts_graph_markdown(graph)
     assert (
         ret
-        == r""" - `MultiAnd(cvs=(1, 1, 1, 1, 1, 1))`
-   - `And(cv1=1, cv2=1, uncompute=False)`: $\displaystyle 5$
+        == """\
+ - `MultiAnd(cvs=(1, 1, 1, 1, 1, 1))`
+   - `And(cv1=1, cv2=1, uncompute=False)`: $\\displaystyle 5$
  - `And(cv1=1, cv2=1, uncompute=False)`
-   - `ArbitraryClifford(n=2)`: $\displaystyle 9$
-   - `TGate()`: $\displaystyle 4$
+   - `ArbitraryClifford(n=2)`: $\\displaystyle 9$
+   - `TGate()`: $\\displaystyle 4$
 """
     )
 

qualtran/drawing/flame_graph.py

@@ -24,7 +24,7 @@
 import sympy
 
 from qualtran import Bloq
-from qualtran.resource_counting.bloq_counts import _compute_sigma
+from qualtran.resource_counting._call_graph import _compute_sigma
 from qualtran.resource_counting.t_counts_from_sigma import t_counts_from_sigma
 
 

qualtran/resource_counting/__init__.py

@@ -17,11 +17,13 @@
 isort:skip_file
 """
 
-from .bloq_counts import (
+from ._generalization import GeneralizerT
+
+from ._call_graph import (
     BloqCountT,
-    GeneralizerT,
     big_O,
     SympySymbolAllocator,
+    get_bloq_callee_counts,
     get_bloq_call_graph,
     print_counts_graph,
     build_cbloq_call_graph,

qualtran/resource_counting/bloq_counts.py qualtran/resource_counting/_call_graph.pyqualtran/resource_counting/bloq_counts.py renamed to qualtran/resource_counting/_call_graph.py

@@ -23,7 +23,7 @@
 from qualtran import Bloq, CompositeBloq, DecomposeNotImplementedError, DecomposeTypeError
 
 BloqCountT = Tuple[Bloq, Union[int, sympy.Expr]]
-GeneralizerT = Callable[[Bloq], Optional[Bloq]]
+from ._generalization import _make_composite_generalizer, GeneralizerT
 
 
 def big_O(expr) -> sympy.Order:
@@ -85,6 +85,38 @@ def _generalize_callees(
     return callee_counts
 
 
+def get_bloq_callee_counts(
+    bloq: 'Bloq', generalizer: 'GeneralizerT' = None, ssa: SympySymbolAllocator = None
+) -> List[BloqCountT]:
+    """Get the direct callees of a bloq and the number of times they are called.
+
+    This calls `bloq.build_call_graph()` with the correct configuration options.
+
+    Args:
+        bloq: The bloq.
+        generalizer: If provided, run this function on each callee to consolidate attributes
+            that do not affect resource estimates. If the callable
+            returns `None`, the bloq is omitted from the counts graph. If a sequence of
+            generalizers is provided, each generalizer will be run in order.
+        ssa: A sympy symbol allocator that can be provided if one already exists in your
+            computation.
+
+    Returns:
+        A list of (bloq, n) bloq counts.
+    """
+    if generalizer is None:
+        generalizer = lambda b: b
+    if isinstance(generalizer, (list, tuple)):
+        generalizer = _make_composite_generalizer(*generalizer)
+    if ssa is None:
+        ssa = SympySymbolAllocator()
+
+    try:
+        return _generalize_callees(bloq.build_call_graph(ssa), generalizer)
+    except (DecomposeNotImplementedError, DecomposeTypeError):
+        return []
+
+
 def _build_call_graph(
     bloq: Bloq,
     generalizer: GeneralizerT,
@@ -103,8 +135,7 @@ def _build_call_graph(
         # We already visited this node.
         return
 
-    # Make sure this node is present in the graph. You could annotate
-    # additional node properties here, too.
+    # Make sure this node is present in the graph.
     g.add_node(bloq)
 
     # Base case 1: This node is requested by the user to be a leaf node via the `keep` parameter.
@@ -116,12 +147,7 @@ def _build_call_graph(
         return
 
     # Prep for recursion: get the callees and modify them according to `generalizer`.
-    try:
-        callee_counts = _generalize_callees(bloq.build_call_graph(ssa), generalizer)
-    except (DecomposeNotImplementedError, DecomposeTypeError):
-        # Base case 3: Decomposition (or `bloq_counts`) is not implemented. This is left as a
-        #              leaf node.
-        return
+    callee_counts = get_bloq_callee_counts(bloq, generalizer)
 
     # Base case 3: Empty list of callees
     if not callee_counts:
@@ -165,19 +191,6 @@ def _compute_sigma(root_bloq: Bloq, g: nx.DiGraph) -> Dict[Bloq, Union[int, symp
     return dict(bloq_sigmas[root_bloq])
 
 
-def _make_composite_generalizer(*funcs: GeneralizerT) -> GeneralizerT:
-    """Return a generalizer that calls each `*funcs` generalizers in order."""
-
-    def _composite_generalize(b: Bloq) -> Optional[Bloq]:
-        for func in funcs:
-            b = func(b)
-            if b is None:
-                return
-        return b
-
-    return _composite_generalize
-
-
 def get_bloq_call_graph(
     bloq: Bloq,
     generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,

…an/resource_counting/bloq_counts_test.py …an/resource_counting/_call_graph_test.pyqualtran/resource_counting/bloq_counts_test.py renamed to qualtran/resource_counting/_call_graph_test.py qualtran/resource_counting/bloq_counts_test.py renamed to qualtran/resource_counting/_call_graph_test.py

@@ -25,7 +25,12 @@
 from qualtran import Bloq, BloqBuilder, Signature, SoquetT
 from qualtran.bloqs.basic_gates import TGate
 from qualtran.bloqs.util_bloqs import ArbitraryClifford, Join, Split
-from qualtran.resource_counting import BloqCountT, get_bloq_call_graph, SympySymbolAllocator
+from qualtran.resource_counting import (
+    BloqCountT,
+    get_bloq_call_graph,
+    get_bloq_callee_counts,
+    SympySymbolAllocator,
+)
 
 
 @frozen
@@ -88,6 +93,23 @@ def test_bloq_counts_method():
     assert str(expr) == '3*log(100)'
 
 
+def test_get_bloq_callee_counts():
+    bloq = BigBloq(100)
+    callee_counts = get_bloq_callee_counts(bloq)
+    assert callee_counts == [(SubBloq(unrelated_param=0.5), sympy.log(100))]
+
+    bloq = DecompBloq(10)
+    callee_counts = get_bloq_callee_counts(bloq)
+    assert len(callee_counts) == 10 + 2  # 2 for split/join
+
+    bloq = SubBloq(unrelated_param=0.5)
+    callee_counts = get_bloq_callee_counts(bloq)
+    assert callee_counts == [(TGate(), 3)]
+
+    callee_counts = get_bloq_callee_counts(TGate())
+    assert callee_counts == []
+
+
 def test_bloq_counts_decomp():
     graph, sigma = get_bloq_call_graph(DecompBloq(10))
     assert len(sigma) == 3  # includes split and join
@@ -107,7 +129,7 @@ def generalize(bloq):
 
 @pytest.mark.notebook
 def test_notebook():
-    qlt_testing.execute_notebook('bloq_counts')
+    qlt_testing.execute_notebook('call_graph')
 
 
 def _to_tuple(x: Iterable[BloqCountT]) -> Sequence[BloqCountT]:

qualtran/resource_counting/_generalization.py

@@ -0,0 +1,32 @@
+#  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, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from qualtran import Bloq
+
+GeneralizerT = Callable[['Bloq'], Optional['Bloq']]
+
+
+def _make_composite_generalizer(*funcs: 'GeneralizerT') -> 'GeneralizerT':
+    """Return a generalizer that calls each `*funcs` generalizers in order."""
+
+    def _composite_generalize(b: 'Bloq') -> Optional['Bloq']:
+        for func in funcs:
+            b = func(b)
+            if b is None:
+                return
+        return b
+
+    return _composite_generalize

qualtran/resource_counting/_generalization_test.py

@@ -0,0 +1,46 @@
+#  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 Optional
+
+from qualtran import Bloq
+from qualtran.bloqs.for_testing import TestAtom
+from qualtran.resource_counting._generalization import _make_composite_generalizer
+
+
+def test_make_composite_generalizer():
+    def func1(b: Bloq) -> Optional[Bloq]:
+        if isinstance(b, TestAtom):
+            return TestAtom()
+        return b
+
+    def func2(b: Bloq) -> Optional[Bloq]:
+        if isinstance(b, TestAtom):
+            return
+        return b
+
+    b = TestAtom(tag='test')
+    assert func1(b) == TestAtom()
+    assert func2(b) is None
+
+    g00 = _make_composite_generalizer()
+    g10 = _make_composite_generalizer(func1)
+    g01 = _make_composite_generalizer(func2)
+    g11 = _make_composite_generalizer(func1, func2)
+    g11_r = _make_composite_generalizer(func2, func1)
+
+    assert g00(b) == b
+    assert g10(b) == TestAtom()
+    assert g01(b) is None
+    assert g11(b) is None
+    assert g11_r(b) is None