[Costs] Framework (#913)

* costing framework

* docstrings and linting

* lint

* python 3.10 type issues

* lint

* part b

* copyright

* mypy

* lint

Author: mpharrigan

qualtran/_infra/bloq.py

@@ -40,7 +40,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
 
 
@@ -296,6 +296,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'):
+        """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,61 @@
+#  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 Any, Sequence, Set, Tuple
+
+from attrs import field, frozen
+
+from qualtran import Bloq, Signature
+from qualtran.resource_counting import BloqCountT, CostKey, SympySymbolAllocator
+
+
+def _convert_callees(callees: Sequence[BloqCountT]) -> Tuple[BloqCountT, ...]:
+    # Convert to tuples in a type-checked way.
+    return tuple(callees)
+
+
+@frozen
+class CostingBloq(Bloq):
+    """A bloq that lets you set the costs via attributes."""
+
+    name: str
+    num_qubits: int
+    callees: Sequence[BloqCountT] = field(converter=_convert_callees, 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, cost_key: 'CostKey'):
+        return dict(self.static_costs).get(cost_key, 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 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 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/bloqs/phase_estimation/lp_resource_state.py

@@ -20,7 +20,7 @@
 import cirq
 import numpy as np
 import sympy
-from numpy._typing import NDArray
+from numpy.typing import NDArray
 
 from qualtran import (
     Bloq,

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

@@ -14,7 +14,7 @@
 
 """Functionality for the `Bloq.call_graph()` protocol."""
 
-import collections.abc as abc
+import collections.abc
 from collections import defaultdict
 from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 
@@ -231,7 +231,7 @@ def get_bloq_call_graph(
         keep = lambda b: False
     if generalizer is None:
         generalizer = lambda b: b
-    if isinstance(generalizer, abc.Sequence):
+    if isinstance(generalizer, collections.abc.Sequence):
         generalizer = _make_composite_generalizer(*generalizer)
 
     g = nx.DiGraph()
@@ -243,8 +243,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)