Optimize swap gates (#4169)

* split

* Allow config param split_entangled_states

* default split to off

* ensure consistent_act_on circuits have a qubit.

* lint

* lint

* mps

* lint

* lint

* run sparse by default

* fix tests

* fix tests

* fix tests

* mps

* tableau

* test simulator

* test simulator

* Update simulator_base.py

* Drop mps/join

* Fix clifford extract

* lint

* remove split/join from ch-form

* Add default arg for zero qubit circuits

* Have last repetition reuse original state repr

* Remove cast

* Split all pure initial states by default

* Detangle on reset channels

* docstrings

* docstrings

* docstrings

* docstrings

* fix merge

* lint

* Add unit test for integer states

* format

* Add tests for splitting and joining

* remove unnecessary qubits param

* Clean up default args

* Fix failing test

* Add ActOnArgsContainer

* Add ActOnArgsContainer

* Clean up tests

* Clean up tests

* Clean up tests

* format

* Fix tests and coverage

* Add OperationTarget interface

* Fix unit tests

* mypy, lint, mocks, coverage

* coverage

* add log to container

* fix logs

* dead code

* EmptyActOnArgs

* simplify dummyargs

* Optimize swap gate

* lint

* cover

* lint

* inplace option

* Add [] to actonargs

* rename _create_act_on_arg

* coverage

* coverage

* Default sparse sim to split=false

* format

* Default sparse sim to split=false

* Default density matrix sim to split=false

* lint

* lint

* lint

* lint

* address review comments

* lint

* Defaults back to split=false

* add error if setting state when split is enabled

* Unit tests

* coverage

* coverage

* coverage

* docs

* fix merge

* tests

* better tests

* lint

* Update docstrings.

Co-authored-by: Cirq Bot <craiggidney+github+cirqbot@google.com>
Co-authored-by: Orion Martin <40585662+95-martin-orion@users.noreply.github.com>

cirq-core/cirq/sim/act_on_args.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Objects and methods for acting efficiently on a state tensor."""
 import abc
+import copy
 from typing import (
     Any,
     Iterable,
@@ -139,6 +140,64 @@ def log_of_measurement_results(self) -> Dict[str, Any]:
     def qubits(self) -> Tuple['cirq.Qid', ...]:
         return self._qubits
 
+    def swap(self, q1: 'cirq.Qid', q2: 'cirq.Qid', *, inplace=False):
+        """Swaps two qubits.
+
+        This only affects the index, and does not modify the underlying
+        state.
+
+        Args:
+            q1: The first qubit to swap.
+            q2: The second qubit to swap.
+            inplace: True to swap the qubits in the current object, False to
+                create a copy with the qubits swapped.
+
+        Returns:
+            The original object with the qubits swapped if inplace is
+            requested, or a copy of the original object with the qubits swapped
+            otherwise.
+
+        Raises:
+            ValueError if the qubits are of different dimensionality."""
+        if q1.dimension != q2.dimension:
+            raise ValueError(f'Cannot swap different dimensions: q1={q1}, q2={q2}')
+
+        args = self if inplace else copy.copy(self)
+        i1 = self.qubits.index(q1)
+        i2 = self.qubits.index(q2)
+        qubits = list(args.qubits)
+        qubits[i1], qubits[i2] = qubits[i2], qubits[i1]
+        args._qubits = tuple(qubits)
+        args.qubit_map = {q: i for i, q in enumerate(qubits)}
+        return args
+
+    def rename(self, q1: 'cirq.Qid', q2: 'cirq.Qid', *, inplace=False):
+        """Renames `q1` to `q2`.
+
+        Args:
+            q1: The qubit to rename.
+            q2: The new name.
+            inplace: True to rename the qubit in the current object, False to
+                create a copy with the qubit renamed.
+
+        Returns:
+            The original object with the qubits renamed if inplace is
+            requested, or a copy of the original object with the qubits renamed
+            otherwise.
+
+        Raises:
+            ValueError if the qubits are of different dimensionality."""
+        if q1.dimension != q2.dimension:
+            raise ValueError(f'Cannot rename to different dimensions: q1={q1}, q2={q2}')
+
+        args = self if inplace else copy.copy(self)
+        i1 = self.qubits.index(q1)
+        qubits = list(args.qubits)
+        qubits[i1] = q2
+        args._qubits = tuple(qubits)
+        args.qubit_map = {q: i for i, q in enumerate(qubits)}
+        return args
+
     def __getitem__(self: TSelf, item: Optional['cirq.Qid']) -> TSelf:
         if item not in self.qubit_map:
             raise IndexError(f'{item} not in {self.qubits}')

cirq-core/cirq/sim/act_on_args_container.py

@@ -81,6 +81,18 @@ def apply_operation(
         self,
         op: 'cirq.Operation',
     ):
+        gate = op.gate
+        if isinstance(gate, ops.SwapPowGate) and gate.exponent % 2 == 1 and gate.global_shift == 0:
+            q0, q1 = op.qubits
+            args0 = self.args[q0]
+            args1 = self.args[q1]
+            if args0 is args1:
+                args0.swap(q0, q1, inplace=True)
+            else:
+                self.args[q0] = args1.rename(q1, q0, inplace=True)
+                self.args[q1] = args0.rename(q0, q1, inplace=True)
+            return
+
         # Go through the op's qubits and join any disparate ActOnArgs states
         # into a new combined state.
         op_args_opt: Optional[TActOnArgs] = None

cirq-core/cirq/sim/act_on_args_container_test.py

@@ -172,3 +172,35 @@ def test_merge_succeeds():
     args = create_container(qs2, False)
     merged = args.create_merged_state()
     assert merged.qubits == (q0, q1)
+
+
+def test_swap_does_not_merge():
+    args = create_container(qs2)
+    old_q0 = args[q0]
+    old_q1 = args[q1]
+    args.apply_operation(cirq.SWAP(q0, q1))
+    assert len(set(args.values())) == 3
+    assert args[q0] is not old_q0
+    assert args[q1] is old_q0
+    assert args[q1] is not old_q1
+    assert args[q0] is old_q1
+    assert args[q0].qubits == (q0,)
+    assert args[q1].qubits == (q1,)
+
+
+def test_half_swap_does_merge():
+    args = create_container(qs2)
+    args.apply_operation(cirq.SWAP(q0, q1) ** 0.5)
+    assert len(set(args.values())) == 2
+    assert args[q0] is args[q1]
+
+
+def test_swap_after_entangle_reorders():
+    args = create_container(qs2)
+    args.apply_operation(cirq.CX(q0, q1))
+    assert len(set(args.values())) == 2
+    assert args[q0].qubits == (q0, q1)
+    args.apply_operation(cirq.SWAP(q0, q1))
+    assert len(set(args.values())) == 2
+    assert args[q0] is args[q1]
+    assert args[q0].qubits == (q1, q0)

cirq-core/cirq/sim/act_on_args_test.py

@@ -59,3 +59,19 @@ def test_mapping():
     assert args is r1
     with pytest.raises(IndexError):
         _ = args[cirq.LineQubit(2)]
+
+
+def test_swap_bad_dimensions():
+    q0 = cirq.LineQubit(0)
+    q1 = cirq.LineQid(1, 3)
+    args = DummyArgs()
+    with pytest.raises(ValueError, match='Cannot swap different dimensions'):
+        args.swap(q0, q1)
+
+
+def test_rename_bad_dimensions():
+    q0 = cirq.LineQubit(0)
+    q1 = cirq.LineQid(1, 3)
+    args = DummyArgs()
+    with pytest.raises(ValueError, match='Cannot rename to different dimensions'):
+        args.rename(q0, q1)