Create a Subtract bloq

qualtran/bloqs/arithmetic/__init__.py

@@ -17,6 +17,7 @@
     AddConstantMod,
     OutOfPlaceAdder,
     SimpleAddConstant,
+    Subtract,
 )
 from qualtran.bloqs.arithmetic.comparison import (
     BiQubitsMixer,

qualtran/bloqs/arithmetic/addition.py

@@ -51,10 +51,10 @@
     SoquetT,
 )
 from qualtran._infra.data_types import QMontgomeryUInt
+from qualtran.bloqs import util_bloqs
 from qualtran.bloqs.basic_gates import CNOT, XGate
 from qualtran.bloqs.mcmt.and_bloq import And
 from qualtran.bloqs.mcmt.multi_control_multi_target_pauli import MultiControlX
-from qualtran.bloqs.util_bloqs import ArbitraryClifford
 from qualtran.cirq_interop import decompose_from_cirq_style_method
 from qualtran.cirq_interop.bit_tools import iter_bits, iter_bits_twos_complement
 from qualtran.cirq_interop.t_complexity_protocol import TComplexity
@@ -348,7 +348,7 @@ def _t_complexity_(self) -> TComplexity:
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
         return {
             (And(uncompute=self.is_adjoint), self.bitsize),
-            (ArbitraryClifford(n=2), 5 * self.bitsize),
+            (util_bloqs.ArbitraryClifford(n=2), 5 * self.bitsize),
         }
 
     def __pow__(self, power: int):
@@ -408,7 +408,7 @@ class SimpleAddConstant(Bloq):
         [Improved quantum circuits for elliptic curve discrete logarithms](https://arxiv.org/abs/2001.09580) Fig 2a
     """
 
-    bitsize: int
+    bitsize: Union[int, sympy.Expr]
     k: int
     cvs: Tuple[int, ...] = field(
         converter=lambda v: (v,) if isinstance(v, int) else tuple(v), default=()
@@ -445,6 +445,10 @@ def on_classical_vals(
     def build_composite_bloq(
         self, bb: 'BloqBuilder', x: Soquet, **regs: SoquetT
     ) -> Dict[str, 'SoquetT']:
+        if isinstance(self.bitsize, sympy.Expr):
+            raise ValueError(
+                f'Symbolic bitsize {self.bitsize} not supported for SimpleAddConstant.build_composite_bloq'
+            )
         # Assign registers to variables and allocate ancilla bits for classical integer k.
         if len(self.cvs) > 0:
             ctrls = regs['ctrls']
@@ -628,3 +632,145 @@ def _add_k_large() -> AddConstantMod:
 _ADD_K_DOC = BloqDocSpec(
     bloq_cls=AddConstantMod, examples=[_add_k_symb, _add_k_small, _add_k_large]
 )
+
+
+@frozen
+class Subtract(Bloq):
+    r"""An n-bit subtraction gate.
+
+    Implements $U|a\rangle|b\rangle \rightarrow |a\rangle|a-b\rangle$ using $4n - 4 T$ gates.
+
+    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.
+
+    Registers:
+        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()
+
+    @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.")
+
+    @property
+    def dtype(self):
+        if self.a_dtype != self.b_dtype:
+            raise ValueError(
+                "Add.dtype is only supported when both operands have the same dtype: "
+                f"{self.a_dtype=}, {self.b_dtype=}"
+            )
+        return self.a_dtype
+
+    @property
+    def signature(self):
+        return Signature([Register("a", self.a_dtype), Register("b", self.b_dtype)])
+
+    def on_classical_vals(
+        self, a: 'ClassicalValT', b: 'ClassicalValT'
+    ) -> Dict[str, 'ClassicalValT']:
+        unsigned = isinstance(self.a_dtype, (QUInt, QMontgomeryUInt))
+        b_bitsize = self.b_dtype.bitsize
+        N = 2**b_bitsize if unsigned else 2 ** (b_bitsize - 1)
+        return {'a': a, 'b': int(math.fmod(a - b, N))}
+
+    def short_name(self) -> str:
+        return "a-b"
+
+    def _circuit_diagram_info_(self, _) -> cirq.CircuitDiagramInfo:
+        wire_symbols = ["In(x)"] * int(self.a_dtype.bitsize)
+        wire_symbols += ["In(y)/Out(x-y)"] * int(self.b_dtype.bitsize)
+        return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols)
+
+    def wire_symbol(self, soq: 'Soquet') -> 'WireSymbol':
+        from qualtran.drawing import directional_text_box
+
+        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 build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+        a_dtype = (
+            self.a_dtype if not isinstance(self.a_dtype, QInt) else QUInt(self.a_dtype.bitsize)
+        )
+        b_dtype = (
+            self.b_dtype if not isinstance(self.b_dtype, QInt) else QUInt(self.b_dtype.bitsize)
+        )
+        return {
+            (XGate(), self.b_dtype.bitsize),
+            (SimpleAddConstant(self.b_dtype.bitsize, k=1), 1),
+            (Add(a_dtype, b_dtype), 1),
+            (util_bloqs.Split(self.b_dtype), 1),
+            (util_bloqs.Join(self.b_dtype), 1),
+        }
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', a: Soquet, b: Soquet) -> Dict[str, 'SoquetT']:
+        b = np.array([bb.add(XGate(), q=q) for q in bb.split(b)])  # 1s complement of b.
+        b = bb.add(
+            SimpleAddConstant(self.b_dtype.bitsize, k=1), x=bb.join(b, self.b_dtype)
+        )  # 2s complement of b.
+
+        a_dtype = (
+            self.a_dtype if not isinstance(self.a_dtype, QInt) else QUInt(self.a_dtype.bitsize)
+        )
+        b_dtype = (
+            self.b_dtype if not isinstance(self.b_dtype, QInt) else QUInt(self.b_dtype.bitsize)
+        )
+
+        a, b = bb.add(Add(a_dtype, b_dtype), a=a, b=b)  # a - b
+        return {'a': a, 'b': b}
+
+
+@bloq_example
+def _sub_symb() -> Subtract:
+    n = sympy.Symbol('n')
+    sub_symb = Subtract(QInt(bitsize=n))
+    return sub_symb
+
+
+@bloq_example
+def _sub_small() -> Subtract:
+    sub_small = Subtract(QInt(bitsize=4))
+    return sub_small
+
+
+@bloq_example
+def _sub_large() -> Subtract:
+    sub_large = Subtract(QInt(bitsize=64))
+    return sub_large
+
+
+@bloq_example
+def _sub_diff_size_regs() -> Subtract:
+    sub_diff_size_regs = Subtract(QInt(bitsize=4), QInt(bitsize=16))
+    return sub_diff_size_regs
+
+
+_SUB_DOC = BloqDocSpec(
+    bloq_cls=Subtract, examples=[_sub_symb, _sub_small, _sub_large, _sub_diff_size_regs]
+)

qualtran/bloqs/arithmetic/addition_test.py

@@ -26,6 +26,7 @@
     AddConstantMod,
     OutOfPlaceAdder,
     SimpleAddConstant,
+    Subtract,
 )
 from qualtran.bloqs.arithmetic.comparison_test import identity_map
 from qualtran.cirq_interop.bit_tools import iter_bits, iter_bits_twos_complement
@@ -381,6 +382,31 @@ def test_classical_simple_add_constant_signed(bitsize, k, x, cvs, ctrls, result)
     assert bloq_classical[-1] == result
 
 
+@pytest.mark.slow
+def test_subtract_bloq_decomposition():
+    gate = Subtract(QInt(3), QInt(5))
+    qlt_testing.assert_valid_bloq_decomposition(gate)
+
+    want = np.zeros((256, 256))
+    for a_b in range(256):
+        a, b = a_b >> 5, a_b & 31
+        c = (a - b) % 32
+        want[(a << 5) | c][a_b] = 1
+    got = gate.tensor_contract()
+    np.testing.assert_equal(got, want)
+
+
+def test_subtract_bloq_validation():
+    assert Subtract(QUInt(3)) == Subtract(QUInt(3), QUInt(3))
+    with pytest.raises(ValueError, match='bitsize must be less'):
+        _ = Subtract(QInt(5), QInt(3))
+    assert Subtract(QUInt(3)).dtype == QUInt(3)
+
+
+def test_subtract_bloq_consitant_counts():
+    qlt_testing.assert_equivalent_bloq_counts(Subtract(QInt(3), QInt(4)))
+
+
 @pytest.mark.notebook
 def test_notebook():
     qlt_testing.execute_notebook('addition')