Add Always metabloq to simplify decompositions with compute-uncompute pairs

dev_tools/qualtran_dev_tools/notebook_specs.py

@@ -251,6 +251,7 @@
             qualtran.bloqs.bookkeeping.partition._PARTITION_DOC,
             qualtran.bloqs.bookkeeping.auto_partition._AUTO_PARTITION_DOC,
             qualtran.bloqs.bookkeeping.cast._CAST_DOC,
+            qualtran.bloqs.bookkeeping.always._ALWAYS_DOC,
         ],
     ),
     NotebookSpecV2(

qualtran/bloqs/arithmetic/addition.ipynb

@@ -325,8 +325,7 @@
     "\n",
     "#### Parameters\n",
     " - `dtype`: data type of the input register `x`\n",
-    " - `k`: The classical integer value to be added to x.\n",
-    " - `is_controlled`: if True, construct a singly-controlled bloq. \n",
+    " - `k`: The classical integer value to be added to x. \n",
     "\n",
     "#### Registers\n",
     " - `x`: register of type `self.dtype` \n",

qualtran/bloqs/arithmetic/addition.py

@@ -15,7 +15,6 @@
 from functools import cached_property
 from typing import Dict, Iterator, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
-import attrs
 import cirq
 import numpy as np
 import sympy
@@ -32,7 +31,6 @@
     CtrlSpec,
     DecomposeTypeError,
     GateWithRegisters,
-    QAny,
     QInt,
     QMontgomeryUInt,
     QUInt,
@@ -43,6 +41,7 @@
     SoquetT,
 )
 from qualtran.bloqs.basic_gates import CNOT
+from qualtran.bloqs.bookkeeping import Always
 from qualtran.bloqs.mcmt.and_bloq import And
 from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv
 from qualtran.cirq_interop import decompose_from_cirq_style_method
@@ -404,7 +403,6 @@ class AddK(Bloq):
     Args:
         dtype: data type of the input register `x`
         k: The classical integer value to be added to x.
-        is_controlled: if True, construct a singly-controlled bloq.
 
     Registers:
         x: register of type `self.dtype`
@@ -416,7 +414,6 @@ class AddK(Bloq):
 
     dtype: Union[QInt, QUInt, QMontgomeryUInt]
     k: 'SymbolicInt'
-    is_controlled: bool = False
 
     def __attrs_post_init__(self):
         if not isinstance(self.dtype, (QInt, QUInt, QMontgomeryUInt)):
@@ -426,19 +423,18 @@ def __attrs_post_init__(self):
 
     @cached_property
     def signature(self) -> 'Signature':
-        return Signature.build_from_dtypes(ctrl=QAny(1 if self.is_controlled else 0), x=self.dtype)
+        return Signature.build_from_dtypes(x=self.dtype)
 
     def on_classical_vals(
         self, x: 'ClassicalValT', **vals: 'ClassicalValT'
     ) -> Dict[str, 'ClassicalValT']:
         if is_symbolic(self.k) or is_symbolic(self.dtype):
             raise ValueError(f"Classical simulation isn't supported for symbolic block {self}")
 
-        if not self.is_controlled or vals['ctrl']:
-            is_signed = isinstance(self.dtype, QInt)
-            x = add_ints(int(x), int(self.k), num_bits=self.dtype.num_qubits, is_signed=is_signed)
+        is_signed = isinstance(self.dtype, QInt)
+        x = add_ints(int(x), int(self.k), num_bits=self.dtype.num_qubits, is_signed=is_signed)
 
-        return vals | {'x': x}
+        return {'x': x}
 
     @cached_property
     def _load_k_bloq(self) -> Bloq:
@@ -449,55 +445,36 @@ def _load_k_bloq(self) -> Bloq:
             # Since this is unsigned addition, adding `-v` is equivalent to adding `2**bitsize - v`
             k %= 2**self.dtype.bitsize
 
-        xork = XorK(self.dtype, k)
-        return xork.controlled() if self.is_controlled else xork
+        return XorK(self.dtype, k)
 
-    def build_composite_bloq(
-        self, bb: 'BloqBuilder', x: Soquet, **soqs: Soquet
-    ) -> Dict[str, 'SoquetT']:
+    def build_composite_bloq(self, bb: 'BloqBuilder', x: Soquet) -> Dict[str, 'SoquetT']:
         if is_symbolic(self.k) or is_symbolic(self.dtype):
             raise DecomposeTypeError(f"Cannot decompose symbolic {self}.")
 
-        # load `k` (conditional on ctrl if present)
+        # load `k`
         k = bb.allocate(dtype=self.dtype)
-        load_soqs = {'x': k}
-        if self.is_controlled:
-            load_soqs |= {'ctrl': soqs.pop('ctrl')}
-        load_soqs = bb.add_d(self._load_k_bloq, **load_soqs)
-        k = load_soqs.pop('x')
+        k = bb.add(self._load_k_bloq, x=k)
 
-        # quantum-quantum addition
-        k, x = bb.add(Add(self.dtype, self.dtype), a=k, b=x)
+        # perform the quantum-quantum addition
+        # we always perform this addition (even when controlled), so we wrap in `Always`
+        # controlling the data loading is sufficient to control this bloq.
+        k, x = bb.add(Always(Add(self.dtype, self.dtype)), a=k, b=x)
 
         # unload `k`
-        load_soqs['x'] = k
-        load_soqs = bb.add_d(self._load_k_bloq.adjoint(), **load_soqs)
-        k = load_soqs.pop('x')
-        assert isinstance(k, Soquet)
+        k = bb.add(self._load_k_bloq.adjoint(), x=k)
         bb.free(k)
 
-        return {'x': x} | load_soqs
+        return {'x': x}
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         counts = Counter[Bloq]()
 
         counts[self._load_k_bloq] += 1
-        counts[Add(self.dtype, self.dtype)] += 1
+        counts[Always(Add(self.dtype, self.dtype))] += 1
         counts[self._load_k_bloq.adjoint()] += 1
 
         return counts
 
-    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
-        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
-
-        return get_ctrl_system_1bit_cv_from_bloqs(
-            self,
-            ctrl_spec,
-            current_ctrl_bit=1 if self.is_controlled else None,
-            bloq_with_ctrl=attrs.evolve(self, is_controlled=True),
-            ctrl_reg_name='ctrl',
-        )
-
 
 @bloq_example(generalizer=ignore_split_join)
 def _add_k() -> AddK:

qualtran/bloqs/arithmetic/addition_test.py

@@ -317,7 +317,6 @@ def test_out_of_place_adder():
 def test_controlled_add_k():
     n, k = sympy.symbols('n k')
     addk = AddK(QUInt(n), k)
-    assert addk.controlled() == AddK(QUInt(n), k, is_controlled=True)
     _, sigma = addk.controlled(CtrlSpec(cvs=0)).call_graph(max_depth=1)
     assert sigma == {addk.controlled(): 1, XGate(): 2}
 
@@ -346,7 +345,7 @@ def test_add_k_decomp_signed(bitsize, k, cvs):
     'bitsize,k,x,cvs,ctrls,result',
     [
         (5, 1, 2, (), (), 3),
-        (5, 3, 2, (1,), 1, 5),
+        (5, 3, 2, (1,), (1,), 5),
         (5, 2, 0, (1, 0), (1, 0), 2),
         (5, 1, 2, (1, 0, 1), (0, 0, 0), 2),
     ],

qualtran/bloqs/bookkeeping/__init__.py

@@ -14,6 +14,7 @@
 """Bloqs for virtual operations and register reshaping."""
 
 from qualtran.bloqs.bookkeeping.allocate import Allocate
+from qualtran.bloqs.bookkeeping.always import Always
 from qualtran.bloqs.bookkeeping.arbitrary_clifford import ArbitraryClifford
 from qualtran.bloqs.bookkeeping.auto_partition import AutoPartition
 from qualtran.bloqs.bookkeeping.cast import Cast

qualtran/bloqs/bookkeeping/always.py

@@ -0,0 +1,99 @@
+#  Copyright 2025 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.
+from typing import Iterable, Optional, Sequence
+
+import attrs
+
+from qualtran import (
+    AddControlledT,
+    Bloq,
+    bloq_example,
+    BloqBuilder,
+    BloqDocSpec,
+    CtrlSpec,
+    Signature,
+    SoquetT,
+)
+
+
+@attrs.frozen
+class Always(Bloq):
+    """Always execute the wrapped bloq, even when a controlled version is requested
+
+    A controlled version of a composite bloq in turn controls each subbloq in the decomposition.
+    Wrapping a particular subbloq with `Always` lets it bypass the controls,
+    i.e. it is "always" executed, irrespective of what the controls are.
+
+    This is useful when writing decompositions for two known patterns:
+
+    1. Compute-uncompute pairs: If a decomposition contains a compute-uncompute pair,
+    then for a controlled version, we only need to control the rest of the bloqs.
+    Wrapping both the compute and uncompute bloqs in `Always` lets them bypass the controls.
+
+    2. Controlled data-loading: For example, in the `AddK` bloq which adds a constant `k` to the
+    register, we (controlled) load the value `k` into a quantum register, and "always" perform an
+    quantum-quantum addition using `Add`, and unload `k`. Here wrapping the middle `Add` with
+    `Always` lets it bypass controls, e.g. when using `AddK.controlled()`.
+
+    This simplifies the decompositions by avoiding the need to explicitly define the decomposition
+    for the controlled version of bloq.
+
+    **Caution:** This wrapper should be used with care. It is up to the bloq author to ensure that
+    the controlled version of a decomposition containing `Always` bloqs still respects the
+    controlled protocol. That is, ignoring controls on these subbloqs wrapped in `Always` should not
+    change the action of the overall bloq with respect to the reference controlled implementation.
+
+    Args:
+        subbloq: The bloq to always apply, irrespective of any controls.
+    """
+
+    subbloq: Bloq
+
+    @property
+    def signature(self) -> 'Signature':
+        return self.subbloq.signature
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> dict[str, 'SoquetT']:
+        return bb.add_d(self.subbloq, **soqs)
+
+    def get_ctrl_system(
+        self, ctrl_spec: Optional['CtrlSpec'] = None
+    ) -> tuple['Bloq', 'AddControlledT']:
+        """Pass-through the control registers as-is"""
+
+        def add_controlled(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: dict[str, 'SoquetT']
+        ) -> tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            out_soqs = bb.add_t(self, **in_soqs)
+            return ctrl_soqs, out_soqs
+
+        return self, add_controlled
+
+    def adjoint(self) -> 'Always':
+        return Always(self.subbloq.adjoint())
+
+    def __str__(self) -> str:
+        return str(self.subbloq)
+
+
+@bloq_example
+def _always_and() -> Always:
+    from qualtran.bloqs.mcmt.and_bloq import And
+
+    always_and = Always(And())
+
+    return always_and
+
+
+_ALWAYS_DOC = BloqDocSpec(bloq_cls=Always, examples=[_always_and])

qualtran/bloqs/bookkeeping/always_test.py

@@ -0,0 +1,24 @@
+#  Copyright 2025 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.
+from qualtran.bloqs.bookkeeping.always import _always_and, Always
+from qualtran.bloqs.for_testing import TestAtom
+
+
+def test_example(bloq_autotester):
+    bloq_autotester(_always_and)
+
+
+def test_always():
+    bloq = Always(TestAtom())
+    assert bloq.controlled() == bloq