Permutation Bloq (#1110)

* Permutation Bloq

* address comments:
use `x` instead of `q`
invert code in `build_call_graph`
add note for neg-ctrl CNOT

* move `XorK` to `bitwise.py`

* tests

* XorK bloq examples

* undo trivial changes

* mypy

* cleanup EqualsAConstant

* fix decomp

* use dictionaries for sparse/partial permutations, simplify code

* fix bug

* regen notebooks

* docstring

* add notebook

* docstring unitary

dev_tools/autogenerate-bloqs-notebooks-v2.py

@@ -53,6 +53,7 @@
 from qualtran_dev_tools.jupyter_autogen import NotebookSpecV2, render_notebook
 
 import qualtran.bloqs.arithmetic.addition
+import qualtran.bloqs.arithmetic.permutation
 import qualtran.bloqs.arithmetic.sorting
 import qualtran.bloqs.arithmetic.subtraction
 import qualtran.bloqs.basic_gates.swap
@@ -380,6 +381,14 @@
             qualtran.bloqs.arithmetic.conversions._TO_CONTG_INDX,
         ],
     ),
+    NotebookSpecV2(
+        title='Permutations',
+        module=qualtran.bloqs.arithmetic.permutation,
+        bloq_specs=[
+            qualtran.bloqs.arithmetic.permutation._PERMUTATION_DOC,
+            qualtran.bloqs.arithmetic.permutation._PERMUTATION_CYCLE_DOC,
+        ],
+    ),
 ]
 
 MOD_ARITHMETIC = [

docs/bloqs/index.rst

@@ -64,6 +64,7 @@ Bloqs Library
     arithmetic/comparison.ipynb
     arithmetic/sorting.ipynb
     arithmetic/conversions.ipynb
+    arithmetic/permutation.ipynb
 
 .. toctree::
     :maxdepth: 2

qualtran/bloqs/arithmetic/bitwise.py

@@ -0,0 +1,120 @@
+#  Copyright 2024 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 functools import cached_property
+from typing import cast, Dict, Optional, Set, Tuple, TYPE_CHECKING
+
+import numpy as np
+from attrs import frozen
+
+from qualtran import (
+    bloq_example,
+    BloqBuilder,
+    DecomposeTypeError,
+    QBit,
+    QDType,
+    QUInt,
+    Register,
+    Signature,
+    Soquet,
+    SoquetT,
+)
+from qualtran._infra.gate_with_registers import SpecializedSingleQubitControlledGate
+from qualtran.bloqs.basic_gates import CNOT, XGate
+from qualtran.resource_counting.generalizers import ignore_split_join
+from qualtran.symbolics import is_symbolic, SymbolicInt
+
+if TYPE_CHECKING:
+    from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+
+
+def _cvs_converter(vv):
+    if isinstance(vv, (int, np.integer)):
+        return (int(vv),)
+    return tuple(int(v) for v in vv)
+
+
+@frozen
+class XorK(SpecializedSingleQubitControlledGate):
+    r"""Maps |x> to |x \oplus k> for a constant k.
+
+    Args:
+        dtype: Data type of the input register `x`.
+        k: The classical integer value to be XOR-ed to x.
+        control_val: an optional single bit control, apply the operation when
+                    the control qubit equals the `control_val`.
+
+    Registers:
+        x: A quantum register of type `self.dtype` (see above).
+        ctrl: A sequence of control qubits (only when `control_val` is not None).
+    """
+    dtype: QDType
+    k: SymbolicInt
+    control_val: Optional[int] = None
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature([*self.control_registers, Register('x', self.dtype)])
+
+    @cached_property
+    def control_registers(self) -> Tuple[Register, ...]:
+        if self.control_val is not None:
+            return (Register('ctrl', QBit()),)
+        return ()
+
+    @cached_property
+    def bitsize(self) -> SymbolicInt:
+        return self.dtype.num_qubits
+
+    def is_symbolic(self):
+        return is_symbolic(self.k, self.dtype)
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str, 'SoquetT']:
+        if self.is_symbolic():
+            raise DecomposeTypeError(f"cannot decompose symbolic {self}")
+
+        # TODO clean this up once https://github.com/quantumlib/Qualtran/pull/1137 is merged
+        ctrl = soqs.pop('ctrl', None)
+
+        xs = bb.split(cast(Soquet, soqs.pop('x')))
+
+        for i, bit in enumerate(self.dtype.to_bits(self.k)):
+            if bit == 1:
+                if ctrl is not None:
+                    ctrl, xs[i] = bb.add(CNOT(), ctrl=ctrl, target=xs[i])
+                else:
+                    xs[i] = bb.add(XGate(), q=xs[i])
+
+        soqs['x'] = bb.join(xs)
+
+        if ctrl is not None:
+            soqs['ctrl'] = ctrl
+        return soqs
+
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
+        bit_flip_bloq = CNOT() if self.control_val is not None else XGate()
+        num_flips = self.bitsize if self.is_symbolic() else sum(self.dtype.to_bits(self.k))
+        return {(bit_flip_bloq, num_flips)}
+
+
+@bloq_example(generalizer=ignore_split_join)
+def _xork() -> XorK:
+    xork = XorK(QUInt(8), 0b01010111)
+    return xork
+
+
+@bloq_example(generalizer=ignore_split_join)
+def _cxork() -> XorK:
+    cxork = XorK(QUInt(8), 0b01010111).controlled()
+    assert isinstance(cxork, XorK)
+    return cxork

qualtran/bloqs/arithmetic/bitwise_test.py

@@ -0,0 +1,38 @@
+#  Copyright 2024 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 import QUInt
+from qualtran.bloqs.arithmetic.bitwise import _cxork, _xork, XorK
+
+
+def test_examples(bloq_autotester):
+    bloq_autotester(_cxork)
+    bloq_autotester(_xork)
+
+
+def test_xork_classical_sim():
+    k = 0b01101010
+    bloq = XorK(QUInt(9), k)
+    cbloq = bloq.controlled()
+
+    for x in bloq.dtype.get_classical_domain():
+        (x_out,) = bloq.call_classically(x=x)
+        assert x_out == x ^ k
+
+        ctrl_out, x_out = cbloq.call_classically(ctrl=0, x=x)
+        assert ctrl_out == 0
+        assert x_out == x
+
+        ctrl_out, x_out = cbloq.call_classically(ctrl=1, x=x)
+        assert ctrl_out == 1
+        assert x_out == x ^ k

qualtran/bloqs/arithmetic/comparison.ipynb

@@ -256,7 +256,7 @@
    },
    "source": [
     "## `EqualsAConstant`\n",
-    "Implements $U_a|x\\rangle = U_a|x\\rangle|z\\rangle = |x\\rangle |z \\land (x = a)\\rangle$\n",
+    "Implements $U_a|x\\rangle|z\\rangle = |x\\rangle |z \\oplus (x = a)\\rangle$\n",
     "\n",
     "The bloq_counts and t_complexity are derived from:\n",
     "https://qualtran.readthedocs.io/en/latest/bloqs/comparison_gates.html#equality-as-a-special-case\n",

qualtran/bloqs/arithmetic/comparison.py

@@ -38,6 +38,7 @@
     Bloq,
     bloq_example,
     BloqDocSpec,
+    DecomposeTypeError,
     GateWithRegisters,
     QAny,
     QBit,
@@ -48,12 +49,12 @@
     Soquet,
     SoquetT,
 )
-from qualtran.bloqs.basic_gates import CNOT, TGate, XGate
+from qualtran.bloqs.basic_gates import CNOT, XGate
 from qualtran.bloqs.mcmt.and_bloq import And, MultiAnd
-from qualtran.bloqs.mcmt.multi_control_multi_target_pauli import MultiControlX
+from qualtran.bloqs.mcmt.multi_control_multi_target_pauli import MultiControlPauli, MultiControlX
 from qualtran.drawing import WireSymbol
 from qualtran.drawing.musical_score import Text, TextBox
-from qualtran.symbolics import is_symbolic, SymbolicInt
+from qualtran.symbolics import HasLength, is_symbolic, SymbolicInt
 
 if TYPE_CHECKING:
     from qualtran import BloqBuilder
@@ -926,7 +927,7 @@ def _gt_k() -> GreaterThanConstant:
 
 @frozen
 class EqualsAConstant(Bloq):
-    r"""Implements $U_a|x\rangle = U_a|x\rangle|z\rangle = |x\rangle |z \land (x = a)\rangle$
+    r"""Implements $U_a|x\rangle|z\rangle = |x\rangle |z \oplus (x = a)\rangle$
 
     The bloq_counts and t_complexity are derived from:
     https://qualtran.readthedocs.io/en/latest/bloqs/comparison_gates.html#equality-as-a-special-case
@@ -940,8 +941,8 @@ class EqualsAConstant(Bloq):
         target: Register to hold result of comparison.
     """
 
-    bitsize: int
-    val: int
+    bitsize: SymbolicInt
+    val: SymbolicInt
 
     @cached_property
     def signature(self) -> Signature:
@@ -956,10 +957,33 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
             return TextBox(f"⨁(x = {self.val})")
         raise ValueError(f'Unknown register symbol {reg.name}')
 
+    def is_symbolic(self):
+        return is_symbolic(self.bitsize, self.val)
+
+    @property
+    def bits_k(self) -> Union[tuple[int, ...], HasLength]:
+        if self.is_symbolic():
+            return HasLength(self.bitsize)
+
+        assert not isinstance(self.bitsize, sympy.Expr)
+        assert not isinstance(self.val, sympy.Expr)
+        return tuple(QUInt(self.bitsize).to_bits(self.val))
+
+    def build_composite_bloq(
+        self, bb: 'BloqBuilder', x: 'Soquet', target: 'Soquet'
+    ) -> Dict[str, 'SoquetT']:
+        if is_symbolic(self.bitsize):
+            raise DecomposeTypeError(f"Cannot decompose {self} with symbolic {self.bitsize=}")
+
+        xs = bb.split(x)
+        xs, target = bb.add(
+            MultiControlPauli(self.bits_k, target_gate=cirq.X), controls=xs, target=target
+        )
+        x = bb.join(xs)
+        return {'x': x, 'target': target}
+
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
-        # See: https://github.com/quantumlib/Qualtran/issues/219
-        # See: https://github.com/quantumlib/Qualtran/issues/217
-        return {(TGate(), 4 * (self.bitsize - 1))}
+        return {(MultiControlPauli(self.bits_k, target_gate=cirq.X), 1)}
 
 
 def _make_equals_a_constant():

qualtran/bloqs/arithmetic/comparison_test.py

@@ -298,7 +298,9 @@ def test_equals_a_constant():
     qlt_testing.assert_wire_symbols_match_expected(
         EqualsAConstant(bitsize, 17), ['In(x)', '⨁(x = 17)']
     )
-    assert t_complexity(EqualsAConstant(bitsize, 17)) == TComplexity(t=4 * (bitsize - 1))
+    assert t_complexity(EqualsAConstant(bitsize, 17)) == TComplexity(
+        t=4 * (bitsize - 1), clifford=65
+    )
 
 
 @pytest.mark.notebook