From 6c6483593ee3bf7431113ffbf2ee0fa781c7afde Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Wed, 24 Apr 2024 11:56:02 -0700
Subject: [PATCH 1/4] bloq counts

---
 qualtran/resource_counting/__init__.py       |   2 +
 qualtran/resource_counting/_bloq_counts.py   | 219 +++++++++++++++++++
 qualtran/resource_counting/classify_bloqs.py |  16 ++
 qualtran/resource_counting/generalizers.py   |  14 +-
 4 files changed, 240 insertions(+), 11 deletions(-)
 create mode 100644 qualtran/resource_counting/_bloq_counts.py

diff --git a/qualtran/resource_counting/__init__.py b/qualtran/resource_counting/__init__.py
index 0f9d1885f2..c3914b5dd7 100644
--- a/qualtran/resource_counting/__init__.py
+++ b/qualtran/resource_counting/__init__.py
@@ -31,4 +31,6 @@
 
 from ._costing import GeneralizerT, get_cost_value, get_cost_cache, query_costs, CostKey, CostValT
 
+from ._bloq_counts import BloqCount, QECGatesCost
+
 from . import generalizers
diff --git a/qualtran/resource_counting/_bloq_counts.py b/qualtran/resource_counting/_bloq_counts.py
new file mode 100644
index 0000000000..64714113fb
--- /dev/null
+++ b/qualtran/resource_counting/_bloq_counts.py
@@ -0,0 +1,219 @@
+#  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 collections import defaultdict
+from typing import Callable, Dict, Optional, Tuple, TYPE_CHECKING
+
+import attrs
+import networkx as nx
+from attrs import frozen
+
+from ._call_graph import get_bloq_callee_counts
+from ._costing import CostKey
+from .classify_bloqs import bloq_is_clifford
+
+if TYPE_CHECKING:
+    from qualtran import Bloq
+
+logger = logging.getLogger(__name__)
+
+BloqCountDict = Dict['Bloq', int]
+
+
+@frozen
+class BloqCount(CostKey[BloqCountDict]):
+    gateset_bloqs: Tuple['Bloq', ...]
+    gateset_name: str
+
+    @classmethod
+    def for_gateset(cls, gateset_name: str):
+        from qualtran.bloqs.basic_gates import TGate, Toffoli, TwoBitCSwap
+
+        if gateset_name == 't':
+            bloqs = (TGate(), TGate(is_adjoint=True))
+        elif gateset_name == 't+tof':
+            bloqs = (TGate(), TGate(is_adjoint=True), Toffoli())
+        elif gateset_name == 't+tof+cswap':
+            bloqs = (TGate(), TGate(is_adjoint=True), Toffoli(), TwoBitCSwap())
+        else:
+            raise ValueError(f"Unknown gateset name {gateset_name}")
+
+        return cls(bloqs, gateset_name=gateset_name)
+
+    @classmethod
+    def for_call_graph_leaf_bloqs(cls, g: nx.DiGraph):
+        leaf_bloqs = {node for node in g.nodes if not g.succ[node]}
+        return cls(tuple(leaf_bloqs), gateset_name='leaf')
+
+    def compute(
+        self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], BloqCountDict]
+    ) -> BloqCountDict:
+        if bloq in self.gateset_bloqs:
+            logger.info("Computing %s: %s is in the target gateset.", self, bloq)
+            return {bloq: 1}
+
+        totals: BloqCountDict = defaultdict(lambda: 0)
+        callees = get_bloq_callee_counts(bloq)
+        logger.info("Computing %s for %s from %d callee(s)", self, bloq, len(callees))
+        for callee, n_times_called in callees:
+            callee_cost = get_callee_cost(callee)
+            for gateset_bloq, count in callee_cost.items():
+                totals[gateset_bloq] += n_times_called * count
+
+        return dict(totals)
+
+    def zero(self) -> BloqCountDict:
+        return {}
+
+    def __str__(self):
+        return f'{self.gateset_name} counts'
+
+
+@frozen(kw_only=True)
+class GateCounts:
+    t: int = 0
+    toffoli: int = 0
+    cswap: int = 0
+    and_bloq: int = 0
+    clifford: int = 0
+
+    def __add__(self, other):
+        if not isinstance(other, GateCounts):
+            raise TypeError(f"Can only add other `GateCounts` objects, not {self}")
+
+        return GateCounts(
+            t=self.t + other.t,
+            toffoli=self.toffoli + other.toffoli,
+            cswap=self.cswap + other.cswap,
+            and_bloq=self.and_bloq + other.and_bloq,
+            clifford=self.clifford + other.clifford,
+        )
+
+    def __mul__(self, other):
+        if not isinstance(other, int):
+            raise TypeError(f"Can only multiply `GateCounts` objects by integers, not {self}")
+
+        return GateCounts(
+            t=other * self.t,
+            toffoli=other * self.toffoli,
+            cswap=other * self.cswap,
+            and_bloq=other * self.and_bloq,
+            clifford=other * self.clifford,
+        )
+
+    def __rmul__(self, other):
+        return self.__mul__(other)
+
+    def __str__(self):
+        strs = []
+        for f in attrs.fields(self.__class__):
+            val = getattr(self, f.name)
+            if val != 0:
+                strs.append(f'{f.name}: {val}')
+
+        if strs:
+            return ', '.join(strs)
+        return '-'
+
+    @property
+    def total_n_magic(self):
+        """The total number of magic states.
+
+        This can be used as a rough proxy for total cost. It is the sum of all the attributes
+        other than `clifford`.
+        """
+        return self.t + self.toffoli + self.cswap + self.and_bloq
+
+    @property
+    def total_n_magic(self):
+        """The total number of magic states.
+
+        This can be used as a rough proxy for total cost. It is the sum of all the attributes
+        other than `clifford`.
+        """
+        return self.t + self.toffoli + self.cswap + self.and_bloq
+
+
+@frozen
+class QECGatesCost(CostKey[GateCounts]):
+    """Counts specifically for 'expensive' gates in a surface code error correction scheme."""
+
+    ts_per_toffoli: Optional[int] = None
+    toffolis_per_and: Optional[int] = None
+    ts_per_and: Optional[int] = None
+    toffolis_per_cswap: Optional[int] = None
+    ts_per_cswap: Optional[int] = None
+
+    def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], GateCounts]) -> GateCounts:
+        from qualtran.bloqs.basic_gates import TGate, Toffoli, TwoBitCSwap
+        from qualtran.bloqs.mcmt.and_bloq import And
+
+        # T gates
+        if isinstance(bloq, TGate):
+            return GateCounts(t=1)
+
+        # Toffolis
+        if isinstance(bloq, Toffoli):
+            if self.ts_per_toffoli is not None:
+                return GateCounts(t=self.ts_per_toffoli)
+            else:
+                return GateCounts(toffoli=1)
+
+        # 'And' bloqs
+        if isinstance(bloq, And) and not bloq.uncompute:
+            if self.toffolis_per_and is not None:
+                return GateCounts(toffoli=self.toffolis_per_and * self.ts_per_toffoli)
+            elif self.ts_per_and is not None:
+                return GateCounts(t=self.ts_per_and)
+            else:
+                return GateCounts(and_bloq=1)
+
+        # CSwaps aka Fredkin
+        if isinstance(bloq, TwoBitCSwap):
+            if self.toffolis_per_cswap is not None:
+                return GateCounts(toffoli=self.toffolis_per_cswap)
+            elif self.ts_per_cswap is not None:
+                return GateCounts(t=self.ts_per_cswap)
+            else:
+                return GateCounts(cswap=1)
+
+        # Cliffords
+        if bloq_is_clifford(bloq):
+            return GateCounts(clifford=1)
+
+        # Recursive case
+        totals = GateCounts()
+        callees = get_bloq_callee_counts(bloq)
+        logger.info("Computing %s for %s from %d callee(s)", self, bloq, len(callees))
+        for callee, n_times_called in callees:
+            callee_cost = get_callee_cost(callee)
+            totals += n_times_called * callee_cost
+        return totals
+
+    def zero(self) -> GateCounts:
+        return GateCounts()
+
+    def validate_val(self, val: GateCounts):
+        if not isinstance(val, GateCounts):
+            raise TypeError(f"{self} values should be `GateCounts`, got {val}")
+
+    def __str__(self):
+        gates = ['t']
+        if self.ts_per_toffoli is None:
+            gates.append('tof')
+        if self.toffolis_per_and is None and self.ts_per_and is None:
+            gates.append('and')
+        if self.toffolis_per_cswap is None and self.ts_per_cswap is None:
+            gates.append('cswap')
+        return ','.join(gates) + ' counts'
diff --git a/qualtran/resource_counting/classify_bloqs.py b/qualtran/resource_counting/classify_bloqs.py
index c88c818f28..924c2b772e 100644
--- a/qualtran/resource_counting/classify_bloqs.py
+++ b/qualtran/resource_counting/classify_bloqs.py
@@ -105,3 +105,19 @@ def classify_t_count_by_bloq_type(
         classification = classify_bloq(k, bloq_classification)
         classified_bloqs[classification] += v * t_counts_from_sigma(k.call_graph()[1])
     return classified_bloqs
+
+
+def bloq_is_clifford(b: Bloq):
+    from qualtran.bloqs.basic_gates import CNOT, Hadamard, SGate, TwoBitSwap, XGate, ZGate
+    from qualtran.bloqs.mcmt.multi_control_multi_target_pauli import MultiTargetCNOT
+    from qualtran.bloqs.util_bloqs import ArbitraryClifford
+
+    if isinstance(b, Adjoint):
+        b = b.subbloq
+
+    if isinstance(
+        b, (TwoBitSwap, Hadamard, XGate, ZGate, ArbitraryClifford, CNOT, MultiTargetCNOT, SGate)
+    ):
+        return True
+
+    return False
diff --git a/qualtran/resource_counting/generalizers.py b/qualtran/resource_counting/generalizers.py
index 0e164039b5..4ca451d09b 100644
--- a/qualtran/resource_counting/generalizers.py
+++ b/qualtran/resource_counting/generalizers.py
@@ -24,7 +24,7 @@
 import attrs
 import sympy
 
-from qualtran import Adjoint, Bloq
+from qualtran import Bloq
 
 PHI = sympy.Symbol(r'\phi')
 CV = sympy.Symbol("cv")
@@ -79,18 +79,10 @@ def generalize_cvs(b: Bloq) -> Optional[Bloq]:
 
 def ignore_cliffords(b: Bloq) -> Optional[Bloq]:
     """A generalizer that ignores known clifford bloqs."""
-    from qualtran.bloqs.basic_gates import CNOT, Hadamard, SGate, TwoBitSwap, XGate, ZGate
-    from qualtran.bloqs.mcmt.multi_control_multi_target_pauli import MultiTargetCNOT
-    from qualtran.bloqs.util_bloqs import ArbitraryClifford
+    from qualtran.resource_counting.classify_bloqs import bloq_is_clifford
 
-    if isinstance(b, Adjoint):
-        b = b.subbloq
-
-    if isinstance(
-        b, (TwoBitSwap, Hadamard, XGate, ZGate, ArbitraryClifford, CNOT, MultiTargetCNOT, SGate)
-    ):
+    if bloq_is_clifford(b):
         return None
-
     return b
 
 

From 6a1d49d30817ae149661e34fbbae4294fecd9b6a Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Wed, 24 Apr 2024 11:56:13 -0700
Subject: [PATCH 2/4] qubit counts

---
 qualtran/resource_counting/__init__.py        |   1 +
 qualtran/resource_counting/_qubit_counting.py |  61 --------
 qualtran/resource_counting/_qubit_counts.py   | 133 ++++++++++++++++++
 ...counting_test.py => _qubit_counts_test.py} |  19 ++-
 qualtran/resource_counting/qubit_counts.ipynb | 117 +++++++++++++++
 5 files changed, 269 insertions(+), 62 deletions(-)
 delete mode 100644 qualtran/resource_counting/_qubit_counting.py
 create mode 100644 qualtran/resource_counting/_qubit_counts.py
 rename qualtran/resource_counting/{_qubit_counting_test.py => _qubit_counts_test.py} (72%)
 create mode 100644 qualtran/resource_counting/qubit_counts.ipynb

diff --git a/qualtran/resource_counting/__init__.py b/qualtran/resource_counting/__init__.py
index c3914b5dd7..228b3d1932 100644
--- a/qualtran/resource_counting/__init__.py
+++ b/qualtran/resource_counting/__init__.py
@@ -32,5 +32,6 @@
 from ._costing import GeneralizerT, get_cost_value, get_cost_cache, query_costs, CostKey, CostValT
 
 from ._bloq_counts import BloqCount, QECGatesCost
+from ._qubit_counts import QubitCount
 
 from . import generalizers
diff --git a/qualtran/resource_counting/_qubit_counting.py b/qualtran/resource_counting/_qubit_counting.py
deleted file mode 100644
index eb3e5ebf32..0000000000
--- a/qualtran/resource_counting/_qubit_counting.py
+++ /dev/null
@@ -1,61 +0,0 @@
-#  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, Set, Union
-
-import networkx as nx
-import sympy
-
-from qualtran import Bloq, Connection, DanglingT
-from qualtran._infra.composite_bloq import _binst_to_cxns
-
-
-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 = sympy.Max(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 = sympy.Max(max_width, after_size)
-
-    return max_width
diff --git a/qualtran/resource_counting/_qubit_counts.py b/qualtran/resource_counting/_qubit_counts.py
new file mode 100644
index 0000000000..6f5f8ff3f0
--- /dev/null
+++ b/qualtran/resource_counting/_qubit_counts.py
@@ -0,0 +1,133 @@
+#  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 abc
+import logging
+from typing import Callable, Dict, Generic, Sequence, Set, Tuple, TYPE_CHECKING, 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 ._call_graph import get_bloq_callee_counts
+from ._costing import CostKey
+from .symbolic_counting_utils import smax
+
+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'
diff --git a/qualtran/resource_counting/_qubit_counting_test.py b/qualtran/resource_counting/_qubit_counts_test.py
similarity index 72%
rename from qualtran/resource_counting/_qubit_counting_test.py
rename to qualtran/resource_counting/_qubit_counts_test.py
index 47c9a83011..7ea9314d8f 100644
--- a/qualtran/resource_counting/_qubit_counting_test.py
+++ b/qualtran/resource_counting/_qubit_counts_test.py
@@ -14,13 +14,18 @@
 
 import sympy
 
+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 +59,15 @@ 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,
+    }
diff --git a/qualtran/resource_counting/qubit_counts.ipynb b/qualtran/resource_counting/qubit_counts.ipynb
new file mode 100644
index 0000000000..092e61a711
--- /dev/null
+++ b/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.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

From 234c020d9188828e51615fcda62fdd69ae69cb6e Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Wed, 24 Apr 2024 11:56:20 -0700
Subject: [PATCH 3/4] success prob

---
 qualtran/resource_counting/__init__.py        |  1 +
 qualtran/resource_counting/_success_prob.py   | 45 +++++++++++++++++++
 .../resource_counting/_success_prob_test.py   | 26 +++++++++++
 3 files changed, 72 insertions(+)
 create mode 100644 qualtran/resource_counting/_success_prob.py
 create mode 100644 qualtran/resource_counting/_success_prob_test.py

diff --git a/qualtran/resource_counting/__init__.py b/qualtran/resource_counting/__init__.py
index 228b3d1932..8f13b7cbd5 100644
--- a/qualtran/resource_counting/__init__.py
+++ b/qualtran/resource_counting/__init__.py
@@ -33,5 +33,6 @@
 
 from ._bloq_counts import BloqCount, QECGatesCost
 from ._qubit_counts import QubitCount
+from ._success_prob import SuccessProb
 
 from . import generalizers
diff --git a/qualtran/resource_counting/_success_prob.py b/qualtran/resource_counting/_success_prob.py
new file mode 100644
index 0000000000..24c6826772
--- /dev/null
+++ b/qualtran/resource_counting/_success_prob.py
@@ -0,0 +1,45 @@
+#  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 abc
+import logging
+from typing import Callable, Dict, Generic, Sequence, Tuple, TYPE_CHECKING
+
+import sympy
+from attrs import frozen
+
+from qualtran import Bloq, DecomposeNotImplementedError, DecomposeTypeError
+
+from . import CostValT
+from ._call_graph import get_bloq_callee_counts
+from ._costing import CostKey
+
+logger = logging.getLogger(__name__)
+
+
+@frozen
+class SuccessProb(CostKey[float]):
+    def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], float]) -> float:
+        tot: float = 1.0
+        callees = get_bloq_callee_counts(bloq)
+        logger.info("Computing %s for %s from %d callee(s)", self, bloq, len(callees))
+        for callee, n in callees:
+            v = get_callee_cost(callee)
+            tot *= v**n
+        return tot
+
+    def zero(self) -> CostValT:
+        return 1.0  # under multiplication, 1 is the identity.
+
+    def __str__(self):
+        return 'success prob'
diff --git a/qualtran/resource_counting/_success_prob_test.py b/qualtran/resource_counting/_success_prob_test.py
new file mode 100644
index 0000000000..8b1fdef54a
--- /dev/null
+++ b/qualtran/resource_counting/_success_prob_test.py
@@ -0,0 +1,26 @@
+#  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 qualtran.bloqs.for_testing.costing import CostingBloq
+from qualtran.resource_counting import get_cost_cache, get_cost_value, SuccessProb
+
+
+def test_coin_flip():
+    flip = CostingBloq('CoinFlip', num_qubits=1, static_costs=[(SuccessProb(), 0.5)])
+    algo = CostingBloq('Algo', num_qubits=0, callees=[(flip, 4)])
+
+    p = get_cost_value(algo, SuccessProb())
+    assert p == 0.5**4
+
+    costs = get_cost_cache(algo, SuccessProb())
+    assert costs == {algo: p, flip: 0.5}

From d82ba062b917b656f65793139cecea4d5482455d Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Thu, 25 Apr 2024 10:56:48 -0700
Subject: [PATCH 4/4] finale

---
 qualtran/_infra/registers.py                  |   4 +-
 qualtran/_infra/registers_test.py             |   9 +
 qualtran/bloqs/basic_gates/hadamard.py        |   3 +
 qualtran/bloqs/basic_gates/on_each_test.py    |   3 +-
 qualtran/bloqs/basic_gates/t_gate.py          |   4 +-
 qualtran/bloqs/basic_gates/toffoli.py         |   3 +
 .../bloqs/chemistry/resource_estimation.ipynb |  42 +++-
 qualtran/bloqs/data_loading/qrom.py           |   8 +
 .../bloqs/data_loading/select_swap_qrom.py    |   3 +
 qualtran/bloqs/for_testing/costing_test.py    |  10 +-
 qualtran/bloqs/mcmt/and_bloq.py               |   7 +
 .../multiplexers/selected_majorana_fermion.py |   8 +
 qualtran/drawing/bloq_counts_graph_test.py    |   8 +-
 qualtran/drawing/flame_graph.py               |   2 +-
 qualtran/resource_counting/_call_graph.py     |  29 +--
 qualtran/resource_counting/costs.ipynb        | 200 ++++++++++++++++++
 16 files changed, 306 insertions(+), 37 deletions(-)
 create mode 100644 qualtran/resource_counting/costs.ipynb

diff --git a/qualtran/_infra/registers.py b/qualtran/_infra/registers.py
index 101e57a57e..c8152c3bd4 100644
--- a/qualtran/_infra/registers.py
+++ b/qualtran/_infra/registers.py
@@ -203,9 +203,11 @@ def n_qubits(self) -> int:
         is taken to be the greater of the number of left or right qubits. A bloq with this
         signature uses at least this many qubits.
         """
+        from qualtran.resource_counting.symbolic_counting_utils import smax
+
         left_size = sum(reg.total_bits() for reg in self.lefts())
         right_size = sum(reg.total_bits() for reg in self.rights())
-        return max(left_size, right_size)
+        return smax(left_size, right_size)
 
     def __repr__(self):
         return f'Signature({repr(self._registers)})'
diff --git a/qualtran/_infra/registers_test.py b/qualtran/_infra/registers_test.py
index cb4c8380a2..9bb55d81f0 100644
--- a/qualtran/_infra/registers_test.py
+++ b/qualtran/_infra/registers_test.py
@@ -205,3 +205,12 @@ def test_is_symbolic():
     assert is_symbolic(r)
     r = Register("my_reg", QAny(2), shape=sympy.symbols("x y"))
     assert is_symbolic(r)
+
+
+def test_symbolic_reg():
+    n = sympy.Symbol('n', positive=True, integer=True)
+    sig = Signature(
+        [Register('x', QAny(n), side=Side.LEFT), Register('y', QAny(2 * n), side=Side.RIGHT)]
+    )
+
+    assert sig.n_qubits() == 2 * n
diff --git a/qualtran/bloqs/basic_gates/hadamard.py b/qualtran/bloqs/basic_gates/hadamard.py
index acb7a7234f..6a834d7803 100644
--- a/qualtran/bloqs/basic_gates/hadamard.py
+++ b/qualtran/bloqs/basic_gates/hadamard.py
@@ -93,6 +93,9 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
             return Text('')
         return TextBox('H')
 
+    def __str__(self):
+        return 'H'
+
 
 @bloq_example
 def _hadamard() -> Hadamard:
diff --git a/qualtran/bloqs/basic_gates/on_each_test.py b/qualtran/bloqs/basic_gates/on_each_test.py
index 6782a24ad9..7ae02d5013 100644
--- a/qualtran/bloqs/basic_gates/on_each_test.py
+++ b/qualtran/bloqs/basic_gates/on_each_test.py
@@ -42,8 +42,7 @@ def test_classical_simulation():
     h_on_each = OnEach(10, Hadamard())
     with pytest.raises(
         NotImplementedError,
-        match=r'.*does not support classical simulation: '
-        r'Hadamard\(\) is not classically simulable\.',
+        match=r'.*does not support classical simulation: ' r'H is not classically simulable\.',
     ):
         h_on_each.call_classically(q=0)
 
diff --git a/qualtran/bloqs/basic_gates/t_gate.py b/qualtran/bloqs/basic_gates/t_gate.py
index 928488f05a..6a55fa2b85 100644
--- a/qualtran/bloqs/basic_gates/t_gate.py
+++ b/qualtran/bloqs/basic_gates/t_gate.py
@@ -108,8 +108,8 @@ def pretty_name(self) -> str:
         return f'T{maybe_dag}'
 
     def __str__(self):
-        maybe_dag = 'is_adjoint=True' if self.is_adjoint else ''
-        return f'TGate({maybe_dag})'
+        maybe_dag = '†' if self.is_adjoint else ''
+        return f'T{maybe_dag}'
 
     def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         if reg is None:
diff --git a/qualtran/bloqs/basic_gates/toffoli.py b/qualtran/bloqs/basic_gates/toffoli.py
index 3c8599c316..4559e31b43 100644
--- a/qualtran/bloqs/basic_gates/toffoli.py
+++ b/qualtran/bloqs/basic_gates/toffoli.py
@@ -125,6 +125,9 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
             return ModPlus()
         raise ValueError(f'Unknown wire symbol register name: {reg.name}')
 
+    def __str__(self):
+        return 'Toffoli'
+
 
 @bloq_example
 def _toffoli() -> Toffoli:
diff --git a/qualtran/bloqs/chemistry/resource_estimation.ipynb b/qualtran/bloqs/chemistry/resource_estimation.ipynb
index 3518c879fa..6c86049b44 100644
--- a/qualtran/bloqs/chemistry/resource_estimation.ipynb
+++ b/qualtran/bloqs/chemistry/resource_estimation.ipynb
@@ -133,7 +133,6 @@
     "from qualtran.drawing.musical_score import get_musical_score_data, draw_musical_score\n",
     "msd = get_musical_score_data(block_encoding_bloq.decompose_bloq())\n",
     "fig, ax = draw_musical_score(msd)\n",
-    "plt.tick_params(left=False, right=False, labelleft=False, labelbottom=False, bottom=False)\n",
     "fig.set_size_inches(8, 4)"
    ]
   },
@@ -185,6 +184,47 @@
     "print(f'qualtran = {num_toff} vs. ref = 10880, delta = {num_toff - 10880}')"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e79d3c99-cd23-4333-a177-6d6ab3dca72a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# qualtran = 26749.0 vs. ref = 10880, delta = 15869.0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c61a4b30-b875-4414-b198-e08774df0c4a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting import BloqCount, query_costs, get_cost_value, QECGatesCost\n",
+    "from qualtran.resource_counting.generalizers import ignore_alloc_free, ignore_split_join, generalize_rotation_angle"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a126c934-1528-425a-aa4d-93a4bb880236",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "get_cost_value(block_encoding_bloq, BloqCount.for_gateset(\"t+tof+cswap\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e68450ff-d582-400f-abd1-f3d24dd43979",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "46976/4 + 30480/4 + 7105 + 280"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "dbd1615f",
diff --git a/qualtran/bloqs/data_loading/qrom.py b/qualtran/bloqs/data_loading/qrom.py
index 4a7388431b..58971d6161 100644
--- a/qualtran/bloqs/data_loading/qrom.py
+++ b/qualtran/bloqs/data_loading/qrom.py
@@ -327,6 +327,11 @@ def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT
         targets = {k: v ^ vals[k] for k, v in targets.items()}
         return controls | selections | targets
 
+    def my_static_costs(self, cost_key: 'CostKey') -> Union[Any, NotImplemented]:
+        if cost_key == QubitCount():
+            return self.num_controls + 2 * sum(self.selection_bitsizes) + sum(self.target_bitsizes)
+        return super().my_static_costs(cost_key)
+
     def _circuit_diagram_info_(self, args) -> cirq.CircuitDiagramInfo:
         from qualtran.cirq_interop._bloq_to_cirq import _wire_symbol_to_cirq_diagram_info
 
@@ -374,6 +379,9 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         n_cnot = prod(*self.target_bitsizes, *self.data_shape)
         return {(And(), n_and), (And().adjoint(), n_and), (CNOT(), n_cnot)}
 
+    def __str__(self):
+        return 'QROM'
+
 
 @bloq_example
 def _qrom_small() -> QROM:
diff --git a/qualtran/bloqs/data_loading/select_swap_qrom.py b/qualtran/bloqs/data_loading/select_swap_qrom.py
index 67008785f2..69762c90e5 100644
--- a/qualtran/bloqs/data_loading/select_swap_qrom.py
+++ b/qualtran/bloqs/data_loading/select_swap_qrom.py
@@ -259,3 +259,6 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
 
     def _value_equality_values_(self):
         return self.block_size, self._target_bitsizes, self.data
+
+    def __str__(self):
+        return 'SelectSwapQROM'
diff --git a/qualtran/bloqs/for_testing/costing_test.py b/qualtran/bloqs/for_testing/costing_test.py
index fb8340b746..b4972a4adb 100644
--- a/qualtran/bloqs/for_testing/costing_test.py
+++ b/qualtran/bloqs/for_testing/costing_test.py
@@ -24,9 +24,9 @@ def test_costing_bloqs():
         == """\
 Algo -- 1 -> Func1
 Algo -- 1 -> Func2
-Func1 -- 10 -> Hadamard()
-Func1 -- 10 -> TGate()
-Func1 -- 10 -> TGate(is_adjoint=True)
-Func2 -- 100 -> Toffoli()
-Toffoli() -- 4 -> TGate()"""
+Func1 -- 10 -> H
+Func1 -- 10 -> T
+Func1 -- 10 -> T†
+Func2 -- 100 -> Toffoli
+Toffoli -- 4 -> T"""
     )
diff --git a/qualtran/bloqs/mcmt/and_bloq.py b/qualtran/bloqs/mcmt/and_bloq.py
index 95dd44b575..56a4bf6ef3 100644
--- a/qualtran/bloqs/mcmt/and_bloq.py
+++ b/qualtran/bloqs/mcmt/and_bloq.py
@@ -235,6 +235,13 @@ def _t_complexity_(self) -> TComplexity:
         else:
             return TComplexity(t=4 * 1, clifford=9 + 2 * pre_post_cliffords)
 
+    def __str__(self):
+        dag = '†' if self.uncompute else ''
+
+        if self.cv1 == 0 or self.cv2 == 0:
+            return f'And{dag}_{self.cv1}{self.cv2}'
+        return f'And{dag}'
+
 
 @bloq_example(
     generalizer=[cirq_to_bloqs, ignore_cliffords, ignore_alloc_free, generalize_rotation_angle]
diff --git a/qualtran/bloqs/multiplexers/selected_majorana_fermion.py b/qualtran/bloqs/multiplexers/selected_majorana_fermion.py
index a659adf4e0..ae17882839 100644
--- a/qualtran/bloqs/multiplexers/selected_majorana_fermion.py
+++ b/qualtran/bloqs/multiplexers/selected_majorana_fermion.py
@@ -25,6 +25,7 @@
 from qualtran._infra.data_types import BoundedQUInt
 from qualtran._infra.gate_with_registers import total_bits
 from qualtran.bloqs.multiplexers.unary_iteration_bloq import UnaryIterationGate
+from qualtran.resource_counting import CostKey
 
 
 @attrs.frozen
@@ -134,3 +135,10 @@ def nth_operation(  # type: ignore[override]
         yield cirq.CNOT(control, *accumulator)
         yield self.target_gate(target[target_idx]).controlled_by(control)
         yield cirq.CZ(*accumulator, target[target_idx])
+
+    def my_static_costs(self, cost_key: 'CostKey') -> Union[Any, NotImplemented]:
+        from qualtran.resource_counting import QubitCount
+
+        if isinstance(cost_key, QubitCount):
+            return self.signature.n_qubits()
+        return super().my_static_costs(cost_key)
diff --git a/qualtran/drawing/bloq_counts_graph_test.py b/qualtran/drawing/bloq_counts_graph_test.py
index c46d725ed4..466160dbbc 100644
--- a/qualtran/drawing/bloq_counts_graph_test.py
+++ b/qualtran/drawing/bloq_counts_graph_test.py
@@ -35,7 +35,7 @@ def test_format_counts_sigma():
         == """\
 #### Counts totals:
  - `ArbitraryClifford(n=2)`: 45
- - `TGate()`: 20"""
+ - `T`: 20"""
     )
 
 
@@ -46,10 +46,10 @@ def test_format_counts_graph_markdown():
         ret
         == """\
  - `MultiAnd(cvs=(1, 1, 1, 1, 1, 1))`
-   - `And(cv1=1, cv2=1, uncompute=False)`: $\\displaystyle 5$
- - `And(cv1=1, cv2=1, uncompute=False)`
+   - `And`: $\\displaystyle 5$
+ - `And`
    - `ArbitraryClifford(n=2)`: $\\displaystyle 9$
-   - `TGate()`: $\\displaystyle 4$
+   - `T`: $\\displaystyle 4$
 """
     )
 
diff --git a/qualtran/drawing/flame_graph.py b/qualtran/drawing/flame_graph.py
index f0dbf2495a..27efcdf275 100644
--- a/qualtran/drawing/flame_graph.py
+++ b/qualtran/drawing/flame_graph.py
@@ -58,7 +58,7 @@ def _pretty_name(bloq: Bloq) -> str:
 
 @functools.lru_cache(maxsize=1024)
 def _t_counts_for_bloq(bloq: Bloq, graph: nx.DiGraph) -> Union[int, sympy.Expr]:
-    sigma = _compute_sigma(bloq, graph)
+    sigma = _compute_sigma(bloq, graph, generalizer=lambda b: b)
     return t_counts_from_sigma(sigma)
 
 
diff --git a/qualtran/resource_counting/_call_graph.py b/qualtran/resource_counting/_call_graph.py
index 0e6017a971..625092a1df 100644
--- a/qualtran/resource_counting/_call_graph.py
+++ b/qualtran/resource_counting/_call_graph.py
@@ -171,27 +171,14 @@ def _build_call_graph(
             g.add_edge(bloq, callee, n=n)
 
 
-def _compute_sigma(root_bloq: Bloq, g: nx.DiGraph) -> Dict[Bloq, Union[int, sympy.Expr]]:
-    """Iterate over nodes to sum up the counts of leaf bloqs."""
-    bloq_sigmas: Dict[Bloq, Dict[Bloq, Union[int, sympy.Expr]]] = defaultdict(
-        lambda: defaultdict(lambda: 0)
-    )
-    for bloq in reversed(list(nx.topological_sort(g))):
-        callees = list(g.successors(bloq))
-        sigma = bloq_sigmas[bloq]
-        if not callees:
-            # 1. `bloq` is a leaf node. Its count is one of itself.
-            sigma[bloq] = 1
-            continue
-
-        for callee in callees:
-            callee_sigma = bloq_sigmas[callee]
-            # 2. Otherwise, sigma of the caller is sum(n * sigma of callee) for all the callees.
-            n = g.edges[bloq, callee]['n']
-            for k in callee_sigma.keys():
-                sigma[k] += callee_sigma[k] * n
+def _compute_sigma(
+    root_bloq: Bloq, g: nx.DiGraph, generalizer: 'GeneralizerT'
+) -> Dict[Bloq, Union[int, sympy.Expr]]:
+    """Shim for compatibility with old 'sigma' that used the call graph to count leaf bloqs."""
+    from qualtran.resource_counting import BloqCount, get_cost_value
 
-    return dict(bloq_sigmas[root_bloq])
+    leaf_counts = BloqCount.for_call_graph_leaf_bloqs(g)
+    return get_cost_value(root_bloq, leaf_counts, generalizer=generalizer)
 
 
 def get_bloq_call_graph(
@@ -239,7 +226,7 @@ def get_bloq_call_graph(
     if bloq is None:
         raise ValueError("You can't generalize away the root bloq.")
     _build_call_graph(bloq, generalizer, ssa, keep, max_depth, g=g, depth=0)
-    sigma = _compute_sigma(bloq, g)
+    sigma = _compute_sigma(bloq, g, generalizer)
     return g, sigma
 
 
diff --git a/qualtran/resource_counting/costs.ipynb b/qualtran/resource_counting/costs.ipynb
new file mode 100644
index 0000000000..c1de1e6c76
--- /dev/null
+++ b/qualtran/resource_counting/costs.ipynb
@@ -0,0 +1,200 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d4d7bfac-b2c7-414e-921e-92da858fe216",
+   "metadata": {},
+   "source": [
+    "# Costs\n",
+    "\n",
+    "We cannot (yet) run the large ('fault-tolerant scale') quantum algorithms expressed in Qualtran. Instead, a major research goal is to estimate the resources required to implement interesting quantum algorithms. Resources can be time, qubits, T gates, precision, depth, or any number of things. Throughout Qualtran, we use the less ambiguous term \"costs\" to represent these quantities.\n",
+    "\n",
+    "In the research literature, there is a surface level agreement on what costs to count: Often gate counts of T, Toffoli, Clifford, and Rotation; as well as the number of qubits are tabulated as the most probable limiting factors of implementing these algorithms on real quantum computers. When it comes to details, agreement on costs is not guaranteed. Researchers may want to consider additional target gatesets or architectures; disagree on how to lump or count cliffords of various sizes; and treat rotations in a variety of ways.\n",
+    "\n",
+    "In Qualtran, we provide a set of configurable `CostKey`s that can be used to query algorithms expressed as bloqs. Developers can also implement their own costs by overriding `CostKey` to provide even more customization."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "42306c58-584f-4b85-89e7-9a64a1398a64",
+   "metadata": {},
+   "source": [
+    "## Getting costs\n",
+    "\n",
+    "The following functions can be used to query costs for a bloq. Each takes a bloq and a cost key:\n",
+    "\n",
+    " - `get_cost_value`\n",
+    " - `get_cost_cache`\n",
+    "\n",
+    "The former will return a single value whereas the latter will return a dictionary mapping every bloq for which a cost was computed to its value.\n",
+    "\n",
+    "Before we start, we'll create an example bloq with enough complexity to be interesting:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0e4801d3-2e55-4c0a-a113-5d6c9263f0e4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Construct a PrepTHC bloq for demonstration purposes.\n",
+    "\n",
+    "import numpy as np\n",
+    "from qualtran.bloqs.chemistry.thc import PrepareTHC\n",
+    "\n",
+    "\n",
+    "num_spin_orb = 25\n",
+    "num_mu = 100\n",
+    "t_l = np.random.normal(0, 1, size=num_spin_orb//2)\n",
+    "zeta = np.random.normal(0, 1, size=(num_mu, num_mu))\n",
+    "zeta = 0.5 * (zeta + zeta.T)\n",
+    "\n",
+    "prep_thc = PrepareTHC.from_hamiltonian_coeffs(t_l, zeta, num_bits_state_prep=10)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b0023aa8-7cfe-4f22-8625-e6bbe509c721",
+   "metadata": {},
+   "source": [
+    "`get_cost_value` takes a `Bloq` and a `CostKey` and returns a value. The type of value depends on the particular cost key. For example, if we ask for the qubit count by using the `QubitCount` cost key, we get an integer."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "600e6e48-e888-4021-aecc-3fbc6b26ab6a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting import get_cost_value, QubitCount\n",
+    "\n",
+    "get_cost_value(prep_thc, QubitCount())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "94920338-1145-4f81-8b30-c306c44f470a",
+   "metadata": {},
+   "source": [
+    "`get_cost_cache` has the same input arguments, but instead returns a dictionary containing the (sub-)costs for each bloq that was encountered during the recursive computation. In the following demo, we only print the top 5 entries to avoid overwhelming the output."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c077411d-8f4f-4aa1-b5a4-c3e8fca90e69",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting import get_cost_cache\n",
+    "\n",
+    "cost_cache = get_cost_cache(prep_thc, QubitCount())\n",
+    "print(f'cost_cache contains {len(cost_cache)} entries. Displaying top five.\\n')\n",
+    "\n",
+    "top_five = sorted(cost_cache.items(), key=lambda x: x[1])[-5:]\n",
+    "for bloq, val in top_five:\n",
+    "    print(f'{bloq}: {val}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dd7c069d-3820-4a1b-95b4-326dbda87fce",
+   "metadata": {},
+   "source": [
+    "## Configurable Cost Keys\n",
+    "\n",
+    "The behavior of a cost computation can be modified by arguments to the cost key. The `QubitCount()` cost is simple: it has no parameters and the value of the cost is a simple integer. `QECGatesCost` returns a richer cost value type and has more options for controlling the cost query."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1963547-c294-4a0f-bb57-5e2239968eb4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting import QECGatesCost\n",
+    "\n",
+    "get_cost_value(prep_thc, QECGatesCost())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "68bfe8db-e5f2-41ac-bf7d-26f38022eb67",
+   "metadata": {},
+   "source": [
+    "Here, we can see that the output is a `GateCounts` data class that tallys up the number of T, Toffoli, CSwap (aka Fredkin), AND, and Clifford Gates. We can provide additional configuration to count things differently."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9ff68d5f-5a78-4c46-ac69-eb738c083af5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "t_cost = QECGatesCost(ts_per_toffoli=4, ts_per_cswap=7, ts_per_and=4)\n",
+    "get_cost_value(prep_thc, QECGatesCost(ts_per_toffoli=4, ts_per_cswap=7, ts_per_and=4))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ecfbf912-fae9-4dd9-a6d3-a8bf1532e084",
+   "metadata": {},
+   "source": [
+    "Now, the result is entirely in terms of T gate counts."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a375c726-4a97-4640-9ced-d7a1bcb700a4",
+   "metadata": {},
+   "source": [
+    "## Querying Multiple Costs\n",
+    "\n",
+    "The `query_costs` function will get multiple costs for multiple bloqs simultaneously. Its output is suitable for annotating a call graph diagram"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "20a95e27-9b69-459f-9f39-56a98ec3ad67",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting import query_costs\n",
+    "\n",
+    "from qualtran.drawing import GraphvizCallGraph\n",
+    "from qualtran.resource_counting.generalizers import ignore_split_join, ignore_alloc_free\n",
+    "\n",
+    "costs = query_costs(prep_thc, [t_cost, QubitCount()])\n",
+    "g, _ = prep_thc.call_graph(max_depth=3, generalizer=[ignore_split_join, ignore_alloc_free])\n",
+    "\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.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}