Y, Z: explicitly implement ctrl sys (avoids AutoPartition in decomp)

qualtran/bloqs/basic_gates/y_gate.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 from functools import cached_property
-from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, Union
+from typing import cast, Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 from attrs import frozen
@@ -22,14 +22,16 @@
     AddControlledT,
     Bloq,
     bloq_example,
+    BloqBuilder,
     BloqDocSpec,
     CompositeBloq,
     ConnectionT,
     CtrlSpec,
     DecomposeTypeError,
-    QBit,
     Register,
     Signature,
+    Soquet,
+    SoquetT,
 )
 from qualtran.drawing import Circle, Text, TextBox, WireSymbol
 
@@ -74,16 +76,20 @@ def my_tensors(
         ]
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
-        from qualtran.bloqs.bookkeeping import AutoPartition
-        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
+        if ctrl_spec != CtrlSpec():
+            return super().get_ctrl_system(ctrl_spec)
 
-        cy_wrapped = AutoPartition(
-            CYGate(), [(Register('ctrl', QBit()), ['ctrl']), (Register('q', QBit()), ['target'])]
-        )
+        cy = CYGate()
 
-        return get_ctrl_system_1bit_cv_from_bloqs(
-            self, ctrl_spec, current_ctrl_bit=None, bloq_with_ctrl=cy_wrapped, ctrl_reg_name='ctrl'
-        )
+        def _adder(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: dict[str, 'SoquetT']
+        ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            (ctrl,) = ctrl_soqs
+            target = in_soqs.pop('q')
+            ctrl, target = bb.add(cy, ctrl=ctrl, target=target)
+            return [ctrl], [target]
+
+        return cy, _adder
 
     def as_cirq_op(
         self, qubit_manager: 'cirq.QubitManager', q: 'CirqQuregT'
@@ -170,17 +176,32 @@ def wire_symbol(
         raise ValueError(f"Unknown register {reg}.")
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
-        from qualtran.bloqs.bookkeeping import AutoPartition
-        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
+        from qualtran.bloqs.mcmt.specialized_ctrl import _MultiControlledFromSinglyControlled
 
-        cy_wrapped = AutoPartition(
-            self, [(Register('ctrl', QBit()), ['ctrl']), (Register('q', QBit()), ['target'])]
-        )
+        if ctrl_spec != CtrlSpec():
+            return super().get_ctrl_system(ctrl_spec)
 
-        return get_ctrl_system_1bit_cv_from_bloqs(
-            self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=cy_wrapped, ctrl_reg_name='ctrl'
+        # controlled-CY
+        ctrl_cy = _MultiControlledFromSinglyControlled(
+            cvs=(1, 1), ctrl_bloq=self, ctrl_reg_name='ctrl'
         )
 
+        def _adder(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: dict[str, 'SoquetT']
+        ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            (ctrl1,) = ctrl_soqs
+            ctrl2 = in_soqs.pop('ctrl')
+            target = in_soqs.pop('target')
+
+            ctrl1 = cast(Soquet, ctrl1)
+            ctrl2 = cast(Soquet, ctrl2)
+
+            [ctrl1, ctrl2], target = bb.add(ctrl_cy, ctrl=[ctrl1, ctrl2], target=target)
+
+            return [ctrl1], [ctrl2, target]
+
+        return ctrl_cy, _adder
+
 
 @bloq_example
 def _cy_gate() -> CYGate:

qualtran/bloqs/basic_gates/z_basis.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 from functools import cached_property
-from typing import cast, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
+from typing import cast, Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import numpy as np
@@ -254,17 +254,20 @@ def my_tensors(
         ]
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
-        from qualtran.bloqs.bookkeeping import AutoPartition
-        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
+        if ctrl_spec != CtrlSpec():
+            return super().get_ctrl_system(ctrl_spec)
 
         cz = CZ()
-        cz_wrapped = AutoPartition(
-            cz, [(Register('ctrl', QBit()), ['q1']), (Register('q', QBit()), ['q2'])]
-        )
 
-        return get_ctrl_system_1bit_cv_from_bloqs(
-            self, ctrl_spec, current_ctrl_bit=None, bloq_with_ctrl=cz_wrapped, ctrl_reg_name='ctrl'
-        )
+        def _adder(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: dict[str, 'SoquetT']
+        ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            (ctrl,) = ctrl_soqs
+            target = in_soqs.pop('q')
+            ctrl, target = bb.add(cz, q1=ctrl, q2=target)
+            return [ctrl], [target]
+
+        return cz, _adder
 
     def as_cirq_op(
         self, qubit_manager: 'cirq.QubitManager', q: 'CirqQuregT'
@@ -338,17 +341,32 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
         raise ValueError(f'Unknown wire symbol register name: {reg.name}')
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
-        from qualtran.bloqs.bookkeeping import AutoPartition
-        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
+        from qualtran.bloqs.mcmt.specialized_ctrl import _MultiControlledFromSinglyControlled
 
-        cz_wrapped = AutoPartition(
-            self, [(Register('ctrl', QBit()), ['q1']), (Register('q', QBit()), ['q2'])]
-        )
+        if ctrl_spec != CtrlSpec():
+            return super().get_ctrl_system(ctrl_spec)
 
-        return get_ctrl_system_1bit_cv_from_bloqs(
-            self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=cz_wrapped, ctrl_reg_name='ctrl'
+        # controlled-CZ
+        ctrl_cz = _MultiControlledFromSinglyControlled(
+            cvs=(1, 1), ctrl_bloq=self, ctrl_reg_name='q1'
         )
 
+        def _adder(
+            bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: dict[str, 'SoquetT']
+        ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+            (ctrl,) = ctrl_soqs
+            q1 = in_soqs.pop('q1')
+            q2 = in_soqs.pop('q2')
+
+            ctrl = cast(Soquet, ctrl)
+            q1 = cast(Soquet, q1)
+
+            [ctrl, q1], q2 = bb.add(ctrl_cz, q1=[ctrl, q1], q2=q2)
+
+            return [ctrl], [q1, q2]
+
+        return ctrl_cz, _adder
+
 
 @bloq_example
 def _cz() -> CZ: