IdentityGate optimizations for act_on and independent states (S) (#4340)

This exponentially improves perf on density matrix simulator for identity gates with several qubits. It also prevents merging independent states when an identity gate has multiple qubits.

Author: daxfohl

cirq-core/cirq/ops/identity.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """IdentityGate."""
 
-from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
+from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Sequence
 
 import numpy as np
 import sympy
@@ -59,6 +59,9 @@ def __init__(
         if len(self._qid_shape) != num_qubits:
             raise ValueError('len(qid_shape) != num_qubits')
 
+    def _act_on_(self, args: 'cirq.ActOnArgs', qubits: Sequence['cirq.Qid']):
+        return True
+
     def _qid_shape_(self) -> Tuple[int, ...]:
         return self._qid_shape
 

cirq-core/cirq/ops/identity_test.py

@@ -13,10 +13,12 @@
 # limitations under the License.
 import itertools
 from typing import Any
+from unittest import mock
 
 import numpy as np
 import pytest
 import sympy
+
 import cirq
 
 
@@ -202,3 +204,9 @@ def with_qubits(self, *new_qubits):
 
     assert i * 2 == cirq.PauliString(coefficient=2)
     assert 1j * i == cirq.PauliString(coefficient=1j)
+
+
+def test_identity_short_circuits_act_on():
+    args = mock.Mock(cirq.ActOnArgs)
+    args._act_on_fallback_.side_effect = mock.Mock(side_effect=Exception('No!'))
+    cirq.act_on(cirq.IdentityGate(1)(cirq.LineQubit(0)), args)

cirq-core/cirq/sim/act_on_args_container.py

@@ -82,6 +82,9 @@ def apply_operation(
         op: 'cirq.Operation',
     ):
         gate = op.gate
+        if isinstance(gate, ops.IdentityGate):
+            return
+
         if isinstance(gate, ops.SwapPowGate) and gate.exponent % 2 == 1 and gate.global_shift == 0:
             q0, q1 = op.qubits
             args0 = self.args[q0]

cirq-core/cirq/sim/act_on_args_container_test.py

@@ -98,6 +98,15 @@ def test_entanglement_causes_join():
     assert args[None] is not args[q0]
 
 
+def test_identity_does_not_join():
+    args = create_container(qs2)
+    assert len(set(args.values())) == 3
+    args.apply_operation(cirq.IdentityGate(2)(q0, q1))
+    assert len(set(args.values())) == 3
+    assert args[q0] is not args[q1]
+    assert args[q0] is not args[None]
+
+
 def test_measurement_causes_split():
     args = create_container(qs2)
     args.apply_operation(cirq.CNOT(q0, q1))