costing framework

Author: mpharrigan

qualtran/_infra/bloq.py

@@ -37,7 +37,7 @@
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.cirq_interop.t_complexity_protocol import TComplexity
     from qualtran.drawing import WireSymbol
-    from qualtran.resource_counting import BloqCountT, GeneralizerT, SympySymbolAllocator
+    from qualtran.resource_counting import BloqCountT, CostKey, GeneralizerT, SympySymbolAllocator
     from qualtran.simulation.classical_sim import ClassicalValT
 
 
@@ -294,6 +294,20 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         """
         return self.decompose_bloq().build_call_graph(ssa)
 
+    def my_static_costs(self, cost_key: 'CostKey') -> Union[Any, NotImplemented]:
+        """Override this method to provide static costs.
+
+        The system will query a particular cost by asking for a `cost_key`. This method
+        can optionally provide a value, which will be preferred over a computed cost.
+
+        Static costs can be provided if the particular cost cannot be easily computed or
+        as a performance optimization.
+
+        This method must return `NotImplemented` if a value cannot be provided for the specified
+        CostKey.
+        """
+        return NotImplemented
+
     def call_graph(
         self,
         generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,

qualtran/bloqs/for_testing/costing.py

@@ -0,0 +1,56 @@
+#  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 Any, Sequence, Set, Tuple, Union
+
+from attrs import field, frozen
+
+from qualtran import Bloq, Signature
+from qualtran.resource_counting import BloqCountT, CostKey, SympySymbolAllocator
+
+
+@frozen
+class CostingBloq(Bloq):
+    """A bloq that lets you set the costs via attributes."""
+
+    name: str
+    num_qubits: int
+    callees: Sequence[BloqCountT] = field(converter=tuple, factory=tuple)
+    static_costs: Sequence[Tuple[CostKey, Any]] = field(converter=tuple, factory=tuple)
+
+    @property
+    def signature(self) -> 'Signature':
+        return Signature.build(register=self.num_qubits)
+
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+        return set(self.callees)
+
+    def my_static_costs(self, k: 'CostKey') -> Union[Any, NotImplemented]:
+        return dict(self.static_costs).get(k, NotImplemented)
+
+    def pretty_name(self):
+        return self.name
+
+    def __str__(self):
+        return self.name
+
+
+def make_example_costing_bloqs():
+    from qualtran.bloqs.basic_gates import Hadamard, TGate, Toffoli
+
+    func1 = CostingBloq(
+        'Func1', num_qubits=10, callees=[(TGate(), 10), (TGate().adjoint(), 10), (Hadamard(), 10)]
+    )
+    func2 = CostingBloq('Func2', num_qubits=3, callees=[(Toffoli(), 100)])
+    algo = CostingBloq('Algo', num_qubits=100, callees=[(func1, 1), (func2, 1)])
+    return algo

qualtran/bloqs/for_testing/costing_test.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 qualtran.bloqs.for_testing.costing import make_example_costing_bloqs
+from qualtran.resource_counting import format_call_graph_debug_text
+
+
+def test_costing_bloqs():
+    algo = make_example_costing_bloqs()
+    g, _ = algo.call_graph()
+    assert (
+        format_call_graph_debug_text(g)
+        == """\
+Algo -- 1 -> Func1
+Algo -- 1 -> Func2
+Func1 -- 10 -> Hadamard()
+Func1 -- 10 -> TGate()
+Func1 -- 10 -> TGate(is_adjoint=True)
+Func2 -- 100 -> Toffoli()
+Toffoli() -- 4 -> TGate()"""
+    )

qualtran/resource_counting/__init__.py

@@ -25,8 +25,10 @@
     SympySymbolAllocator,
     get_bloq_callee_counts,
     get_bloq_call_graph,
-    print_counts_graph,
     build_cbloq_call_graph,
+    format_call_graph_debug_text,
 )
 
+from ._costing import GeneralizerT, get_cost_value, get_cost_cache, query_costs, CostKey, CostValT
+
 from . import generalizers

qualtran/resource_counting/_call_graph.py

@@ -240,8 +240,11 @@ def get_bloq_call_graph(
     return g, sigma
 
 
-def print_counts_graph(g: nx.DiGraph):
+def format_call_graph_debug_text(g: nx.DiGraph) -> str:
     """Print the graph returned from `get_bloq_counts_graph`."""
-    for b in nx.topological_sort(g):
-        for succ in g.succ[b]:
-            print(b, '--', g.edges[b, succ]['n'], '->', succ)
+    lines = []
+    for gen in nx.topological_generations(g):
+        for b in sorted(gen, key=str):
+            for succ in sorted(g.succ[b], key=str):
+                lines.append(f"{b} -- {g.edges[b, succ]['n']} -> {succ}")
+    return '\n'.join(lines)

qualtran/resource_counting/_costing.py

@@ -0,0 +1,197 @@
+#  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
+import time
+from collections import defaultdict
+from typing import (
+    Callable,
+    Dict,
+    Generic,
+    Iterable,
+    Optional,
+    Sequence,
+    Tuple,
+    TYPE_CHECKING,
+    TypeVar,
+    Union,
+)
+
+from ._generalization import _make_composite_generalizer, GeneralizerT
+
+if TYPE_CHECKING:
+    from qualtran import Bloq
+
+logger = logging.getLogger(__name__)
+
+CostValT = TypeVar('CostValT')
+
+
+class CostKey(Generic[CostValT], metaclass=abc.ABCMeta):
+    """Abstract base class for different types of costs.
+
+    One important aspect of a bloq is the resources required to execute it on an error
+    corrected quantum computer. Since we're usually trying to minimize these resource requirements
+    we will generally use the catch-all term "costs".
+
+    There are a variety of different types or flavors of costs. Each is represented by an
+    instance of a sublcass of `CostKey`. For example, gate counts (including T-gate counts),
+    qubit requirements, and circuit depth are all cost metrics that may be of interest.
+
+    Each `CostKey` primarily encodes the behavior required to compute a cost value from a
+    bloq. Often, these costs are defined recursively: a bloq's costs is some combination
+    of the costs of the bloqs in its decomposition (i.e. the bloq 'callees'). Implementors
+    must override the `compute` method to define the cost computation.
+
+    Each cost key has an associated CostValT. For example, the CostValT of a "t count"
+    CostKey could be an integer. For a more complicated gateset, the value could be a mapping
+    from gate to count. This abstract base class is generic w.r.t. `CostValT`. Subclasses
+    should have a concrete value type. The `validate_val` method can optionally be overridden
+    to raise an exception if a bad value type is encountered. The `zero` method must return
+    the zero (additive identity) cost value of the correct type.
+    """
+
+    @abc.abstractmethod
+    def compute(self, bloq: 'Bloq', get_callee_cost: Callable[['Bloq'], CostValT]) -> CostValT:
+        """Compute this type of cost.
+
+        When implementing a new CostKey, this method must be overridden.
+        Users should not call this method directly. Instead: use the `qualtran.resource_counting`
+        functions like `get_cost_value`, `get_cost_cache`, or `query_costs`. These provide
+        caching, logging, generalizers, and support for static costs.
+
+        For recursive computations, use the provided callable to recurse.
+
+        Args:
+            bloq: The bloq to compute the cost of.
+            get_callee_cost: A qualtran-provided function for computing costs for "callees"
+                of the bloq; i.e. bloqs in the decomposition. Use this function to accurately
+                cache intermediate cost values and respect bloqs' static costs.
+
+        Returns:
+            A value of the generic type `CostValT`. Subclasses should define their value type.
+        """
+
+    @abc.abstractmethod
+    def zero(self) -> CostValT:
+        """The value corresponding to zero cost."""
+
+    def validate_val(self, val: CostValT):
+        """Assert that `val` is a valid `CostValT`.
+
+        This method can be optionally overridden to raise an error if an invalid value
+        is encountered. By default, no validation is performed.
+        """
+
+
+def _get_cost_value(
+    bloq: 'Bloq',
+    cost_key: CostKey[CostValT],
+    *,
+    costs_cache: Dict['Bloq', CostValT],
+    generalizer: 'GeneralizerT',
+) -> CostValT:
+    """Helper function for `query_costs`.
+
+    This function tries the following strategies
+     1. Use the value found in `costs_cache`, if it exists.
+     2. Use the value returned by `Bloq.my_static_costs` if one is returned.
+     3. Use `cost_key.compute()` and cache the result in `costs_cache`.
+
+    Args:
+        bloq: The bloq.
+        cost_key: The cost key to get the value for.
+        costs_cache: A dictionary to use as a cache for computed bloq costs. This cache
+            will be mutated by this function.
+        generalizer: The generalizer to operate on each bloq before computing its cost.
+    """
+    bloq = generalizer(bloq)
+    if bloq is None:
+        return cost_key.zero()
+
+    # Strategy 1: Use cached value
+    if bloq in costs_cache:
+        logger.debug("Using cached %s for %s", cost_key, bloq)
+        return costs_cache[bloq]
+
+    # Strategy 2: Static costs
+    static_cost = bloq.my_static_costs(cost_key)
+    if static_cost is not NotImplemented:
+        cost_key.validate_val(static_cost)
+        logger.info("Using static %s for %s", cost_key, bloq)
+        costs_cache[bloq] = static_cost
+        return static_cost
+
+    # Strategy 3: Compute
+    # part a. set up caching of computed costs. Using the callable will use the cache if possible
+    #         and only recurse if the bloq has not been seen before. The result of a computation
+    #         will be cached.
+    def _get_cost_val_internal(callee: 'Bloq'):
+        return _get_cost_value(callee, cost_key, costs_cache=costs_cache, generalizer=generalizer)
+
+    tstart = time.perf_counter()
+    computed_cost = cost_key.compute(bloq, _get_cost_val_internal)
+    tdur = time.perf_counter() - tstart
+    logger.info("Computed %s for %s in %g s", cost_key, bloq, tdur)
+    costs_cache[bloq] = computed_cost
+    return computed_cost
+
+
+def get_cost_value(
+    bloq: 'Bloq',
+    cost_key: CostKey[CostValT],
+    costs_cache: Optional[Dict['Bloq', CostValT]] = None,
+    generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
+) -> CostValT:
+    if costs_cache is None:
+        costs_cache = {}
+    if generalizer is None:
+        generalizer = lambda b: b
+    if isinstance(generalizer, (list, tuple)):
+        generalizer = _make_composite_generalizer(*generalizer)
+
+    cost_val = _get_cost_value(bloq, cost_key, costs_cache=costs_cache, generalizer=generalizer)
+    return cost_val
+
+
+def get_cost_cache(
+    bloq: 'Bloq',
+    cost_key: CostKey[CostValT],
+    costs_cache: Optional[Dict['Bloq', CostValT]] = None,
+    generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
+) -> Dict['Bloq', CostValT]:
+    if costs_cache is None:
+        costs_cache = {}
+    if generalizer is None:
+        generalizer = lambda b: b
+    if isinstance(generalizer, (list, tuple)):
+        generalizer = _make_composite_generalizer(*generalizer)
+
+    _get_cost_value(bloq, cost_key, costs_cache=costs_cache, generalizer=generalizer)
+    return costs_cache
+
+
+def query_costs(
+    bloq: 'Bloq',
+    cost_keys: Iterable[CostKey],
+    generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
+) -> Dict['Bloq', Dict[CostKey, CostValT]]:
+
+    costs = defaultdict(dict)
+    for cost_key in cost_keys:
+        cost_for_bloqs = get_cost_cache(bloq, cost_key, generalizer=generalizer)
+        for bloq, val in cost_for_bloqs.items():
+            costs[bloq][cost_key] = val
+    return dict(costs)