From 21ef526b2b0021bff5ee9e8263568195cf0f8665 Mon Sep 17 00:00:00 2001 From: Tanuj Khattar Date: Tue, 24 Sep 2024 02:35:30 +0200 Subject: [PATCH] Add `my_static_costs` override for accurate `QubitCount`s for qrom bloqs (#1414) * Add my_static_costs for accurate QubitCounts for qrom bloqs * Fix mypy, pylint and autogenerate notebooks * Fix failing tests --- .../bloqs/chemistry/sparse/sparse_test.py | 7 +++- qualtran/bloqs/data_loading/qroam_clean.py | 18 ++++++++-- .../bloqs/data_loading/qroam_clean_test.py | 21 ++++++++++-- qualtran/bloqs/data_loading/qrom.py | 13 ++++++- qualtran/bloqs/data_loading/qrom_test.py | 18 +++++++++- .../bloqs/data_loading/select_swap_qrom.ipynb | 8 ++--- .../bloqs/data_loading/select_swap_qrom.py | 34 +++++++++++++++---- .../data_loading/select_swap_qrom_test.py | 18 +++++++++- .../mod_arithmetic/mod_multiplication_test.py | 4 +-- 9 files changed, 118 insertions(+), 23 deletions(-) diff --git a/qualtran/bloqs/chemistry/sparse/sparse_test.py b/qualtran/bloqs/chemistry/sparse/sparse_test.py index bb7823e8f..4b0c04ec1 100644 --- a/qualtran/bloqs/chemistry/sparse/sparse_test.py +++ b/qualtran/bloqs/chemistry/sparse/sparse_test.py @@ -79,6 +79,9 @@ def test_sparse_costs_against_openfermion(num_spin_orb, num_bits_rot_aa): # Qualtran (constants are not ignored). The difference arises from # uncontrolled unary iteration used by QROM, which QROAMClean delegates to. delta_qrom = -2 + # The -4 comes from QROAMCleanAdjoint, which delegates to a QROM and SwapWithZero + # and each of them contributes a -2 factor. + delta_qrom_adjoint = -4 # inequality test difference # https://github.com/quantumlib/Qualtran/issues/235 lte = LessThanEqual(prep_sparse.num_bits_state_prep, prep_sparse.num_bits_state_prep) @@ -86,7 +89,9 @@ def test_sparse_costs_against_openfermion(num_spin_orb, num_bits_rot_aa): lte_cost_paper = prep_sparse.num_bits_state_prep # inverted at zero cost delta_ineq = lte_cost - lte_cost_paper swap_cost = 8 * (num_spin_orb // 2 - 1).bit_length() + 1 # inverted at zero cost - adjusted_cost_qualtran = cost - delta_qrom - delta_uni_prep - delta_ineq - swap_cost + adjusted_cost_qualtran = ( + cost - delta_qrom - delta_uni_prep - delta_ineq - swap_cost - delta_qrom_adjoint + ) cost_of = cost_sparse( num_spin_orb, unused_lambda, num_non_zero, unused_de, num_bits_state_prep, unused_stps )[0] diff --git a/qualtran/bloqs/data_loading/qroam_clean.py b/qualtran/bloqs/data_loading/qroam_clean.py index e7dc9c78e..32b46bb70 100644 --- a/qualtran/bloqs/data_loading/qroam_clean.py +++ b/qualtran/bloqs/data_loading/qroam_clean.py @@ -30,7 +30,7 @@ if TYPE_CHECKING: from qualtran import Bloq, BloqBuilder, SoquetT, QDType from qualtran.simulation.classical_sim import ClassicalValT - from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator, CostKey from qualtran.bloqs.data_loading.select_swap_qrom import _alloc_anc_for_reg, SelectSwapQROM @@ -179,7 +179,7 @@ def with_log_block_sizes( def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': block_sizes = prod([2**k for k in self.log_block_sizes]) data_size = prod(self.data_shape) - n_toffoli = ceil(data_size / block_sizes) + block_sizes + n_toffoli = ceil(data_size / block_sizes) + block_sizes - 4 + self.num_controls return {Toffoli(): n_toffoli} @cached_property @@ -268,7 +268,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': block_sizes = prod([2**k for k in self.log_block_sizes]) data_size = prod(self.qroam_clean.data_shape) - n_toffoli = ceil(data_size / block_sizes) + block_sizes + n_toffoli = ceil(data_size / block_sizes) + block_sizes - 4 + self.qroam_clean.num_controls return {Toffoli(): n_toffoli} def adjoint(self) -> 'QROAMClean': @@ -299,6 +299,9 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) - return Circle() raise ValueError(f'Unknown register name {name}') + def __str__(self): + return 'QROAMCleanAdjoint' + @attrs.frozen class QROAMClean(SelectSwapQROM): @@ -450,6 +453,15 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT': ret[swz] += 1 return ret + def my_static_costs(self, cost_key: "CostKey"): + from qualtran.resource_counting import get_cost_value, QubitCount + + if isinstance(cost_key, QubitCount): + qrom_costs = get_cost_value(self.qrom_bloq, QubitCount()) + return qrom_costs + sum(self.log_block_sizes) + + return NotImplemented + def _build_composite_bloq_with_swz_clean( self, bb: 'BloqBuilder', diff --git a/qualtran/bloqs/data_loading/qroam_clean_test.py b/qualtran/bloqs/data_loading/qroam_clean_test.py index 96f30af53..1f8cd0f0a 100644 --- a/qualtran/bloqs/data_loading/qroam_clean_test.py +++ b/qualtran/bloqs/data_loading/qroam_clean_test.py @@ -22,7 +22,8 @@ QROAMClean, QROAMCleanAdjointWrapper, ) -from qualtran.symbolics import ceil +from qualtran.resource_counting import get_cost_value, QubitCount +from qualtran.symbolics import ceil, log2 def test_bloq_examples(bloq_autotester): @@ -30,6 +31,20 @@ def test_bloq_examples(bloq_autotester): bloq_autotester(_qroam_clean_multi_dim) +def test_qroam_clean_qubit_counts(): + bloq = _qroam_clean_multi_data.make() + assert get_cost_value(bloq, QubitCount()) == get_cost_value(bloq.decompose_bloq(), QubitCount()) + bloq = _qroam_clean_multi_dim.make() + assert get_cost_value(bloq, QubitCount()) == get_cost_value(bloq.decompose_bloq(), QubitCount()) + # Symbolic + N, b, k = sympy.symbols('N b k', positive=True, integer=True) + bloq = QROAMClean.build_from_bitsize((N,), (b,), log_block_sizes=(k,)) + K = 2**k + # log(N) - k ancilla are required for the nested unary iteration. + expected_qubits = K * b + 2 * ceil(log2(N)) - k - 1 + assert sympy.simplify(get_cost_value(bloq, QubitCount()) - expected_qubits) == 0 + + def test_t_complexity_1d_data_symbolic(): # 1D data, 1 dataset N, b, k = sympy.symbols('N b k') @@ -42,7 +57,7 @@ def test_t_complexity_1d_data_symbolic(): inv_k = sympy.symbols('kinv') inv_K = 2**inv_k bloq_inv = bloq_inv.with_log_block_sizes(log_block_sizes=(inv_k,)) - expected_toffoli_inv = ceil(N / inv_K) + inv_K + expected_toffoli_inv = ceil(N / inv_K) + inv_K - 4 assert bloq_inv.t_complexity().t == 4 * expected_toffoli_inv @@ -58,7 +73,7 @@ def test_t_complexity_2d_data_symbolic(): inv_k1, inv_k2 = sympy.symbols('kinv1, kinv2') inv_K1, inv_K2 = 2**inv_k1, 2**inv_k2 bloq_inv = bloq_inv.with_log_block_sizes(log_block_sizes=(inv_k1, inv_k2)) - expected_toffoli_inv = ceil(N1 * N2 / (inv_K1 * inv_K2)) + inv_K1 * inv_K2 + expected_toffoli_inv = ceil(N1 * N2 / (inv_K1 * inv_K2)) + inv_K1 * inv_K2 - 4 assert bloq_inv.t_complexity().t == 4 * expected_toffoli_inv diff --git a/qualtran/bloqs/data_loading/qrom.py b/qualtran/bloqs/data_loading/qrom.py index 59ace8544..363fc560b 100644 --- a/qualtran/bloqs/data_loading/qrom.py +++ b/qualtran/bloqs/data_loading/qrom.py @@ -33,7 +33,7 @@ from qualtran.symbolics import prod, SymbolicInt if TYPE_CHECKING: - from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, BloqCountT, CostKey, SympySymbolAllocator def _to_tuple(x: Iterable[NDArray]) -> Sequence[NDArray]: @@ -184,6 +184,17 @@ def _circuit_diagram_info_(self, args) -> cirq.CircuitDiagramInfo: return _wire_symbol_to_cirq_diagram_info(self, args) + def my_static_costs(self, cost_key: "CostKey"): + from qualtran.resource_counting import QubitCount + + if isinstance(cost_key, QubitCount): + return self.signature.n_qubits() + sum(self.selection_bitsizes) - 1 + self.num_controls + + return NotImplemented + + def __str__(self): + return f'QROM({self.data_shape}, {self.target_shapes}, {self.target_bitsizes})' + def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol': if reg is None: return Text('QROM') diff --git a/qualtran/bloqs/data_loading/qrom_test.py b/qualtran/bloqs/data_loading/qrom_test.py index 271346c34..af990ed80 100644 --- a/qualtran/bloqs/data_loading/qrom_test.py +++ b/qualtran/bloqs/data_loading/qrom_test.py @@ -25,7 +25,8 @@ from qualtran.bloqs.data_loading.qrom import _qrom_multi_data, _qrom_multi_dim, _qrom_small, QROM from qualtran.cirq_interop.t_complexity_protocol import t_complexity from qualtran.cirq_interop.testing import assert_circuit_inp_out_cirqsim, GateHelper -from qualtran.resource_counting import get_cost_value, QECGatesCost +from qualtran.resource_counting import get_cost_value, QECGatesCost, QubitCount +from qualtran.symbolics import ceil, log2 def test_qrom_small(bloq_autotester): @@ -40,6 +41,21 @@ def test_qrom_multi_dim(bloq_autotester): bloq_autotester(_qrom_multi_dim) +def test_qrom_qubit_counts(): + bloq = _qrom_small.make() + assert get_cost_value(bloq, QubitCount()) == get_cost_value(bloq.decompose_bloq(), QubitCount()) + bloq = _qrom_multi_data.make() + assert get_cost_value(bloq, QubitCount()) == get_cost_value(bloq.decompose_bloq(), QubitCount()) + bloq = _qrom_multi_dim.make() + assert get_cost_value(bloq, QubitCount()) == get_cost_value(bloq.decompose_bloq(), QubitCount()) + # Symbolic + N, b, c = sympy.symbols('N b c', positive=True, integer=True) + bloq = QROM.build_from_bitsize((N,), (b,), num_controls=c) + # log(N) ancilla are required for the ancilla used in unary iteration. + expected_qubits = 2 * ceil(log2(N)) + b + 2 * c - 1 + assert sympy.simplify(get_cost_value(bloq, QubitCount()) - expected_qubits) == 0 + + @pytest.mark.slow @pytest.mark.parametrize( "data,num_controls", diff --git a/qualtran/bloqs/data_loading/select_swap_qrom.ipynb b/qualtran/bloqs/data_loading/select_swap_qrom.ipynb index 4bd012a63..0c7dc4117 100644 --- a/qualtran/bloqs/data_loading/select_swap_qrom.ipynb +++ b/qualtran/bloqs/data_loading/select_swap_qrom.ipynb @@ -241,7 +241,7 @@ }, "outputs": [], "source": [ - "N, b, k, c = sympy.symbols('N b k c')\n", + "N, b, k, c = sympy.symbols('N b k c', positive=True, integers=True)\n", "qroam_symb_dirty_1d = SelectSwapQROM.build_from_bitsize(\n", " (N,), (b,), log_block_sizes=(k,), num_controls=c\n", ")" @@ -256,7 +256,7 @@ }, "outputs": [], "source": [ - "N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c')\n", + "N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c', positive=True, integers=True)\n", "log_block_sizes = (k1, k2)\n", "qroam_symb_dirty_2d = SelectSwapQROM.build_from_bitsize(\n", " (N, M), (b1, b2), log_block_sizes=log_block_sizes, num_controls=c\n", @@ -272,7 +272,7 @@ }, "outputs": [], "source": [ - "N, b, k, c = sympy.symbols('N b k c')\n", + "N, b, k, c = sympy.symbols('N b k c', positive=True, integers=True)\n", "qroam_symb_clean_1d = SelectSwapQROM.build_from_bitsize(\n", " (N,), (b,), log_block_sizes=(k,), num_controls=c, use_dirty_ancilla=False\n", ")" @@ -287,7 +287,7 @@ }, "outputs": [], "source": [ - "N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c')\n", + "N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c', positive=True, integers=True)\n", "log_block_sizes = (k1, k2)\n", "qroam_symb_clean_2d = SelectSwapQROM.build_from_bitsize(\n", " (N, M), (b1, b2), log_block_sizes=log_block_sizes, num_controls=c, use_dirty_ancilla=False\n", diff --git a/qualtran/bloqs/data_loading/select_swap_qrom.py b/qualtran/bloqs/data_loading/select_swap_qrom.py index 74b86c269..b3ab05eef 100644 --- a/qualtran/bloqs/data_loading/select_swap_qrom.py +++ b/qualtran/bloqs/data_loading/select_swap_qrom.py @@ -22,7 +22,15 @@ import sympy from numpy.typing import ArrayLike -from qualtran import bloq_example, BloqDocSpec, BQUInt, GateWithRegisters, Register, Signature +from qualtran import ( + bloq_example, + BloqDocSpec, + BQUInt, + DecomposeTypeError, + GateWithRegisters, + Register, + Signature, +) from qualtran.bloqs.arithmetic.bitwise import Xor from qualtran.bloqs.bookkeeping import Partition from qualtran.bloqs.data_loading.qrom import QROM @@ -33,7 +41,7 @@ if TYPE_CHECKING: from qualtran import Bloq, BloqBuilder, QDType, SoquetT - from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator + from qualtran.resource_counting import BloqCountDictT, CostKey, SympySymbolAllocator SelSwapQROM_T = TypeVar('SelSwapQROM_T', bound='SelectSwapQROM') @@ -408,7 +416,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str target = [soqs.pop(reg.name) for reg in self.target_registers] # Allocate intermediate clean/dirty ancilla for the underlying QROM call. if is_symbolic(*self.block_sizes): - raise ValueError( + raise DecomposeTypeError( f"Cannot decompose SelectSwapQROM bloq with symbolic block sizes. Found {self.block_sizes=}" ) block_sizes = cast(Tuple[int, ...], self.block_sizes) @@ -448,6 +456,18 @@ def _circuit_diagram_info_(self, args) -> cirq.CircuitDiagramInfo: return _wire_symbol_to_cirq_diagram_info(self, args) + def my_static_costs(self, cost_key: "CostKey"): + from qualtran.resource_counting import get_cost_value, QubitCount + + if isinstance(cost_key, QubitCount): + return ( + get_cost_value(self.qrom_bloq, QubitCount()) + + sum(self.log_block_sizes) + + sum(self.target_bitsizes) + ) + + return NotImplemented + def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol': if reg is None: return Text('QROAM') @@ -482,7 +502,7 @@ def _qroam_multi_dim() -> SelectSwapQROM: @bloq_example def _qroam_symb_dirty_1d() -> SelectSwapQROM: - N, b, k, c = sympy.symbols('N b k c') + N, b, k, c = sympy.symbols('N b k c', positive=True, integers=True) qroam_symb_dirty_1d = SelectSwapQROM.build_from_bitsize( (N,), (b,), log_block_sizes=(k,), num_controls=c ) @@ -491,7 +511,7 @@ def _qroam_symb_dirty_1d() -> SelectSwapQROM: @bloq_example def _qroam_symb_dirty_2d() -> SelectSwapQROM: - N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c') + N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c', positive=True, integers=True) log_block_sizes = (k1, k2) qroam_symb_dirty_2d = SelectSwapQROM.build_from_bitsize( (N, M), (b1, b2), log_block_sizes=log_block_sizes, num_controls=c @@ -501,7 +521,7 @@ def _qroam_symb_dirty_2d() -> SelectSwapQROM: @bloq_example def _qroam_symb_clean_1d() -> SelectSwapQROM: - N, b, k, c = sympy.symbols('N b k c') + N, b, k, c = sympy.symbols('N b k c', positive=True, integers=True) qroam_symb_clean_1d = SelectSwapQROM.build_from_bitsize( (N,), (b,), log_block_sizes=(k,), num_controls=c, use_dirty_ancilla=False ) @@ -510,7 +530,7 @@ def _qroam_symb_clean_1d() -> SelectSwapQROM: @bloq_example def _qroam_symb_clean_2d() -> SelectSwapQROM: - N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c') + N, M, b1, b2, k1, k2, c = sympy.symbols('N M b1 b2 k1 k2 c', positive=True, integers=True) log_block_sizes = (k1, k2) qroam_symb_clean_2d = SelectSwapQROM.build_from_bitsize( (N, M), (b1, b2), log_block_sizes=log_block_sizes, num_controls=c, use_dirty_ancilla=False diff --git a/qualtran/bloqs/data_loading/select_swap_qrom_test.py b/qualtran/bloqs/data_loading/select_swap_qrom_test.py index 975120bbf..8db1c9836 100644 --- a/qualtran/bloqs/data_loading/select_swap_qrom_test.py +++ b/qualtran/bloqs/data_loading/select_swap_qrom_test.py @@ -15,6 +15,7 @@ import cirq import numpy as np import pytest +import sympy from qualtran._infra.data_types import QUInt from qualtran._infra.gate_with_registers import get_named_qubits, split_qubits @@ -27,7 +28,8 @@ ) from qualtran.cirq_interop.t_complexity_protocol import t_complexity, TComplexity from qualtran.cirq_interop.testing import assert_circuit_inp_out_cirqsim -from qualtran.resource_counting import GateCounts, get_cost_value, QECGatesCost +from qualtran.resource_counting import GateCounts, get_cost_value, QECGatesCost, QubitCount +from qualtran.symbolics import ceil, log2 from qualtran.testing import assert_valid_bloq_decomposition @@ -192,6 +194,20 @@ def test_qroam_t_complexity(): assert qroam.t_complexity() == TComplexity(t=192, clifford=1082) +def test_selswap_qubit_counts(): + bloq = _qroam_multi_data.make() + assert get_cost_value(bloq, QubitCount()) == get_cost_value(bloq.decompose_bloq(), QubitCount()) + bloq = _qroam_multi_dim.make() + assert get_cost_value(bloq, QubitCount()) == get_cost_value(bloq.decompose_bloq(), QubitCount()) + # Symbolic + N, b, k = sympy.symbols('N b k', positive=True, integer=True) + bloq = SelectSwapQROM.build_from_bitsize((N,), (b,), log_block_sizes=(k,)) + K = 2**k + # log(N) - k ancilla are required for the nested unary iteration. + expected_qubits = K * b + b + 2 * ceil(log2(N)) - k - 1 + assert sympy.simplify(get_cost_value(bloq, QubitCount()) - expected_qubits) == 0 + + def test_qroam_many_registers(): # Test > 10 registers which resulted in https://github.com/quantumlib/Qualtran/issues/556 target_bitsizes = (3,) * 10 + (1,) * 2 + (3,) diff --git a/qualtran/bloqs/mod_arithmetic/mod_multiplication_test.py b/qualtran/bloqs/mod_arithmetic/mod_multiplication_test.py index 3be8b3793..d6a1caa62 100644 --- a/qualtran/bloqs/mod_arithmetic/mod_multiplication_test.py +++ b/qualtran/bloqs/mod_arithmetic/mod_multiplication_test.py @@ -183,11 +183,11 @@ def test_dirtyoutofplacemontgomerymodmul_symbolic_cost(uncompute): # Litinski 2023 https://arxiv.org/abs/2306.08585 # Figure/Table 8. Lists modular multiplication as 2.25n^2+9n toffoli. - # The following formula is 2.25n^2+8.25n-1 written with rationals because sympy comparison fails with floats. + # The following formula is 2.25n^2+7.25n-1 written with rationals because sympy comparison fails with floats. assert isinstance(cost['n_ccz'], sympy.Expr) assert ( cost['n_ccz'].subs(m, n / 4).expand() - == sympy.Rational(9, 4) * n**2 + sympy.Rational(33, 4) * n - 1 + == sympy.Rational(9, 4) * n**2 + sympy.Rational(29, 4) * n - 1 )