Use XorK bloq instead of CNOTs to load data in QROM (#1335)

* Use XorK bloq instead of CNOTs to load data in QROM

* Remove commented code and fix pylint

Author: tanujkhattar

qualtran/bloqs/data_loading/qrom.py

@@ -14,27 +14,17 @@
 
 """Quantum read-only memory."""
 import numbers
-from typing import (
-    Callable,
-    cast,
-    Iterable,
-    Iterator,
-    Optional,
-    Sequence,
-    Set,
-    Tuple,
-    TYPE_CHECKING,
-    Union,
-)
+from typing import cast, Iterable, Iterator, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import cirq
 import numpy as np
 import sympy
 from numpy.typing import ArrayLike, NDArray
 
-from qualtran import bloq_example, BloqDocSpec, Register
+from qualtran import bloq_example, BloqDocSpec, QUInt, Register
 from qualtran._infra.gate_with_registers import merge_qubits
+from qualtran.bloqs.arithmetic import XorK
 from qualtran.bloqs.basic_gates import CNOT
 from qualtran.bloqs.data_loading.qrom_base import QROMBase
 from qualtran.bloqs.mcmt.and_bloq import And, MultiAnd
@@ -127,7 +117,7 @@ def build_from_bitsize(
     def _load_nth_data(
         self,
         selection_idx: Tuple[int, ...],
-        gate: Callable[[cirq.Qid], cirq.Operation],
+        ctrl_qubits: Tuple[cirq.Qid, ...] = (),
         **target_regs: NDArray[cirq.Qid],  # type: ignore[type-var]
     ) -> Iterator[cirq.OP_TREE]:
         for i, d in enumerate(self.data):
@@ -136,20 +126,18 @@ def _load_nth_data(
             assert all(isinstance(x, (int, numbers.Integral)) for x in target_shape)
             for idx in np.ndindex(cast(Tuple[int, ...], target_shape)):
                 data_to_load = int(d[selection_idx + idx])
-                for q, bit in zip(target[idx], f'{data_to_load:0{target_bitsize}b}'):
-                    if int(bit):
-                        yield gate(q)
+                yield XorK(QUInt(target_bitsize), data_to_load).on(*target[idx]).controlled_by(
+                    *ctrl_qubits
+                )
 
     def decompose_zero_selection(
         self, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
     ) -> Iterator[cirq.OP_TREE]:
-        controls = merge_qubits(self.control_registers, **quregs)
+        controls = tuple(merge_qubits(self.control_registers, **quregs))
         target_regs = {reg.name: quregs[reg.name] for reg in self.target_registers}
         zero_indx = (0,) * len(self.data_shape)
-        if self.num_controls == 0:
-            yield self._load_nth_data(zero_indx, cirq.X, **target_regs)
-        elif self.num_controls == 1:
-            yield self._load_nth_data(zero_indx, lambda q: CNOT().on(controls[0], q), **target_regs)
+        if self.num_controls <= 1:
+            yield self._load_nth_data(zero_indx, ctrl_qubits=controls, **target_regs)
         else:
             ctrl = np.array(controls)[:, np.newaxis]
             junk = np.array(context.qubit_manager.qalloc(len(controls) - 2))[:, np.newaxis]
@@ -161,7 +149,7 @@ def decompose_zero_selection(
                     ctrl=ctrl, junk=junk, target=and_target
                 )
             yield multi_controlled_and
-            yield self._load_nth_data(zero_indx, lambda q: CNOT().on(and_target, q), **target_regs)
+            yield self._load_nth_data(zero_indx, ctrl_qubits=(and_target,), **target_regs)
             yield cirq.inverse(multi_controlled_and)
             context.qubit_manager.qfree(list(junk.flatten()) + [and_target])
 
@@ -182,7 +170,7 @@ def nth_operation(
     ) -> Iterator[cirq.OP_TREE]:
         selection_idx = tuple(kwargs[reg.name] for reg in self.selection_registers)
         target_regs = {reg.name: kwargs[reg.name] for reg in self.target_registers}
-        yield self._load_nth_data(selection_idx, lambda q: CNOT().on(control, q), **target_regs)
+        yield self._load_nth_data(selection_idx, ctrl_qubits=(control,), **target_regs)
 
     def _circuit_diagram_info_(self, args) -> cirq.CircuitDiagramInfo:
         from qualtran.cirq_interop._bloq_to_cirq import _wire_symbol_to_cirq_diagram_info

qualtran/bloqs/data_loading/qrom_test.py

@@ -285,11 +285,11 @@ def test_qrom_variable_spacing():
     _assert_qrom_has_diagram(
         qrom,
         r'''
-selection00: ───X───@───X───@───
-                    │       │
-target0_0: ─────────┼───────X───
-                    │
-target0_1: ─────────X───────────
+selection00: ───X───@────X───@────
+                    │        │
+target0_0: ─────────⊕1───────⊕2───
+                    │        │
+target0_1: ─────────⊕1───────⊕2───
     ''',
     )
     # When inner loop range is not a power of 2, the inner segment tree cannot be skipped.
@@ -302,16 +302,16 @@ def test_qrom_variable_spacing():
     _assert_qrom_has_diagram(
         qrom,
         r'''
-selection00: ───X───@─────────@───────@──────X───@─────────@───────@──────
-                    │         │       │          │         │       │
-selection10: ───────(0)───────┼───────@──────────(0)───────┼───────@──────
-                    │         │       │          │         │       │
-anc_1: ─────────────And───@───X───@───And†───────And───@───X───@───And†───
-                          │       │                    │       │
-target0_0: ───────────────┼───────┼────────────────────X───────X──────────
-                          │       │
-target0_1: ───────────────X───────X───────────────────────────────────────
-        ''',
+selection00: ───X───@──────────@────────@──────X───@──────────@────────@──────
+                    │          │        │          │          │        │
+selection10: ───────(0)────────┼────────@──────────(0)────────┼────────@──────
+                    │          │        │          │          │        │
+anc_1: ─────────────And───@────X───@────And†───────And───@────X───@────And†───
+                          │        │                     │        │
+target0_0: ───────────────⊕1───────⊕1────────────────────⊕2───────⊕2──────────
+                          │        │                     │        │
+target0_1: ───────────────⊕1───────⊕1────────────────────⊕2───────⊕2──────────
+''',
     )
     # No T-gates needed if all elements to load are identical.
     assert t_complexity(QROM.build_from_data([3, 3, 3, 3])).t == 0