cirq-core 1.3.0.dev20231201141002__py3-none-any.whl → 1.4.0__py3-none-any.whl

cirq/transformers/optimize_for_target_gateset.py

@@ -14,7 +14,7 @@
 
 """Transformers to rewrite a circuit using gates from a given target gateset."""
 
-from typing import Optional, Callable, Hashable, Sequence, TYPE_CHECKING
+from typing import Optional, Callable, Hashable, Sequence, TYPE_CHECKING, Union
 
 from cirq import circuits
 from cirq.protocols import decompose_protocol as dp
@@ -102,19 +102,29 @@ def optimize_for_target_gateset(
     context: Optional['cirq.TransformerContext'] = None,
     gateset: Optional['cirq.CompilationTargetGateset'] = None,
     ignore_failures: bool = True,
+    max_num_passes: Union[int, None] = 1,
 ) -> 'cirq.Circuit':
     """Transforms the given circuit into an equivalent circuit using gates accepted by `gateset`.
 
+    Repeat max_num_passes times or when `max_num_passes=None` until no further changes can be done
     1. Run all `gateset.preprocess_transformers`
     2. Convert operations using built-in cirq decompose + `gateset.decompose_to_target_gateset`.
     3. Run all `gateset.postprocess_transformers`
 
+    Note:
+        The optimizer is a heuristic and may not produce optimal results even with
+        max_num_passes=None. The preprocessors and postprocessors of the gate set
+        as well as their order yield different results.
+
+
     Args:
         circuit: Input circuit to transform. It will not be modified.
         context: `cirq.TransformerContext` storing common configurable options for transformers.
         gateset: Target gateset, which should be an instance of `cirq.CompilationTargetGateset`.
         ignore_failures: If set, operations that fail to convert are left unchanged. If not set,
             conversion failures raise a ValueError.
+        max_num_passes: The maximum number of passes to do. A value of `None` means to keep
+            iterating until no more changes happen to the number of moments or operations.
 
     Returns:
         An equivalent circuit containing gates accepted by `gateset`.
@@ -126,20 +136,32 @@ def optimize_for_target_gateset(
         return _decompose_operations_to_target_gateset(
             circuit, context=context, ignore_failures=ignore_failures
         )
-
-    for transformer in gateset.preprocess_transformers:
-        circuit = transformer(circuit, context=context)
-
-    circuit = _decompose_operations_to_target_gateset(
-        circuit,
-        context=context,
-        gateset=gateset,
-        decomposer=gateset.decompose_to_target_gateset,
-        ignore_failures=ignore_failures,
-        tags_to_decompose=(gateset._intermediate_result_tag,),
-    )
-
-    for transformer in gateset.postprocess_transformers:
-        circuit = transformer(circuit, context=context)
-
+    if isinstance(max_num_passes, int):
+        _outerloop = lambda: range(max_num_passes)
+    else:
+
+        def _outerloop():
+            while True:
+                yield 0
+
+    initial_num_moments, initial_num_ops = len(circuit), sum(1 for _ in circuit.all_operations())
+    for _ in _outerloop():
+        for transformer in gateset.preprocess_transformers:
+            circuit = transformer(circuit, context=context)
+        circuit = _decompose_operations_to_target_gateset(
+            circuit,
+            context=context,
+            gateset=gateset,
+            decomposer=gateset.decompose_to_target_gateset,
+            ignore_failures=ignore_failures,
+            tags_to_decompose=(gateset._intermediate_result_tag,),
+        )
+        for transformer in gateset.postprocess_transformers:
+            circuit = transformer(circuit, context=context)
+
+        num_moments, num_ops = len(circuit), sum(1 for _ in circuit.all_operations())
+        if (num_moments, num_ops) == (initial_num_moments, initial_num_ops):
+            # Stop early. No further optimizations can be done.
+            break
+        initial_num_moments, initial_num_ops = num_moments, num_ops
     return circuit.unfreeze(copy=False)

cirq/transformers/optimize_for_target_gateset_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Union
+
 import cirq
 from cirq.protocols.decompose_protocol import DecomposeResult
 from cirq.transformers.optimize_for_target_gateset import _decompose_operations_to_target_gateset
@@ -51,22 +53,22 @@ def test_decompose_operations_to_target_gateset_default():
     cirq.testing.assert_has_diagram(
         c_orig,
         '''
-0: ───T───×───T───×['ignore']───M───────T───×───T───
-          │       │             ║           │
-1: ───────×───────×─────────────╫───X───T───×───T───
-                                ║   ║
-m: ═════════════════════════════@═══^═══════════════''',
+0: ───T───×───T───×[ignore]───M───────T───×───T───
+          │       │           ║           │
+1: ───────×───────×───────────╫───X───T───×───T───
+                              ║   ║
+m: ═══════════════════════════@═══^═══════════════''',
     )
     context = cirq.TransformerContext(tags_to_ignore=("ignore",))
     c_new = _decompose_operations_to_target_gateset(c_orig, context=context)
     cirq.testing.assert_has_diagram(
         c_new,
         '''
-0: ───T────────────@───Y^-0.5───@───Y^0.5────@───────────T───×['ignore']───M───────T────────────@───Y^-0.5───@───Y^0.5────@───────────T───
-                   │            │            │               │             ║                    │            │            │
-1: ───────Y^-0.5───@───Y^0.5────@───Y^-0.5───@───Y^0.5───────×─────────────╫───X───T───Y^-0.5───@───Y^0.5────@───Y^-0.5───@───Y^0.5───T───
-                                                                           ║   ║
-m: ════════════════════════════════════════════════════════════════════════@═══^══════════════════════════════════════════════════════════
+0: ───T────────────@───Y^-0.5───@───Y^0.5────@───────────T───×[ignore]───M───────T────────────@───Y^-0.5───@───Y^0.5────@───────────T───
+                   │            │            │               │           ║                    │            │            │
+1: ───────Y^-0.5───@───Y^0.5────@───Y^-0.5───@───Y^0.5───────×───────────╫───X───T───Y^-0.5───@───Y^0.5────@───Y^-0.5───@───Y^0.5───T───
+                                                                         ║   ║
+m: ══════════════════════════════════════════════════════════════════════@═══^══════════════════════════════════════════════════════════
 ''',
     )
 
@@ -85,8 +87,8 @@ def test_decompose_operations_to_target_gateset():
         cirq.T.on_each(*q),
     )
     gateset = cirq.Gateset(cirq.H, cirq.CNOT)
-    decomposer = (
-        lambda op, _: cirq.H(op.qubits[0])
+    decomposer = lambda op, _: (
+        cirq.H(op.qubits[0])
         if cirq.has_unitary(op) and cirq.num_qubits(op) == 1
         else NotImplemented
     )
@@ -97,11 +99,11 @@ def test_decompose_operations_to_target_gateset():
     cirq.testing.assert_has_diagram(
         c_new,
         '''
-0: ───H───@───X───@───H───×['ignore']───M───────H───@───X───@───H───
-          │   │   │       │             ║           │   │   │
-1: ───────X───@───X───────×─────────────╫───X───H───X───@───X───H───
-                                        ║   ║
-m: ═════════════════════════════════════@═══^═══════════════════════''',
+0: ───H───@───X───@───H───×[ignore]───M───────H───@───X───@───H───
+          │   │   │       │           ║           │   │   │
+1: ───────X───@───X───────×───────────╫───X───H───X───@───X───H───
+                                      ║   ║
+m: ═══════════════════════════════════@═══^═══════════════════════''',
     )
 
     with pytest.raises(ValueError, match="Unable to convert"):
@@ -134,9 +136,9 @@ def test_optimize_for_target_gateset_default():
     cirq.testing.assert_has_diagram(
         c_new,
         '''
-0: ───T────────────@───Y^-0.5───@───Y^0.5────@───────────T───×['ignore']───
+0: ───T────────────@───Y^-0.5───@───Y^0.5────@───────────T───×[ignore]───
                    │            │            │               │
-1: ───────Y^-0.5───@───Y^0.5────@───Y^-0.5───@───Y^0.5───────×─────────────
+1: ───────Y^-0.5───@───Y^0.5────@───Y^-0.5───@───Y^0.5───────×───────────
 ''',
     )
     cirq.testing.assert_circuits_with_terminal_measurements_are_equivalent(c_orig, c_new, atol=1e-6)
@@ -157,15 +159,15 @@ def test_optimize_for_target_gateset():
     cirq.testing.assert_has_diagram(
         c_orig,
         '''
-0: ───qft───Y['ignore']───M───────qft^-1───
-      │                   ║       │
-1: ───#2────Y['ignore']───M───────#2───────
-      │                   ║       │
-2: ───#3────@['ignore']───╫───@───#3───────
-      │     │             ║   ║   │
-3: ───#4────X─────────────╫───@───#4───────
-                          ║   ║
-m: ═══════════════════════@═══^════════════
+0: ───qft───Y[ignore]───M───────qft^-1───
+      │                 ║       │
+1: ───#2────Y[ignore]───M───────#2───────
+      │                 ║       │
+2: ───#3────@[ignore]───╫───@───#3───────
+      │     │           ║   ║   │
+3: ───#4────X───────────╫───@───#4───────
+                        ║   ║
+m: ═════════════════════@═══^════════════
 ''',
     )
     gateset = MatrixGateTargetGateset()
@@ -174,17 +176,17 @@ m: ═══════════════════════@══
     cirq.testing.assert_has_diagram(
         c_new,
         '''
-                                         ┌────────┐                         ┌────────┐                 ┌────────┐
-0: ───M[1]──────────M[1]──────────────────────M[1]────Y['ignore']───M────────────M[1]───────────────────M[1]────────M[1]───M[1]───
-      │             │                         │                     ║            │                      │           │      │
-1: ───M[2]───M[1]───┼─────────────M[1]────M[1]┼───────Y['ignore']───M────────M[1]┼──────────────M[1]────┼───M[1]────┼──────M[2]───
-             │      │             │       │   │                     ║        │   │              │       │   │       │
-2: ──────────M[2]───M[2]───M[1]───┼───────M[2]┼───────@['ignore']───╫───@────M[2]┼───────M[1]───┼───────┼───M[2]────M[2]──────────
-                           │      │           │       │             ║   ║        │       │      │       │
-3: ────────────────────────M[2]───M[2]────────M[2]────X─────────────╫───@────────M[2]────M[2]───M[2]────M[2]──────────────────────
-                                                                    ║   ║
-m: ═════════════════════════════════════════════════════════════════@═══^═════════════════════════════════════════════════════════
-                                         └────────┘                         └────────┘                 └────────┘
+                                         ┌────────┐                       ┌────────┐                 ┌────────┐
+0: ───M[1]──────────M[1]──────────────────────M[1]────Y[ignore]───M────────────M[1]───────────────────M[1]────────M[1]───M[1]───
+      │             │                         │                   ║            │                      │           │      │
+1: ───M[2]───M[1]───┼─────────────M[1]────M[1]┼───────Y[ignore]───M────────M[1]┼──────────────M[1]────┼───M[1]────┼──────M[2]───
+             │      │             │       │   │                   ║        │   │              │       │   │       │
+2: ──────────M[2]───M[2]───M[1]───┼───────M[2]┼───────@[ignore]───╫───@────M[2]┼───────M[1]───┼───────┼───M[2]────M[2]──────────
+                           │      │           │       │           ║   ║        │       │      │       │
+3: ────────────────────────M[2]───M[2]────────M[2]────X───────────╫───@────────M[2]────M[2]───M[2]────M[2]──────────────────────
+                                                                  ║   ║
+m: ═══════════════════════════════════════════════════════════════@═══^═════════════════════════════════════════════════════════
+                                         └────────┘                       └────────┘                 └────────┘
        ''',
     )
 
@@ -243,3 +245,151 @@ def test_optimize_for_target_gateset_deep():
 1: ───#2───────────────────────────────────────────────────────────────────────────
 ''',
     )
+
+
+@pytest.mark.parametrize('max_num_passes', [2, None])
+def test_optimize_for_target_gateset_multiple_passes(max_num_passes: Union[int, None]):
+    gateset = cirq.CZTargetGateset()
+
+    input_circuit = cirq.Circuit(
+        [
+            cirq.Moment(
+                cirq.X(cirq.LineQubit(1)),
+                cirq.X(cirq.LineQubit(2)),
+                cirq.X(cirq.LineQubit(3)),
+                cirq.X(cirq.LineQubit(6)),
+            ),
+            cirq.Moment(
+                cirq.H(cirq.LineQubit(0)),
+                cirq.H(cirq.LineQubit(1)),
+                cirq.H(cirq.LineQubit(2)),
+                cirq.H(cirq.LineQubit(3)),
+                cirq.H(cirq.LineQubit(4)),
+                cirq.H(cirq.LineQubit(5)),
+                cirq.H(cirq.LineQubit(6)),
+            ),
+            cirq.Moment(
+                cirq.H(cirq.LineQubit(1)), cirq.H(cirq.LineQubit(3)), cirq.H(cirq.LineQubit(5))
+            ),
+            cirq.Moment(
+                cirq.CZ(cirq.LineQubit(0), cirq.LineQubit(1)),
+                cirq.CZ(cirq.LineQubit(2), cirq.LineQubit(3)),
+                cirq.CZ(cirq.LineQubit(4), cirq.LineQubit(5)),
+            ),
+            cirq.Moment(
+                cirq.CZ(cirq.LineQubit(2), cirq.LineQubit(1)),
+                cirq.CZ(cirq.LineQubit(4), cirq.LineQubit(3)),
+                cirq.CZ(cirq.LineQubit(6), cirq.LineQubit(5)),
+            ),
+        ]
+    )
+    desired_circuit = cirq.Circuit.from_moments(
+        cirq.Moment(
+            cirq.PhasedXZGate(axis_phase_exponent=0.5, x_exponent=-0.5, z_exponent=1.0).on(
+                cirq.LineQubit(4)
+            )
+        ),
+        cirq.Moment(cirq.CZ(cirq.LineQubit(4), cirq.LineQubit(5))),
+        cirq.Moment(
+            cirq.PhasedXZGate(axis_phase_exponent=-1.0, x_exponent=1, z_exponent=0).on(
+                cirq.LineQubit(1)
+            ),
+            cirq.PhasedXZGate(axis_phase_exponent=0.5, x_exponent=-0.5, z_exponent=1.0).on(
+                cirq.LineQubit(0)
+            ),
+            cirq.PhasedXZGate(axis_phase_exponent=-1.0, x_exponent=1, z_exponent=0).on(
+                cirq.LineQubit(3)
+            ),
+            cirq.PhasedXZGate(axis_phase_exponent=-0.5, x_exponent=0.5, z_exponent=0.0).on(
+                cirq.LineQubit(2)
+            ),
+        ),
+        cirq.Moment(
+            cirq.CZ(cirq.LineQubit(0), cirq.LineQubit(1)),
+            cirq.CZ(cirq.LineQubit(2), cirq.LineQubit(3)),
+        ),
+        cirq.Moment(
+            cirq.CZ(cirq.LineQubit(2), cirq.LineQubit(1)),
+            cirq.CZ(cirq.LineQubit(4), cirq.LineQubit(3)),
+        ),
+        cirq.Moment(
+            cirq.PhasedXZGate(axis_phase_exponent=-0.5, x_exponent=0.5, z_exponent=0.0).on(
+                cirq.LineQubit(6)
+            )
+        ),
+        cirq.Moment(cirq.CZ(cirq.LineQubit(6), cirq.LineQubit(5))),
+    )
+    got = cirq.optimize_for_target_gateset(
+        input_circuit, gateset=gateset, max_num_passes=max_num_passes
+    )
+    cirq.testing.assert_same_circuits(got, desired_circuit)
+
+
+@pytest.mark.parametrize('max_num_passes', [2, None])
+def test_optimize_for_target_gateset_multiple_passes_dont_preserve_moment_structure(
+    max_num_passes: Union[int, None]
+):
+    gateset = cirq.CZTargetGateset(preserve_moment_structure=False)
+
+    input_circuit = cirq.Circuit(
+        [
+            cirq.Moment(
+                cirq.X(cirq.LineQubit(1)),
+                cirq.X(cirq.LineQubit(2)),
+                cirq.X(cirq.LineQubit(3)),
+                cirq.X(cirq.LineQubit(6)),
+            ),
+            cirq.Moment(
+                cirq.H(cirq.LineQubit(0)),
+                cirq.H(cirq.LineQubit(1)),
+                cirq.H(cirq.LineQubit(2)),
+                cirq.H(cirq.LineQubit(3)),
+                cirq.H(cirq.LineQubit(4)),
+                cirq.H(cirq.LineQubit(5)),
+                cirq.H(cirq.LineQubit(6)),
+            ),
+            cirq.Moment(
+                cirq.H(cirq.LineQubit(1)), cirq.H(cirq.LineQubit(3)), cirq.H(cirq.LineQubit(5))
+            ),
+            cirq.Moment(
+                cirq.CZ(cirq.LineQubit(0), cirq.LineQubit(1)),
+                cirq.CZ(cirq.LineQubit(2), cirq.LineQubit(3)),
+                cirq.CZ(cirq.LineQubit(4), cirq.LineQubit(5)),
+            ),
+            cirq.Moment(
+                cirq.CZ(cirq.LineQubit(2), cirq.LineQubit(1)),
+                cirq.CZ(cirq.LineQubit(4), cirq.LineQubit(3)),
+                cirq.CZ(cirq.LineQubit(6), cirq.LineQubit(5)),
+            ),
+        ]
+    )
+    desired_circuit = cirq.Circuit(
+        cirq.PhasedXZGate(axis_phase_exponent=0.5, x_exponent=-0.5, z_exponent=1.0).on(
+            cirq.LineQubit(4)
+        ),
+        cirq.PhasedXZGate(axis_phase_exponent=-1.0, x_exponent=1, z_exponent=0).on(
+            cirq.LineQubit(1)
+        ),
+        cirq.PhasedXZGate(axis_phase_exponent=-0.5, x_exponent=0.5, z_exponent=0.0).on(
+            cirq.LineQubit(2)
+        ),
+        cirq.PhasedXZGate(axis_phase_exponent=0.5, x_exponent=-0.5, z_exponent=1.0).on(
+            cirq.LineQubit(0)
+        ),
+        cirq.PhasedXZGate(axis_phase_exponent=-1.0, x_exponent=1, z_exponent=0).on(
+            cirq.LineQubit(3)
+        ),
+        cirq.PhasedXZGate(axis_phase_exponent=-0.5, x_exponent=0.5, z_exponent=0.0).on(
+            cirq.LineQubit(6)
+        ),
+        cirq.CZ(cirq.LineQubit(4), cirq.LineQubit(5)),
+        cirq.CZ(cirq.LineQubit(0), cirq.LineQubit(1)),
+        cirq.CZ(cirq.LineQubit(2), cirq.LineQubit(3)),
+        cirq.CZ(cirq.LineQubit(2), cirq.LineQubit(1)),
+        cirq.CZ(cirq.LineQubit(4), cirq.LineQubit(3)),
+        cirq.CZ(cirq.LineQubit(6), cirq.LineQubit(5)),
+    )
+    got = cirq.optimize_for_target_gateset(
+        input_circuit, gateset=gateset, max_num_passes=max_num_passes
+    )
+    cirq.testing.assert_same_circuits(got, desired_circuit)

cirq/transformers/qubit_management_transformers.py

@@ -71,7 +71,7 @@ def map_clean_and_borrowable_qubits(
     inserting all mapped operations in a resulting circuit using EARLIEST strategy. The reason
     is that preserving moment structure forces additional constraints on the qubit allocation
     strategy (i.e. if two operations `op1` and `op2` are in the same moment, then we cannot
-    reuse ancilla across `op1` and `op2`). We leave it upto the user to force such constraints
+    reuse ancilla across `op1` and `op2`). We leave it up to the user to force such constraints
     using the qubit manager instead of making it part of the transformer.
         3. However, for borrowable system qubits managed by the transformer, we do not reuse qubits
     within the same moment.

cirq/transformers/routing/visualize_routed_circuit_test.py

@@ -20,15 +20,15 @@ def test_routed_circuit_with_mapping_simple():
     q = cirq.LineQubit.range(2)
     circuit = cirq.Circuit([cirq.Moment(cirq.SWAP(q[0], q[1]).with_tags(cirq.RoutingSwapTag()))])
     expected_diagram = """
-0: ───q(0)───×[cirq.RoutingSwapTag()]───q(1)───
-      │      │                          │
-1: ───q(1)───×──────────────────────────q(0)───"""
+0: ───q(0)───×[<r>]───q(1)───
+      │      │        │
+1: ───q(1)───×────────q(0)───"""
     cirq.testing.assert_has_diagram(cirq.routed_circuit_with_mapping(circuit), expected_diagram)
 
     expected_diagram_with_initial_mapping = """
-0: ───a───×[cirq.RoutingSwapTag()]───b───
-      │   │                          │
-1: ───b───×──────────────────────────a───"""
+0: ───a───×[<r>]───b───
+      │   │        │
+1: ───b───×────────a───"""
     cirq.testing.assert_has_diagram(
         cirq.routed_circuit_with_mapping(
             circuit, {cirq.NamedQubit("a"): q[0], cirq.NamedQubit("b"): q[1]}
@@ -74,16 +74,16 @@ def test_routed_circuit_with_mapping_multi_swaps():
         ]
     )
     expected_diagram = """
-0: ───q(0)──────────────────────────────────────q(0)───×[cirq.RoutingSwapTag()]───q(1)───────X───
-      │                                         │      │                          │          │
-1: ───q(1)───────────X──────────────────────────q(1)───×──────────────────────────q(0)───X───@───
-      │              │                          │                                 │      │
-2: ───q(2)───────@───@──────────────────────────q(2)───×──────────────────────────q(4)───@───────
-      │          │                              │      │                          │
-3: ───q(3)───@───X───×──────────────────────────q(4)───×[cirq.RoutingSwapTag()]───q(2)───────────
-      │      │       │                          │                                 │
-4: ───q(4)───X───X───×[cirq.RoutingSwapTag()]───q(3)──────────────────────────────q(3)───────────
-      │          │                              │                                 │
-5: ───q(5)───────@──────────────────────────────q(5)──────────────────────────────q(5)───────────
+0: ───q(0)────────────────────q(0)───×[<r>]───q(1)───────X───
+      │                       │      │        │          │
+1: ───q(1)───────────X────────q(1)───×────────q(0)───X───@───
+      │              │        │               │      │
+2: ───q(2)───────@───@────────q(2)───×────────q(4)───@───────
+      │          │            │      │        │
+3: ───q(3)───@───X───×────────q(4)───×[<r>]───q(2)───────────
+      │      │       │        │               │
+4: ───q(4)───X───X───×[<r>]───q(3)────────────q(3)───────────
+      │          │            │               │
+5: ───q(5)───────@────────────q(5)────────────q(5)───────────
 """
     cirq.testing.assert_has_diagram(cirq.routed_circuit_with_mapping(circuit), expected_diagram)

cirq/transformers/stratify_test.py

@@ -237,29 +237,29 @@ def test_stratify_respects_no_compile_operations():
     cirq.testing.assert_has_diagram(
         input_circuit,
         '''
-0: ───X['nocompile']───────X───────iSwap───
-                                   │
-1: ───iSwap['nocompile']───────────iSwap───
+0: ───X[nocompile]───────X───────iSwap───
+                                 │
+1: ───iSwap[nocompile]───────────iSwap───
       │
-2: ───iSwap────────────────────────────────
+2: ───iSwap──────────────────────────────
 
-3: ────────────────────────iSwap───X───────
-                           │
-4: ───Z────────────────────iSwap───────────
+3: ──────────────────────iSwap───X───────
+                         │
+4: ───Z──────────────────iSwap───────────
 ''',
     )
     cirq.testing.assert_has_diagram(
         expected,
         '''
-0: ───────────────X['nocompile']───────X───iSwap───
-                                           │
-1: ───────────────iSwap['nocompile']───────iSwap───
+0: ───────────────X[nocompile]───────X───iSwap───
+                                         │
+1: ───────────────iSwap[nocompile]───────iSwap───
                   │
-2: ───────────────iSwap────────────────────────────
+2: ───────────────iSwap──────────────────────────
 
-3: ───────iSwap────────────────────────X───────────
+3: ───────iSwap──────────────────────X───────────
           │
-4: ───Z───iSwap────────────────────────────────────
+4: ───Z───iSwap──────────────────────────────────
 ''',
     )
     cirq.testing.assert_same_circuits(

cirq/transformers/target_gatesets/compilation_target_gateset.py

@@ -14,7 +14,7 @@
 
 """Base class for creating custom target gatesets which can be used for compilation."""
 
-from typing import Optional, List, Hashable, TYPE_CHECKING
+from typing import Optional, List, Hashable, TYPE_CHECKING, Union, Type
 import abc
 
 from cirq import circuits, ops, protocols, transformers
@@ -80,6 +80,27 @@ class CompilationTargetGateset(ops.Gateset, metaclass=abc.ABCMeta):
     which can transform any given circuit to contain gates accepted by this gateset.
     """
 
+    def __init__(
+        self,
+        *gates: Union[Type['cirq.Gate'], 'cirq.Gate', 'cirq.GateFamily'],
+        name: Optional[str] = None,
+        unroll_circuit_op: bool = True,
+        preserve_moment_structure: bool = True,
+    ):
+        """Initializes CompilationTargetGateset.
+
+        Args:
+            *gates: A list of `cirq.Gate` subclasses / `cirq.Gate` instances /
+                `cirq.GateFamily` instances to initialize the Gateset.
+            name: (Optional) Name for the Gateset. Useful for description.
+            unroll_circuit_op: If True, `cirq.CircuitOperation` is recursively
+                validated by validating the underlying `cirq.Circuit`.
+            preserve_moment_structure: Whether to preserve the moment structure of the
+                circuit during compilation or not.
+        """
+        super().__init__(*gates, name=name, unroll_circuit_op=unroll_circuit_op)
+        self._preserve_moment_structure = preserve_moment_structure
+
     @property
     @abc.abstractmethod
     def num_qubits(self) -> int:
@@ -140,11 +161,14 @@ class CompilationTargetGateset(ops.Gateset, metaclass=abc.ABCMeta):
     @property
     def postprocess_transformers(self) -> List['cirq.TRANSFORMER']:
         """List of transformers which should be run after decomposing individual operations."""
-        return [
+        processors: List['cirq.TRANSFORMER'] = [
             merge_single_qubit_gates.merge_single_qubit_moments_to_phxz,
             transformers.drop_negligible_operations,
             transformers.drop_empty_moments,
         ]
+        if not self._preserve_moment_structure:
+            processors.append(transformers.stratified_circuit)
+        return processors
 
 
 class TwoQubitCompilationTargetGateset(CompilationTargetGateset):

cirq/transformers/target_gatesets/compilation_target_gateset_test.py

@@ -138,9 +138,9 @@ def test_two_qubit_compilation_merge_and_replace_to_target_gateset():
     cirq.testing.assert_has_diagram(
         c_orig,
         '''
-0: ───X───@['no_compile']───Z───X───@───Z───X───
-          │                         │
-1: ───Z───@─────────────────Z───────@───Z───────
+0: ───X───@[no_compile]───Z───X───@───Z───X───
+          │                       │
+1: ───Z───@───────────────Z───────@───Z───────
 ''',
     )
     c_new = cirq.optimize_for_target_gateset(
@@ -151,9 +151,9 @@ def test_two_qubit_compilation_merge_and_replace_to_target_gateset():
     cirq.testing.assert_has_diagram(
         c_new,
         '''
-0: ───X───@['no_compile']───X───@───Y───@───Z───
-          │                     │       │
-1: ───Z───@─────────────────X───X───Y───X───Z───
+0: ───X───@[no_compile]───X───@───Y───@───Z───
+          │                   │       │
+1: ───Z───@───────────────X───X───Y───X───Z───
 ''',
     )
 
@@ -178,11 +178,11 @@ def test_two_qubit_compilation_merge_and_replace_inefficient_component():
     cirq.testing.assert_has_diagram(
         c_orig,
         '''
-0: ───X───@───X───@['no_compile']───Z───X───@───@───Z───X───@───M───────
-          │       │                         │   │           │   ║
-1: ───────X───────@─────────────────Z───────X───X───Z───────X───╫───X───
-                                                                ║   ║
-m: ═════════════════════════════════════════════════════════════@═══^═══
+0: ───X───@───X───@[no_compile]───Z───X───@───@───Z───X───@───M───────
+          │       │                       │   │           │   ║
+1: ───────X───────@───────────────Z───────X───X───Z───────X───╫───X───
+                                                              ║   ║
+m: ═══════════════════════════════════════════════════════════@═══^═══
 ''',
     )
     c_new = cirq.optimize_for_target_gateset(
@@ -193,11 +193,11 @@ m: ═════════════════════════
     cirq.testing.assert_has_diagram(
         c_new,
         '''
-0: ───X───@───X───@['no_compile']───X───@───Y───@───Z───M───────
-          │       │                     │       │       ║
-1: ───────X───────@─────────────────X───X───Y───X───Z───╫───X───
-                                                        ║   ║
-m: ═════════════════════════════════════════════════════@═══^═══
+0: ───X───@───X───@[no_compile]───X───@───Y───@───Z───M───────
+          │       │                   │       │       ║
+1: ───────X───────@───────────────X───X───Y───X───Z───╫───X───
+                                                      ║   ║
+m: ═══════════════════════════════════════════════════@═══^═══
 ''',
     )
 

cirq/transformers/target_gatesets/cz_gateset.py

@@ -48,6 +48,7 @@ class CZTargetGateset(compilation_target_gateset.TwoQubitCompilationTargetGatese
         atol: float = 1e-8,
         allow_partial_czs: bool = False,
         additional_gates: Sequence[Union[Type['cirq.Gate'], 'cirq.Gate', 'cirq.GateFamily']] = (),
+        preserve_moment_structure: bool = True,
     ) -> None:
         """Initializes CZTargetGateset
 
@@ -57,6 +58,8 @@ class CZTargetGateset(compilation_target_gateset.TwoQubitCompilationTargetGatese
              `cirq.CZ`, are part of this gateset.
             additional_gates: Sequence of additional gates / gate families which should also
               be "accepted" by this gateset. This is empty by default.
+            preserve_moment_structure: Whether to preserve the moment structure of the
+                circuit during compilation or not.
         """
         super().__init__(
             ops.CZPowGate if allow_partial_czs else ops.CZ,
@@ -65,6 +68,7 @@ class CZTargetGateset(compilation_target_gateset.TwoQubitCompilationTargetGatese
             ops.GlobalPhaseGate,
             *additional_gates,
             name='CZPowTargetGateset' if allow_partial_czs else 'CZTargetGateset',
+            preserve_moment_structure=preserve_moment_structure,
         )
         self.additional_gates = tuple(
             g if isinstance(g, ops.GateFamily) else ops.GateFamily(gate=g) for g in additional_gates

cirq/transformers/transformer_api.py

@@ -262,8 +262,7 @@ class TRANSFORMER(Protocol):
 
     def __call__(
         self, circuit: 'cirq.AbstractCircuit', *, context: Optional[TransformerContext] = None
-    ) -> 'cirq.AbstractCircuit':
-        ...
+    ) -> 'cirq.AbstractCircuit': ...
 
 
 _TRANSFORMER_T = TypeVar('_TRANSFORMER_T', bound=TRANSFORMER)