Add an externally controlled QROM lookup

python/cudaq_algorithms/primitives/__init__.py

@@ -10,7 +10,8 @@
 ``QROM`` builds coherent
 classical-table lookups on top of it, either as the plain
 unary-iteration SELECT or as the SELECT-SWAP / QROAM construction of
-Low, Kliuchnikov and Schaeffer (`arXiv:1812.00954`).
+Low, Kliuchnikov and Schaeffer (`arXiv:1812.00954`), either one
+optionally under an external control (``controlled=True``).
 
 Everything here is strictly unitary: the papers' headline Toffoli counts
 assume measurement-based ancilla uncomputation, which this library does

python/cudaq_algorithms/primitives/_qrom.py

@@ -4,7 +4,8 @@
 
 ``QROM(data, address_bits, output_bits)`` mints the lookup
 ``|k>|y> -> |k>|y XOR data[k]>`` behind one kernel signature
-``(address: qview, ladder: qview, output: qview)`` — all little-endian
+``(address: qview, ladder: qview, output: qview)`` — with a leading
+one-qubit ``control`` view when ``controlled=True`` — all little-endian
 (qubit 0 = LSB, ``docs/conventions.md``), ``ladder`` being ``num_ladder``
 clean ancillas (``|0>`` in, ``|0>`` out). Addresses ``k >= len(data)`` read as
 zero. Two constructions sit behind ``variant=``:
@@ -39,6 +40,28 @@
   binary routing network is ``B - 1`` register swaps per direction, and
   both the routing and the block writes run twice (compute + unitary
   uncompute).
+- ``controlled=True`` adds an external control qubit and replaces ``W``
+  by its controlled count (``_walk_toffoli_count(n, m, controlled)``: the
+  walk's root gadgets become Toffolis and its root crossings lose the
+  unguarded shortcut). Only the walk passes carry the control: the walk
+  sub-kernel is minted ``controlled=True`` and the parent passes the
+  control view to both walk passes, while the routing network and the
+  copy are the plain kernels. At control ``|0>`` a controlled walk on a
+  clean ladder writes nothing, so the block registers stay ``|0>`` and
+  the routing Fredkins, the copy CNOTs and the un-route/un-write act on
+  ``|0>`` as the identity. The lookup is then the identity at control
+  ``|0>`` and the plain lookup at ``|1>``, with the control as a separate
+  one-qubit view first in the signature, as ``unary_iteration_kernels``
+  does — and at a fraction of the price of ``cudaq.control`` over the
+  composition, which would add a control to every gate: for eleven 3-bit
+  entries over four address bits the native lookup compiles to 19 / 26 /
+  28 Toffolis (select, select_swap with ``B = 2`` and ``B = 4``) against
+  15 / 18 / 22 plain, where the wrapped composition costs two to three
+  times the native count, widest on select_swap because the wrapper also
+  controls the Fredkins and the copy CNOTs. The identity at
+  ``|0>`` is a clean-ladder statement (the clean ladder is the kernel's
+  precondition in any case): on a dirty ladder the deeper walk Toffolis
+  are gated by ladder lines, not by the control.
 
 The papers' headline counts (``N - 1`` lookup Toffolis for SELECT,
 ``ceil(N/B) + output_bits * (B - 1)`` for QROAM) price ancilla
@@ -60,7 +83,8 @@
 ------------------------------------------------------------------------
 
 Every minted kernel has the same signature, ``(address, ladder,
-output)``, and ``ladder`` is simply **the clean-ancilla register** —
+output)`` (``(control, address, ladder, output)`` when ``controlled``),
+and ``ladder`` is simply **the clean-ancilla register** —
 its contents differ per variant, and for ``select_swap`` the name is
 looser than the physics:
 
@@ -152,13 +176,16 @@ def primitives_qrom_route(low: cudaq.qview, blocks: cudaq.qview):
     return primitives_qrom_route
 
 
-def _select_swap_cost(address_bits: int, num_entries: int, output_bits: int,
-                      block_size: int) -> int:
+def _select_swap_cost(address_bits: int,
+                      num_entries: int,
+                      output_bits: int,
+                      block_size: int,
+                      controlled: bool = False) -> int:
     """Documented Toffoli price of the select_swap variant."""
     high_bits = address_bits - (block_size.bit_length() - 1)
     num_blocks = -(-num_entries // block_size)
-    return (2 * _walk_toffoli_count(high_bits, num_blocks) + 2 * output_bits *
-            (block_size - 1))
+    return (2 * _walk_toffoli_count(high_bits, num_blocks, controlled) +
+            2 * output_bits * (block_size - 1))
 
 
 class QROM:
@@ -181,8 +208,20 @@ class QROM:
         ``[2, 2^(address_bits - 1)]``. Only valid with
         ``variant="select_swap"``; ``None`` there picks the power of two
         minimizing the documented Toffoli count.
-
-    The three views of the minted kernel may live anywhere (no contiguity
+    controlled
+        ``True`` mints the externally controlled lookup: the kernel takes
+        a leading one-qubit ``control`` view, is the identity at control
+        ``|0>`` and the plain lookup at ``|1>`` (module docstring). The
+        identity at ``|0>`` holds on the clean ladder the kernel requires
+        anyway; a dirty ladder is outside the contract either way.
+        ``variant="auto"`` prices the controlled candidates against each
+        other on Toffolis alone, and the optimum can sit at a larger block
+        than the plain lookup's, with more block registers: size the
+        ladder from this instance's ``num_ladder`` rather than assume the
+        plain and controlled lookups share a layout (seven one-bit entries
+        over three address bits take 5 ladder qubits controlled, 4 plain).
+
+    The views of the minted kernel may live anywhere (no contiguity
     requirement between them), which is what lets callers weave the QROM
     into their own ancilla layouts. Whatever the variant, the kernel is
     exactly self-inverse — apply it again to uncompute.
@@ -194,7 +233,11 @@ def __init__(self,
                  output_bits: int,
                  *,
                  variant: str = "auto",
-                 block_size: int | None = None) -> None:
+                 block_size: int | None = None,
+                 controlled: bool = False) -> None:
+        if not isinstance(controlled, bool):
+            raise ValueError("controlled must be a bool (pass bool(flag) "
+                             "for numpy booleans)")
         if int(address_bits) != address_bits or address_bits < 1:
             raise ValueError("address_bits must be a positive integer")
         if int(output_bits) != output_bits or output_bits < 1:
@@ -244,12 +287,14 @@ def __init__(self,
                     f"index address bit: max is 2^(address_bits - 1) = "
                     f"{1 << (address_bits - 1)}")
 
-        select_cost = _walk_toffoli_count(address_bits, len(entries))
+        select_cost = _walk_toffoli_count(address_bits, len(entries),
+                                          controlled)
         swap_candidates = [1 << s for s in range(1, address_bits)
                            ] if block_size is None else [block_size]
         swap_costs = {
-            size: _select_swap_cost(address_bits, len(entries), output_bits,
-                                    size)
+            size:
+            _select_swap_cost(address_bits, len(entries), output_bits, size,
+                              controlled)
             for size in swap_candidates
         }
         if variant == "auto":
@@ -271,6 +316,7 @@ def __init__(self,
         self._output_bits = output_bits
         self._variant = variant
         self._block_size = block_size
+        self._controlled = controlled
         if variant == "select":
             self._build_select()
         else:
@@ -288,6 +334,7 @@ def body(k: int) -> list[tuple[str, int]]:
         walk = unary_iteration_kernels(self._num_address,
                                        len(entries),
                                        body,
+                                       controlled=self._controlled,
                                        include_adjoint=False)
         self._kernel = walk.kernel
         self._num_ladder = walk.num_ladder
@@ -315,9 +362,13 @@ def block_write(j: int) -> list[tuple[str, int]]:
         # registers (the leaf action writes each visited block), then a
         # low-bit-controlled swap network routes the selected entry.
         # X-only body: the walk is an involution, skip the adjoint mint.
+        # Only the walk passes carry the external control (module
+        # docstring): at control |0> nothing is written, so the routing
+        # and copy below act on |0> blocks as the identity.
         walk = unary_iteration_kernels(high_bits,
                                        num_blocks,
                                        block_write,
+                                       controlled=self._controlled,
                                        include_adjoint=False)
         walk_kernel = walk.kernel
 
@@ -343,24 +394,50 @@ def block_write(j: int) -> list[tuple[str, int]]:
         # W S C S^-1 W: write, route, copy, unroute, unwrite — clean
         # ancillas and an exactly self-inverse lookup. The parent carves
         # the documented sub-views out of the public (address, ladder,
-        # output) signature and composes the sub-kernels.
-        @cudaq.kernel
-        def primitives_qrom_select_swap(address: cudaq.qview,
-                                        ladder: cudaq.qview,
-                                        output: cudaq.qview):
-            low = address[0:low_bits]
-            high = address[low_bits:]
-            lines = ladder[0:high_bits]
-            blocks = ladder[high_bits:]
-            walk_kernel(high, lines, blocks)
-            route_kernel(low, blocks)
-            for t in range(b):
-                cx(blocks[t], output[t])
-            unroute_kernel(low, blocks)
-            walk_kernel(high, lines, blocks)
-
-        _retain(primitives_qrom_select_swap)
-        self._kernel = primitives_qrom_select_swap
+        # output) signature and composes the sub-kernels. Two parents on
+        # purpose: they differ only in the walk calls and share the route
+        # and copy sub-kernels; the plain one keeps the three-view
+        # signature of the composition contract, and a dummy control would
+        # cost every plain caller a qubit.
+        if self._controlled:
+
+            @cudaq.kernel
+            def primitives_qrom_select_swap_ctrl(control: cudaq.qview,
+                                                 address: cudaq.qview,
+                                                 ladder: cudaq.qview,
+                                                 output: cudaq.qview):
+                low = address[0:low_bits]
+                high = address[low_bits:]
+                lines = ladder[0:high_bits]
+                blocks = ladder[high_bits:]
+                walk_kernel(control, high, lines, blocks)
+                route_kernel(low, blocks)
+                for t in range(b):
+                    cx(blocks[t], output[t])
+                unroute_kernel(low, blocks)
+                walk_kernel(control, high, lines, blocks)
+
+            _retain(primitives_qrom_select_swap_ctrl)
+            self._kernel = primitives_qrom_select_swap_ctrl
+        else:
+
+            @cudaq.kernel
+            def primitives_qrom_select_swap(address: cudaq.qview,
+                                            ladder: cudaq.qview,
+                                            output: cudaq.qview):
+                low = address[0:low_bits]
+                high = address[low_bits:]
+                lines = ladder[0:high_bits]
+                blocks = ladder[high_bits:]
+                walk_kernel(high, lines, blocks)
+                route_kernel(low, blocks)
+                for t in range(b):
+                    cx(blocks[t], output[t])
+                unroute_kernel(low, blocks)
+                walk_kernel(high, lines, blocks)
+
+            _retain(primitives_qrom_select_swap)
+            self._kernel = primitives_qrom_select_swap
         self._num_ladder = high_bits + size * b
         self._toffoli_count = 2 * walk.toffoli_count + 2 * len(swap_c)
 
@@ -400,9 +477,15 @@ def block_size(self) -> int | None:
     def toffoli_count(self) -> int:
         return self._toffoli_count
 
+    @property
+    def controlled(self) -> bool:
+        """Whether the kernel takes a leading one-qubit ``control`` view."""
+        return self._controlled
+
     @property
     def kernel(self):
-        """The lookup kernel ``(address, ladder, output)`` — self-inverse."""
+        """The lookup kernel — self-inverse. ``(address, ladder, output)``,
+        or ``(control, address, ladder, output)`` when ``controlled``."""
         return self._kernel
 
     def __repr__(self) -> str:
@@ -411,4 +494,5 @@ def __repr__(self) -> str:
                 f"output_bits={self.num_output}, "
                 f"variant={self._variant!r}, "
                 f"block_size={self._block_size}, "
+                f"controlled={self._controlled}, "
                 f"toffolis={self.toffoli_count})")