From 2917eeb177d036b70912d7581401a8bcb588f81b Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Fri, 9 Aug 2024 09:52:44 -0700 Subject: [PATCH] Use cost framework for sparse chemistry test (#1269) * Remove counting via TGates --- .../bloqs/chemistry/sparse/sparse_test.py | 85 ++++++++++--------- 1 file changed, 43 insertions(+), 42 deletions(-) diff --git a/qualtran/bloqs/chemistry/sparse/sparse_test.py b/qualtran/bloqs/chemistry/sparse/sparse_test.py index 7a794eab9..e910f3455 100644 --- a/qualtran/bloqs/chemistry/sparse/sparse_test.py +++ b/qualtran/bloqs/chemistry/sparse/sparse_test.py @@ -11,20 +11,25 @@ # 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, Union + import attrs import numpy as np import pytest from openfermion.resource_estimates.sparse.costing_sparse import cost_sparse from openfermion.resource_estimates.utils import power_two, QI +from qualtran import Adjoint, Bloq from qualtran.bloqs.arithmetic.comparison import LessThanEqual -from qualtran.bloqs.basic_gates import TGate from qualtran.bloqs.chemistry.sparse import PrepareSparse, SelectSparse from qualtran.bloqs.chemistry.sparse.prepare_test import build_random_test_integrals from qualtran.bloqs.data_loading.select_swap_qrom import find_optimal_log_block_size, SelectSwapQROM from qualtran.bloqs.state_preparation.prepare_uniform_superposition import ( PrepareUniformSuperposition, ) +from qualtran.resource_counting import get_cost_value, QECGatesCost +from qualtran.resource_counting.generalizers import generalize_cswap_approx +from qualtran.symbolics import SymbolicInt from qualtran.testing import execute_notebook @@ -62,7 +67,15 @@ def qrom_cost(prep: PrepareSparse) -> int: return num_toff_qrom -def get_sel_swap_qrom_t_count(prep: PrepareSparse) -> int: +def get_toffoli_count(bloq: Bloq) -> SymbolicInt: + """Get the toffoli count of a bloq assuming no raw Ts.""" + counts = get_cost_value(bloq, QECGatesCost(), generalizer=generalize_cswap_approx) + cost_dict = counts.total_t_and_ccz_count(ts_per_rotation=0) + assert cost_dict['n_t'] == 0 + return cost_dict['n_ccz'] + + +def get_sel_swap_qrom_toff_count(prep: PrepareSparse) -> SymbolicInt: """Utility function to pick out the SelectSwapQROM cost from the prepare call graph.""" def keep_qrom(bloq): @@ -71,76 +84,64 @@ def keep_qrom(bloq): return False _, sigma = prep.call_graph(keep=keep_qrom) - qrom_bloq = None + qrom_bloq: Optional[Union[SelectSwapQROM, Adjoint]] = None for k in sigma.keys(): if isinstance(k, SelectSwapQROM): qrom_bloq = k break + if isinstance(k, Adjoint) and isinstance(k.subbloq, SelectSwapQROM): + qrom_bloq = k + break if qrom_bloq is None: return 0 - return int(qrom_bloq.call_graph()[1].get(TGate(), 0)) + return get_toffoli_count(qrom_bloq) @pytest.mark.parametrize("num_spin_orb, num_bits_rot_aa", ((8, 3), (12, 4), (16, 3))) def test_sparse_costs_against_openfermion(num_spin_orb, num_bits_rot_aa): num_bits_state_prep = 12 - cost = 0 - bloq = SelectSparse(num_spin_orb) - _, sigma = bloq.call_graph() - cost += sigma[TGate()] + sel_sparse = SelectSparse(num_spin_orb) + cost = get_toffoli_count(sel_sparse) prep_sparse, num_non_zero = make_prep_sparse(num_spin_orb, num_bits_state_prep, num_bits_rot_aa) - _, sigma = prep_sparse.call_graph() - cost += sigma[TGate()] + cost += get_toffoli_count(prep_sparse) prep_sparse_adj = attrs.evolve( prep_sparse, is_adjoint=True, qroam_block_size=2 ** QI(num_non_zero)[0] ) - _, sigma = prep_sparse_adj.call_graph() - cost += sigma[TGate()] + cost += get_toffoli_count(prep_sparse_adj) unused_lambda = 10 unused_de = 1e-3 unused_stps = 10 logd = (num_non_zero - 1).bit_length() - refl_cost = 4 * (num_bits_state_prep + logd + 4) # page 40 listing 4 - # Correct the swap cost: - # 1. For prepare swaps are costed as Toffolis which we convert to 4 T gates, but a swap costs 7 T gates. - # 2. The reference inverts the swaps at zero cost for Prep^, so we need to add this cost back. - delta_swap = ( - 8 * (7 - 4) * (num_spin_orb // 2 - 1).bit_length() - + (7 - 4) - + 8 * 7 * (num_spin_orb // 2 - 1).bit_length() - + 7 - ) - cost_of = cost_sparse( - num_spin_orb, unused_lambda, num_non_zero, unused_de, num_bits_state_prep, unused_stps - )[0] + refl_cost = num_bits_state_prep + logd + 4 # page 40 listing 4 # correct the expected cost by using a different uniform superposition algorithm # see: https://github.com/quantumlib/Qualtran/issues/611 eta = power_two(prep_sparse.num_non_zero) - cost_uni_prep = ( - 4 - * 2 - * (3 * (prep_sparse.num_non_zero - 1).bit_length() - 3 * eta + 2 * num_bits_rot_aa - 9) + cost_uni_prep = 2 * ( + 3 * (prep_sparse.num_non_zero - 1).bit_length() - 3 * eta + 2 * num_bits_rot_aa - 9 + ) + prep_uni = PrepareUniformSuperposition(prep_sparse.num_non_zero) + delta_uni_prep = ( + get_toffoli_count(prep_uni) + get_toffoli_count(prep_uni.adjoint()) - cost_uni_prep ) - prep = PrepareUniformSuperposition(prep_sparse.num_non_zero) - cost1a_mod = prep.call_graph()[1][TGate()] - cost1a_mod += prep.adjoint().call_graph()[1][TGate()] # correct for SelectSwapQROM vs QROAM # https://github.com/quantumlib/Qualtran/issues/574 - our_qrom = get_sel_swap_qrom_t_count(prep_sparse) - our_qrom += get_sel_swap_qrom_t_count(prep_sparse_adj) paper_qrom = qrom_cost(prep_sparse) paper_qrom += qrom_cost(prep_sparse_adj) - delta_qrom = our_qrom - paper_qrom * 4 + qual_qrom_cost = get_sel_swap_qrom_toff_count(prep_sparse) + qual_qrom_cost += get_sel_swap_qrom_toff_count(prep_sparse_adj) + delta_qrom = qual_qrom_cost - paper_qrom # inequality test difference # https://github.com/quantumlib/Qualtran/issues/235 lte = LessThanEqual(prep_sparse.num_bits_state_prep, prep_sparse.num_bits_state_prep) - t_count_lte = 2 * lte.call_graph()[1][TGate()] - t_count_lte_paper = 4 * prep_sparse.num_bits_state_prep # inverted at zero cost - delta_ineq = t_count_lte - t_count_lte_paper # 4 * (prep_sparse.num_bits_state_prep + 1) - adjusted_cost_qualtran = ( - cost - cost1a_mod + cost_uni_prep + refl_cost - delta_swap - delta_qrom - delta_ineq - ) // 4 - assert adjusted_cost_qualtran == cost_of + lte_cost = get_toffoli_count(lte) + get_toffoli_count(lte.adjoint()) + 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 + cost_of = cost_sparse( + num_spin_orb, unused_lambda, num_non_zero, unused_de, num_bits_state_prep, unused_stps + )[0] + assert adjusted_cost_qualtran == cost_of - refl_cost @pytest.mark.notebook