diff --git a/qualtran/bloqs/arithmetic/controlled_addition.py b/qualtran/bloqs/arithmetic/controlled_addition.py new file mode 100644 index 000000000..15dd23100 --- /dev/null +++ b/qualtran/bloqs/arithmetic/controlled_addition.py @@ -0,0 +1,231 @@ +# 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. +import itertools +import math +from typing import Any, Dict, Iterator, Set, TYPE_CHECKING, Union + +import cirq +import numpy as np +import sympy +from attrs import field, frozen +from numpy.typing import NDArray + +from qualtran import Bloq, CompositeBloq, QBit, QInt, QUInt, Register, Signature, Soquet, SoquetT +from qualtran._infra.data_types import QMontgomeryUInt +from qualtran.bloqs.basic_gates import CNOT +from qualtran.bloqs.mcmt import MultiControlX +from qualtran.bloqs.mcmt.and_bloq import And +from qualtran.cirq_interop import decompose_from_cirq_style_method +from qualtran.cirq_interop.t_complexity_protocol import TComplexity + +if TYPE_CHECKING: + import quimb.tensor as qtn + + from qualtran.drawing import WireSymbol + from qualtran.resource_counting import BloqCountT, SympySymbolAllocator + from qualtran.simulation.classical_sim import ClassicalValT + + +@frozen +class CAdd(Bloq): + r"""An n-bit controlled-addition gate. + + Args: + a_dtype: Quantum datatype used to represent the integer a. + b_dtype: Quantum datatype used to represent the integer b. Must be large + enough to hold the result in the output register of a + b, or else it simply + drops the most significant bits. If not specified, b_dtype is set to a_dtype. + cv: When controlled=0, this bloq is active when the ctrl register is 0. When + controlled=1, this bloq is active when the ctrl register is 1. + + Registers: + ctrl: the control bit for the addition + a: A a_dtype.bitsize-sized input register (register a above). + b: A b_dtype.bitsize-sized input/output register (register b above). + + References: + [Halving the cost of quantum addition](https://arxiv.org/abs/1709.06648) + """ + + a_dtype: Union[QInt, QUInt, QMontgomeryUInt] = field() + b_dtype: Union[QInt, QUInt, QMontgomeryUInt] = field() + cv: int = field(default=1) + + @b_dtype.default + def b_dtype_default(self): + return self.a_dtype + + @a_dtype.validator + def _a_dtype_validate(self, field, val): + if not isinstance(val, (QInt, QUInt, QMontgomeryUInt)): + raise ValueError("Only QInt, QUInt and QMontgomerUInt types are supported.") + if isinstance(val.num_qubits, sympy.Expr): + return + if val.bitsize > self.b_dtype.bitsize: + raise ValueError("a_dtype bitsize must be less than or equal to b_dtype bitsize") + + @b_dtype.validator + def _b_dtype_validate(self, field, val): + if not isinstance(val, (QInt, QUInt, QMontgomeryUInt)): + raise ValueError("Only QInt, QUInt and QMontgomerUInt types are supported.") + + @cv.validator + def _controlled_validate(self, field, val): + if val not in (0, 1): + raise ValueError("controlled must be either 0 or 1") + + @property + def signature(self): + return Signature( + [Register("ctrl", QBit()), Register("a", self.a_dtype), Register("b", self.b_dtype)] + ) + + def add_my_tensors( + self, + tn: 'qtn.TensorNetwork', + tag: Any, + *, + incoming: Dict[str, 'SoquetT'], + outgoing: Dict[str, 'SoquetT'], + ): + import quimb.tensor as qtn + + if isinstance(self.a_dtype, QInt) or isinstance(self.b_dtype, QInt): + raise TypeError("Tensor contraction for addition is only supported for unsigned ints.") + N_a = 2**self.a_dtype.bitsize + N_b = 2**self.b_dtype.bitsize + inds = ( + incoming['ctrl'], + incoming['a'], + incoming['b'], + outgoing['ctrl'], + outgoing['a'], + outgoing['b'], + ) + unitary = np.zeros((2, N_a, N_b, 2, N_a, N_b), dtype=np.complex128) + for c, a, b in itertools.product(range(2), range(N_a), range(N_b)): + if c == self.cv: + unitary[c, a, b, c, a, int(math.fmod(a + b, N_b))] = 1 + else: + unitary[c, a, b, c, a, b] = 1 + + tn.add(qtn.Tensor(data=unitary, inds=inds, tags=[self.short_name(), tag])) + + def decompose_bloq(self) -> 'CompositeBloq': + return decompose_from_cirq_style_method(self) + + def on_classical_vals(self, **kwargs) -> Dict[str, 'ClassicalValT']: + a, b = kwargs['a'], kwargs['b'] + unsigned = isinstance(self.a_dtype, (QUInt, QMontgomeryUInt)) + b_bitsize = self.b_dtype.bitsize + N = 2**b_bitsize if unsigned else 2 ** (b_bitsize - 1) + ctrl = kwargs['ctrl'] + if ctrl != self.cv: + return {'ctrl': ctrl, 'a': a, 'b': b} + else: + return {'ctrl': ctrl, 'a': a, 'b': int(math.fmod(a + b, N))} + + def short_name(self) -> str: + return "a+b" + + def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol': + from qualtran.drawing import directional_text_box + + if soq.reg.name == 'ctrl': + return directional_text_box('ctrl', side=soq.reg.side) + if soq.reg.name == 'a': + return directional_text_box('a', side=soq.reg.side) + elif soq.reg.name == 'b': + return directional_text_box('a+b', side=soq.reg.side) + else: + raise ValueError() + + def _left_building_block(self, inp, out, anc, depth): + if depth == self.b_dtype.bitsize - 1: + return + else: + if depth < 1: + raise ValueError(f"{depth=} is not a positive integer") + if depth < len(inp): + yield CNOT().on(anc[depth - 1], inp[depth]) + control = inp[depth] + else: + # If inp[depth] doesn't exist, we treat it as a |0>, + # and therefore applying CNOT().on(anc[depth - 1], inp[depth]) + # essentially "copies" anc[depth - 1] into inp[depth] + # in the classical basis. So therefore, on future operations, + # we can use anc[depth - 1] in its place. + control = anc[depth - 1] + yield CNOT().on(anc[depth - 1], out[depth]) + yield And().on(control, out[depth], anc[depth]) + yield CNOT().on(anc[depth - 1], anc[depth]) + yield from self._left_building_block(inp, out, anc, depth + 1) + + def _right_building_block(self, inp, out, anc, control, depth): + if depth == 0: + return + yield CNOT().on(anc[depth - 1], anc[depth]) + if depth < len(inp): + yield And().adjoint().on(inp[depth], out[depth], anc[depth]) + yield MultiControlX((1, 1)).on(control, inp[depth], out[depth]) + yield CNOT().on(anc[depth - 1], inp[depth]) + else: + yield And().adjoint().on(anc[depth - 1], out[depth], anc[depth]) + yield MultiControlX((1, 1)).on(control, anc[depth - 1], out[depth]) + yield CNOT().on(anc[depth - 1], out[depth]) + yield from self._right_building_block(inp, out, anc, control, depth - 1) + + def decompose_from_registers( + self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid] # type: ignore[type-var] + ) -> Iterator[cirq.OP_TREE]: + # reverse the order of qubits for big endian-ness. + input_bits = quregs['a'][::-1] + output_bits = quregs['b'][::-1] + ancillas = context.qubit_manager.qalloc(self.b_dtype.bitsize - 1)[::-1] + control = quregs['ctrl'][0] + if self.cv == 0: + yield cirq.X(control) + # Start off the addition by anding into the ancilla + yield And().on(input_bits[0], output_bits[0], ancillas[0]) + # Left part of Fig.4 + yield from self._left_building_block(input_bits, output_bits, ancillas, 1) + yield CNOT().on(ancillas[-1], output_bits[-1]) + if len(input_bits) == len(output_bits): + yield MultiControlX((1, 1)).on(control, input_bits[-1], output_bits[-1]) + yield CNOT().on(ancillas[-1], output_bits[-1]) + # right part of Fig.4 + yield from self._right_building_block( + input_bits, output_bits, ancillas, control, self.b_dtype.bitsize - 2 + ) + yield And().adjoint().on(input_bits[0], output_bits[0], ancillas[0]) + yield MultiControlX((1, 1)).on(control, input_bits[0], output_bits[0]) + if self.cv == 0: + yield cirq.X(control) + context.qubit_manager.qfree(ancillas) + + def _t_complexity_(self): + n = self.b_dtype.bitsize + num_and = 2 * n - 1 + num_clifford = 33 * (n - 2) + 43 + return TComplexity(t=4 * num_and, clifford=num_clifford) + + def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']: + n = self.b_dtype.bitsize + n_cnot = (n - 2) * 6 + 2 + return { + (MultiControlX((1, 1)), n), + (And(), n - 1), + (And().adjoint(), n - 1), + (CNOT(), n_cnot), + } diff --git a/qualtran/bloqs/arithmetic/controlled_addition_test.py b/qualtran/bloqs/arithmetic/controlled_addition_test.py new file mode 100644 index 000000000..abfa4a1c2 --- /dev/null +++ b/qualtran/bloqs/arithmetic/controlled_addition_test.py @@ -0,0 +1,74 @@ +# 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. + +import cirq +import numpy as np +import pytest + +import qualtran.testing as qlt_testing +from qualtran import QUInt +from qualtran.bloqs.arithmetic.controlled_addition import CAdd +from qualtran.cirq_interop.testing import assert_circuit_inp_out_cirqsim +from qualtran.resource_counting.generalizers import ignore_split_join + + +def iter_bits(x, w): + return [int(b) for b in np.binary_repr(x, width=w)] + + +@pytest.mark.parametrize('a', [1, 2]) +@pytest.mark.parametrize('b', [1, 2, 3]) +@pytest.mark.parametrize('num_bits_a', [2, 3]) +@pytest.mark.parametrize('num_bits_b', [5]) +@pytest.mark.parametrize('controlled_on', [0, 1]) +@pytest.mark.parametrize('control', [0, 1]) +def test_controlled_addition(a, b, num_bits_a, num_bits_b, controlled_on, control): + num_anc = num_bits_b - 1 + gate = CAdd(QUInt(num_bits_a), QUInt(num_bits_b), cv=controlled_on) + qubits = cirq.LineQubit.range(num_bits_a + num_bits_b + 1) + op = gate.on_registers(ctrl=qubits[0], a=qubits[1 : num_bits_a + 1], b=qubits[num_bits_a + 1 :]) + greedy_mm = cirq.GreedyQubitManager(prefix="_a", maximize_reuse=True) + context = cirq.DecompositionContext(greedy_mm) + circuit = cirq.Circuit(cirq.decompose_once(op, context=context)) + circuit0 = cirq.Circuit(op) + ancillas = sorted(circuit.all_qubits())[-num_anc:] + initial_state = [0] * (num_bits_a + num_bits_b + num_anc + 1) + initial_state[0] = control + initial_state[1 : num_bits_a + 1] = list(iter_bits(a, num_bits_a)) + initial_state[num_bits_a + 1 : num_bits_a + num_bits_b + 1] = list(iter_bits(b, num_bits_b)) + final_state = [0] * (num_bits_a + num_bits_b + num_anc + 1) + final_state[0] = control + final_state[1 : num_bits_a + 1] = list(iter_bits(a, num_bits_a)) + if control == controlled_on: + final_state[num_bits_a + 1 : num_bits_a + num_bits_b + 1] = list( + iter_bits(a + b, num_bits_b) + ) + else: + final_state[num_bits_a + 1 : num_bits_a + num_bits_b + 1] = list(iter_bits(b, num_bits_b)) + assert_circuit_inp_out_cirqsim(circuit, qubits + ancillas, initial_state, final_state) + assert_circuit_inp_out_cirqsim( + circuit0, qubits, initial_state[:-num_anc], final_state[:-num_anc] + ) + + +@pytest.mark.parametrize("n", [*range(3, 10)]) +def test_addition_gate_counts_controlled(n: int): + add = CAdd(QUInt(n), cv=1) + num_and = 2 * n - 1 + t_count = 4 * num_and + + qlt_testing.assert_valid_bloq_decomposition(add) + assert add.t_complexity() == add.decompose_bloq().t_complexity() + assert add.bloq_counts() == add.decompose_bloq().bloq_counts(generalizer=ignore_split_join) + assert add.t_complexity().t == t_count