Add GF2MulMBUC Bloq (#1718)

Add Bloq for measurement based uncomputation of GF2 Multiplication
circuits. The `Parity` bloq added to compute the parity of input bits
should be generalized into a more general `CtrlSpec` that takes a bunch
of control registers and uses their parity to control the operation.
I'll open a new issue for it.

Author: tanujkhattar

qualtran/bloqs/basic_gates/diag_gates.ipynb

@@ -192,8 +192,8 @@
     "Two-qubit controlled-Z gate.\n",
     "\n",
     "#### Registers\n",
-    " - `ctrl`: One-bit control register.\n",
-    " - `target`: One-bit target register.\n"
+    " - `q1`: One-bit control register.\n",
+    " - `q2`: One-bit target register.\n"
    ]
   },
   {

qualtran/bloqs/basic_gates/z_basis.py

@@ -320,8 +320,8 @@ class CZ(Bloq):
     """Two-qubit controlled-Z gate.
 
     Registers:
-        ctrl: One-bit control register.
-        target: One-bit target register.
+        q1: One-bit control register.
+        q2: One-bit target register.
     """
 
     @cached_property

qualtran/bloqs/gf_arithmetic/gf2_multiplication.ipynb

@@ -655,8 +655,7 @@
     "The toffoli complexity is $n^{\\log_2{3}}$\n",
     "\n",
     "#### Parameters\n",
-    " - `m_x`: The irreducible polynomial that defines the galois field.\n",
-    " - `uncompute`: Whether to compute or uncompute the product. \n",
+    " - `m_x`: The irreducible polynomial that defines the galois field. \n",
     "\n",
     "#### Registers\n",
     " - `x`: A TRHU register representing the first number (or polynomial).\n",

qualtran/bloqs/gf_arithmetic/gf2_multiplication.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from functools import cached_property
-from typing import Dict, Optional, Sequence, Set, TYPE_CHECKING, Union
+from typing import Dict, Mapping, Optional, Sequence, Set, TYPE_CHECKING, Union
 
 import attrs
 import galois
@@ -24,20 +24,22 @@
     Bloq,
     bloq_example,
     BloqDocSpec,
+    CBit,
+    CtrlSpec,
     DecomposeTypeError,
     QBit,
     QGF,
     Register,
     Side,
     Signature,
 )
-from qualtran.bloqs.basic_gates import CNOT, Toffoli
+from qualtran.bloqs.basic_gates import CNOT, CZ, Discard, MeasureX, Toffoli
 from qualtran.symbolics import ceil, is_symbolic, log2, Shaped, SymbolicInt
 
 if TYPE_CHECKING:
     from qualtran import BloqBuilder, Soquet, SoquetT
     from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator
-    from qualtran.simulation.classical_sim import ClassicalValT
+    from qualtran.simulation.classical_sim import ClassicalValRetT, ClassicalValT
 
 
 def _data_or_shape_to_tuple(data_or_shape: Union[np.ndarray, Shaped]) -> tuple:
@@ -172,21 +174,6 @@ def signature(self) -> 'Signature':
     def bitsize(self) -> SymbolicInt:
         return self.qgf.bitsize
 
-    @cached_property
-    def reduction_matrix_q(self) -> np.ndarray:
-        m = int(self.bitsize)
-        f = self.qgf.gf_type.irreducible_poly
-        M = np.zeros((m, m))
-        alpha = [1] + [0] * m
-        for i in range(m - 1):
-            # x ** (m + i) % f
-            coeffs = (Poly(alpha, GF(2)) % f).coeffs.tolist()[::-1]
-            coeffs = coeffs + [0] * (m - len(coeffs))
-            M[i] = coeffs
-            alpha += [0]
-        M[m - 1][m - 1] = 1
-        return np.transpose(M)
-
     def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'Soquet') -> Dict[str, 'Soquet']:
         if is_symbolic(self.bitsize):
             raise DecomposeTypeError(f"Cannot decompose symbolic {self}")
@@ -261,6 +248,110 @@ def _gf2_multiplication_symbolic() -> GF2Multiplication:
 )
 
 
+@attrs.frozen
+class Parity(Bloq):
+    n: int
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature(
+            [
+                Register('x', dtype=CBit(), shape=(self.n,)),
+                Register('parity', dtype=CBit(), side=Side.RIGHT),
+            ]
+        )
+
+    def on_classical_vals(
+        self, *, x: Union['sympy.Symbol', 'ClassicalValT']
+    ) -> Mapping[str, 'ClassicalValRetT']:
+        assert isinstance(x, np.ndarray)
+        return {'x': x, 'parity': np.sum(x, dtype=int) & 1}
+
+
+@attrs.frozen
+class GF2MulMBUC(Bloq):
+    r"""Measurement based uncomputation of out of place multiplication over GF($2^m$).
+
+    Args:
+        bitsize: The degree $m$ of the galois field $GF(2^m)$. Also corresponds to the number of
+            qubits in each of the two input registers $a$ and $b$ that should be multiplied.
+
+    Registers:
+        x: Input THRU register of size $m$ that stores elements from $GF(2^m)$.
+        y: Input THRU register of size $m$ that stores elements from $GF(2^m)$.
+        result: Register of size $m$ that stores the product $x * y$ in $GF(2^m)$.
+    """
+
+    qgf: QGF = attrs.field(converter=_qgf_converter)
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature(
+            [
+                Register('x', dtype=self.qgf),
+                Register('y', dtype=self.qgf),
+                Register('result', dtype=self.qgf, side=Side.LEFT),
+            ]
+        )
+
+    @cached_property
+    def bitsize(self) -> SymbolicInt:
+        return self.qgf.bitsize
+
+    @cached_property
+    def reduction_matrix_q(self) -> np.ndarray:
+        m = int(self.bitsize)
+        f = self.qgf.gf_type.irreducible_poly
+        M = np.zeros((m, m), dtype=int)
+        alpha = [1] + [0] * m
+        for i in range(m):
+            # x ** (m + i) % f
+            coeffs = (Poly(alpha, GF(2)) % f).coeffs.tolist()[::-1]
+            coeffs = coeffs + [0] * (m - len(coeffs))
+            M[i] = coeffs
+            alpha += [0]
+        return np.transpose(M)
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'Soquet') -> Dict[str, 'Soquet']:
+        if is_symbolic(self.bitsize):
+            raise DecomposeTypeError(f"Cannot decompose symbolic {self}")
+        x, y, result = soqs['x'], soqs['y'], soqs['result']
+        x, y, result = bb.split(x)[::-1], bb.split(y)[::-1], bb.split(result)[::-1]
+        result = np.array([bb.add(MeasureX(), q=q) for q in result])
+        m = int(self.bitsize)
+
+        # Inverse of Step-3: Multiply Monomials
+        ctrl_cz = CZ().controlled(CtrlSpec(qdtypes=[CBit()]))
+        for i in range(m):
+            for j in range(i + 1):
+                result[i], x[j], y[i - j] = bb.add(ctrl_cz, ctrl=result[i], q1=x[j], q2=y[i - j])
+
+        # Inverse of Step-1 & 2: Multiply Monomials.
+        for i in range(m):
+            inp_vec = GF(2).Zeros(m)
+            inp_vec[i] = 1
+            out_vec = GF(2)(self.reduction_matrix_q) @ inp_vec
+            indices = [k for k in range(m) if out_vec[k]]
+            result[indices], parity = bb.add(Parity(len(indices)), x=result[indices])
+            for j in range(i + 1, m):
+                parity, x[m - j + i], y[j] = bb.add(ctrl_cz, ctrl=parity, q1=x[m - j + i], q2=y[j])
+            bb.add(Discard(), c=parity)
+
+        # Done :)
+        for c in result:
+            bb.add(Discard(), c=c)
+        return {'x': bb.join(x[::-1], dtype=self.qgf), 'y': bb.join(y[::-1], dtype=self.qgf)}
+
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
+        m = self.bitsize
+        return {CZ(): m**2}
+
+    def adjoint(self) -> 'Bloq':
+        return GF2MulViaKaratsuba(self.qgf)
+
+
 @attrs.frozen
 class GF2MulK(Bloq):
     r"""Multiply by constant $f(x)$ modulo $m(x)$. Both $f(x)$ and $m(x)$ are constants.
@@ -953,7 +1044,6 @@ class GF2MulViaKaratsuba(Bloq):
 
     Args:
         m_x: The irreducible polynomial that defines the galois field.
-        uncompute: Whether to compute or uncompute the product.
 
     Registers:
         x: A TRHU register representing the first number (or polynomial).
@@ -966,7 +1056,6 @@ class GF2MulViaKaratsuba(Bloq):
     """
 
     dtype: QGF = attrs.field(converter=_qgf_converter)
-    uncompute: bool = False
 
     @cached_property
     def m_x(self):
@@ -988,21 +1077,16 @@ def gf(self):
     def qgf(self):
         return self.dtype
 
-    def adjoint(self) -> 'GF2MulViaKaratsuba':
-        return attrs.evolve(self, uncompute=not self.uncompute)
-
     def __str__(self):
-        return f'{self.__class__.__name__}†' if self.uncompute else f'{self.__class__.__name__}'
+        return f'{self.__class__.__name__}'
 
     @cached_property
     def signature(self) -> 'Signature':
-        # C is directional
-        side = Side.LEFT if self.uncompute else Side.RIGHT
         return Signature(
             [
                 Register('x', dtype=self.qgf),
                 Register('y', dtype=self.qgf),
-                Register('result', dtype=self.qgf, side=side),
+                Register('result', dtype=self.qgf, side=Side.RIGHT),
             ]
         )
 
@@ -1016,8 +1100,6 @@ def k(self):
     @cached_property
     def _GF2MulViaKaratsubamod_impl(self) -> Bloq:
         impl = _GF2MulViaKaratsubaImpl(self.m_x)
-        if self.uncompute:
-            return impl.adjoint()
         return impl
 
     def build_composite_bloq(
@@ -1026,17 +1108,10 @@ def build_composite_bloq(
         if is_symbolic(self.k, self.n):
             raise DecomposeTypeError(f"Symbolic Decomposition is not supported for {self}")
 
-        if self.uncompute:
-            result = soqs['result']
-        else:
-            result = bb.allocate(self.n, self.qgf)
+        result = bb.allocate(self.n, self.qgf)
 
         x, y, result = bb.add_from(self._GF2MulViaKaratsubamod_impl, f=x, g=y, h=result)
 
-        if self.uncompute:
-            bb.free(result)  # type: ignore[arg-type]
-            return {'x': x, 'y': y}
-
         return {'x': x, 'y': y, 'result': result}
 
     def build_call_graph(
@@ -1069,11 +1144,11 @@ def on_classical_vals(
     ) -> Dict[str, 'ClassicalValT']:
         assert isinstance(x, self.gf)
         assert isinstance(y, self.gf)
-        if self.uncompute:
-            assert x * y == result
-            return {'x': x, 'y': y}
         return {'x': x, 'y': y, 'result': x * y}
 
+    def adjoint(self) -> 'Bloq':
+        return GF2MulMBUC(self.qgf)
+
 
 @bloq_example
 def _gf2mulviakaratsuba() -> GF2MulViaKaratsuba:

qualtran/bloqs/gf_arithmetic/gf2_multiplication_test.py

@@ -22,9 +22,11 @@
 from qualtran import QGF
 from qualtran.bloqs.gf_arithmetic.gf2_multiplication import (
     _gf2_multiplication_symbolic,
+    _GF2MulViaKaratsubaImpl,
     _gf16_multiplication,
     BinaryPolynomialMultiplication,
     GF2MulK,
+    GF2MulMBUC,
     GF2Multiplication,
     GF2MulViaKaratsuba,
     GF2ShiftLeft,
@@ -34,6 +36,7 @@
 )
 from qualtran.resource_counting import get_cost_value, QECGatesCost
 from qualtran.resource_counting.generalizers import ignore_alloc_free, ignore_split_join
+from qualtran.simulation.classical_sim import do_phased_classical_simulation
 from qualtran.testing import assert_consistent_classical_action
 
 
@@ -47,7 +50,7 @@ def test_gf2_multiplication_symbolic(bloq_autotester):
 
 @pytest.mark.parametrize('m', [2, 4, 6, 8])
 def test_synthesize_lr_circuit(m: int):
-    matrix = GF2Multiplication(m).reduction_matrix_q
+    matrix = GF2MulMBUC(m).reduction_matrix_q
     bloq = SynthesizeLRCircuit(matrix)
     bloq_adj = bloq.adjoint()
     QGFM, GFM = QGF(2, m), GF(2**m)
@@ -61,7 +64,7 @@ def test_synthesize_lr_circuit(m: int):
 @pytest.mark.slow
 @pytest.mark.parametrize('m', [3, 4, 5])
 def test_synthesize_lr_circuit_slow(m):
-    matrix = GF2Multiplication(m).reduction_matrix_q
+    matrix = GF2MulMBUC(m).reduction_matrix_q
     bloq = SynthesizeLRCircuit(matrix)
     bloq_adj = bloq.adjoint()
     QGFM, GFM = QGF(2, m), GF(2**m)
@@ -349,15 +352,16 @@ def test_gf2mulmod_classical_action_slow():
 
 @pytest.mark.parametrize('m_x', [[2, 1, 0], [3, 1, 0], [5, 2, 0]])
 def test_gf2mulmod_classical_action_adjoint(m_x):
-    blq = GF2MulViaKaratsuba(m_x)
+    blq = _GF2MulViaKaratsubaImpl(m_x)
     adjoint = blq.adjoint()
     rs = np.random.default_rng(42)
-    for i, j in rs.integers(0, len(blq.gf.elements) - 1, (10, 2)):
-        f = blq.gf.elements[i]
-        g = blq.gf.elements[j]
-        a, b, c = blq.call_classically(x=f, y=g)
-        a, b = adjoint.call_classically(x=a, y=b, result=c)
-        assert a == f and b == g
+    zero = blq.qgf.gf_type(0)
+    for i, j in rs.integers(0, len(blq.qgf.gf_type.elements) - 1, (10, 2)):
+        f = blq.qgf.gf_type.elements[i]
+        g = blq.qgf.gf_type.elements[j]
+        a, b, c = blq.call_classically(f=f, g=g, h=zero)
+        a, b, c = adjoint.call_classically(f=a, g=b, h=c)
+        assert a == f and b == g and c == zero
 
 
 @pytest.mark.parametrize('m_x', [[2, 1, 0], [8, 4, 3, 1, 0], [16, 5, 3, 1, 0]])
@@ -374,3 +378,35 @@ def test_gf2mulmod_classical_complexity(m_x):
 def test_gf2mul_invalid_input_raises():
     with pytest.raises(ValueError):
         _ = GF2MulViaKaratsuba([0, 1])  # type: ignore[arg-type]
+
+
+def test_gf2_mul_mbuc_quick():
+    m = 3
+    bloq_mbuc = GF2MulMBUC(m)
+    rng = np.random.default_rng(seed=123)
+    for x in bloq_mbuc.qgf.gf_type.elements:
+        for y in bloq_mbuc.qgf.gf_type.elements:
+            in_vals = {'x': x, 'y': y, 'result': x * y}
+            out_vals, phase = do_phased_classical_simulation(
+                bloq_mbuc.decompose_bloq(), in_vals, rng=rng
+            )
+            assert out_vals['x'] == x
+            assert out_vals['y'] == y
+            assert 'result' not in out_vals
+            assert phase == 1
+
+
+@pytest.mark.parametrize('m', [4, 5])
+def test_gf2_mul_mbuc(m: int):
+    bloq_mbuc = GF2MulMBUC(m)
+    rng = np.random.default_rng(seed=123)
+    for x in bloq_mbuc.qgf.gf_type.elements:
+        for y in bloq_mbuc.qgf.gf_type.elements:
+            in_vals = {'x': x, 'y': y, 'result': x * y}
+            out_vals, phase = do_phased_classical_simulation(
+                bloq_mbuc.decompose_bloq(), in_vals, rng=rng
+            )
+            assert out_vals['x'] == x
+            assert out_vals['y'] == y
+            assert 'result' not in out_vals
+            assert phase == 1