Perf optimization for building circuits using only appends

cirq-core/cirq/circuits/circuit.py

@@ -1673,6 +1673,39 @@ def _concat_ragged_helper(
     return min(c1_offset, c2_offset), max(n1, n2, n1 + n2 - shift)
 
 
+class _OpPlacer:
+    """Maintains qubit and cbit indices for quick op placement.
+
+    Here, we instead keep track of the greatest moment that contains each
+    qubit, measurement key, and control key, and append the operation to
+    the moment after the maximum of these. This avoids having to check each
+    moment.
+    """
+
+    def __init__(self) -> None:
+        # These are dicts from the qubit/key to the greatest moment index that has it.
+        self._qubit_indices: Dict['cirq.Qid', int] = {}
+        self._mkey_indices: Dict['cirq.MeasurementKey', int] = {}
+        self._ckey_indices: Dict['cirq.MeasurementKey', int] = {}
+
+        # For keeping track of length of the circuit thus far.
+        self._length = 0
+
+    def get_earliest_accommodating_moment_index(
+        self, moment_or_operation: Union['cirq.Moment', 'cirq.Operation']
+    ):
+        # Identify the index of the moment to place this into.
+        index = get_earliest_accommodating_moment_index(
+            moment_or_operation,
+            self._qubit_indices,
+            self._mkey_indices,
+            self._ckey_indices,
+            self._length,
+        )
+        self._length = max(self._length, index + 1)
+        return index
+
+
 class Circuit(AbstractCircuit):
     """A mutable list of groups of operations to apply to some qubits.
 
@@ -1769,6 +1802,7 @@ def __init__(
                 together. This option does not affect later insertions into the
                 circuit.
         """
+        self._placer: Optional[_OpPlacer] = _OpPlacer()
         self._moments: List['cirq.Moment'] = []
 
         # Implementation note: the following cached properties are set lazily and then
@@ -1779,9 +1813,11 @@ def __init__(
         self._is_measurement: Optional[bool] = None
         self._is_parameterized: Optional[bool] = None
         self._parameter_names: Optional[AbstractSet[str]] = None
-
+        if not contents:
+            return
         flattened_contents = tuple(ops.flatten_to_ops_or_moments(contents))
         if all(isinstance(c, Moment) for c in flattened_contents):
+            self._placer = None
             self._moments[:] = cast(Iterable[Moment], flattened_contents)
             return
         with _compat.block_overlapping_deprecation('.*'):
@@ -1790,18 +1826,21 @@ def __init__(
             else:
                 self.append(flattened_contents, strategy=strategy)
 
-    def _mutated(self) -> None:
+    def _mutated(self, preserve_placer=False) -> None:
         """Clear cached properties in response to this circuit being mutated."""
         self._all_qubits = None
         self._frozen = None
         self._is_measurement = None
         self._is_parameterized = None
         self._parameter_names = None
+        if not preserve_placer:
+            self._placer = None
 
     @classmethod
     def _from_moments(cls, moments: Iterable['cirq.Moment']) -> 'Circuit':
         new_circuit = Circuit()
         new_circuit._moments[:] = moments
+        new_circuit._placer = None
         return new_circuit
 
     def _load_contents_with_earliest_strategy(self, contents: 'cirq.OP_TREE'):
@@ -1823,35 +1862,24 @@ def _load_contents_with_earliest_strategy(self, contents: 'cirq.OP_TREE'):
                 Non-moment entries will be inserted according to the EARLIEST
                 insertion strategy.
         """
-        # These are dicts from the qubit/key to the greatest moment index that has it.
-        qubit_indices: Dict['cirq.Qid', int] = {}
-        mkey_indices: Dict['cirq.MeasurementKey', int] = {}
-        ckey_indices: Dict['cirq.MeasurementKey', int] = {}
-
-        # We also maintain the dict from moment index to moments/ops that go into it, for use when
+        # We maintain the dict from moment index to moments/ops that go into it, for use when
         # building the actual moments at the end.
         op_lists_by_index: Dict[int, List['cirq.Operation']] = defaultdict(list)
         moments_by_index: Dict[int, 'cirq.Moment'] = {}
-
-        # For keeping track of length of the circuit thus far.
-        length = 0
+        placer = cast(_OpPlacer, self._placer)
 
         # "mop" means current moment-or-operation
         for mop in ops.flatten_to_ops_or_moments(contents):
             # Identify the index of the moment to place this `mop` into.
-            placement_index = get_earliest_accommodating_moment_index(
-                mop, qubit_indices, mkey_indices, ckey_indices, length
-            )
-            length = max(length, placement_index + 1)  # update the length of the circuit thus far
-
+            placement_index = placer.get_earliest_accommodating_moment_index(mop)
             if isinstance(mop, Moment):
                 moments_by_index[placement_index] = mop
             else:
                 op_lists_by_index[placement_index].append(mop)
 
         # Finally, once everything is placed, we can construct and append the actual moments for
         # each index.
-        for i in range(length):
+        for i in range(placer._length):
             if i in moments_by_index:
                 self._moments.append(moments_by_index[i].with_operations(op_lists_by_index[i]))
             else:
@@ -1899,6 +1927,7 @@ def copy(self) -> 'Circuit':
         """Return a copy of this circuit."""
         copied_circuit = Circuit()
         copied_circuit._moments = self._moments[:]
+        copied_circuit._placer = None
         return copied_circuit
 
     # pylint: disable=function-redefined
@@ -2154,20 +2183,38 @@ def insert(
         """
         # limit index to 0..len(self._moments), also deal with indices smaller 0
         k = max(min(index if index >= 0 else len(self._moments) + index, len(self._moments)), 0)
-        for moment_or_op in list(ops.flatten_to_ops_or_moments(moment_or_operation_tree)):
-            if isinstance(moment_or_op, Moment):
-                self._moments.insert(k, moment_or_op)
-                k += 1
-            else:
-                op = moment_or_op
-                p = self._pick_or_create_inserted_op_moment_index(k, op, strategy)
-                while p >= len(self._moments):
-                    self._moments.append(Moment())
-                self._moments[p] = self._moments[p].with_operation(op)
+        moments_or_ops = list(ops.flatten_to_ops_or_moments(moment_or_operation_tree))
+        if self._placer and strategy == InsertStrategy.EARLIEST and index == len(self._moments):
+            # Use `placer` to get placement indices quickly.
+            for moment_or_op in moments_or_ops:
+                p = self._placer.get_earliest_accommodating_moment_index(moment_or_op)
+                if isinstance(moment_or_op, Moment):
+                    self._moments.append(moment_or_op)
+                else:
+                    if p >= len(self._moments):
+                        self._moments.append(Moment(moment_or_op))
+                    else:
+                        self._moments[p] = self._moments[p].with_operation(moment_or_op)
                 k = max(k, p + 1)
-                if strategy is InsertStrategy.NEW_THEN_INLINE:
-                    strategy = InsertStrategy.INLINE
-        self._mutated()
+            self._mutated(preserve_placer=True)
+        else:
+            # Default algorithm. Same behavior as above, but has to search for placement indices.
+            # First invalidate the placer due to unsupported insertion.
+            self._placer = None
+            for moment_or_op in moments_or_ops:
+                if isinstance(moment_or_op, Moment):
+                    self._moments.insert(k, moment_or_op)
+                    k += 1
+                else:
+                    op = moment_or_op
+                    p = self._pick_or_create_inserted_op_moment_index(k, op, strategy)
+                    while p >= len(self._moments):
+                        self._moments.append(Moment())
+                    self._moments[p] = self._moments[p].with_operation(op)
+                    k = max(k, p + 1)
+                    if strategy is InsertStrategy.NEW_THEN_INLINE:
+                        strategy = InsertStrategy.INLINE
+            self._mutated()
         return k
 
     def insert_into_range(self, operations: 'cirq.OP_TREE', start: int, end: int) -> int:

cirq-core/cirq/circuits/circuit_test.py

@@ -4845,18 +4845,35 @@ def _circuit_diagram_info_(self, args) -> str:
 def test_create_speed():
     # Added in https://github.com/quantumlib/Cirq/pull/5332
     # Previously this took ~30s to run. Now it should take ~150ms. However the coverage test can
-    # run this slowly, so allowing 2 sec to account for things like that. Feel free to increase the
+    # run this slowly, so allowing 4 sec to account for things like that. Feel free to increase the
     # buffer time or delete the test entirely if it ends up causing flakes.
-    #
-    # Updated in https://github.com/quantumlib/Cirq/pull/5756
-    # After several tiny overtime failures of the GitHub CI Pytest MacOS (3.7)
-    # the timeout was increased to 4 sec.  A more thorough investigation or test
-    # removal should be considered if this continues to time out.
     qs = 100
     moments = 500
     xs = [cirq.X(cirq.LineQubit(i)) for i in range(qs)]
-    opa = [xs[i] for i in range(qs) for _ in range(moments)]
+    ops = [xs[i] for i in range(qs) for _ in range(moments)]
     t = time.perf_counter()
-    c = cirq.Circuit(opa)
+    c = cirq.Circuit(ops)
     assert len(c) == moments
     assert time.perf_counter() - t < 4
+
+
+def test_append_speed():
+    # Previously this took ~17s to run. Now it should take ~150ms. However the coverage test can
+    # run this slowly, so allowing 4 sec to account for things like that. Feel free to increase the
+    # buffer time or delete the test entirely if it ends up causing flakes.
+    #
+    # The `append` improvement mainly helps for deep circuits. It is less useful for wide circuits
+    # because the Moment (immutable) needs verified and reconstructed each time an op is added.
+    qs = 2
+    moments = 10000
+    xs = [cirq.X(cirq.LineQubit(i)) for i in range(qs)]
+    c = cirq.Circuit()
+    t = time.perf_counter()
+    # Iterating with the moments in the inner loop highlights the improvement: when filling in the
+    # second qubit, we no longer have to search backwards from moment 10000 for a placement index.
+    for q in range(qs):
+        for _ in range(moments):
+            c.append(xs[q])
+    duration = time.perf_counter() - t
+    assert len(c) == moments
+    assert duration < 4

cirq-core/cirq/contrib/acquaintance/executor.py

@@ -75,6 +75,7 @@ def __call__(self, *args, **kwargs):
         strategy = StrategyExecutorTransformer(self)
         final_circuit = strategy(input_circuit, **kwargs)
         input_circuit._moments = final_circuit._moments
+        input_circuit._placer = final_circuit._placer
         return strategy.mapping