optimize_for_target_gateset OOMs on large CircuitOps

[95-martin-orion]

Description of the issue

When a circuit containing a large (in this case, 17 qubits & depth 20) CircuitOp is passed to optimize_for_target_gateset, it generates an OOM, possibly crashing the device. This circuit is well within the realm of "easy to simulate", so we shouldn't be OOMing or should at least warn users that this has a lower threshold than simulation.

How to reproduce the issue

qnum = 17
qubits = cirq.LineQubit.range(qnum)
subcircuit = cirq.FrozenCircuit(
    *[
      cirq.CX(qubits[i], qubits[i+1])
      for i in range(0, qnum-1, 2)
    ],
    *[
      cirq.CZ(qubits[i], qubits[i+1])
      for i in range(1, qnum-1, 2)
    ],
)
circuit = cirq.Circuit(
    cirq.CircuitOperation(subcircuit).repeat(10),
    cirq.measure(*qubits, key='out')
)
print(circuit)
# OOM: cannot allocate 256GB
supported_circuit = cirq.optimize_for_target_gateset(
    circuit,
    context=cirq.TransformerContext(deep=True),
    gateset=cirq.CZTargetGateset(),
)

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_api.py", line 380, in func_with_logging
    return _transform_and_log(
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_api.py", line 439, in _transform_and_log
    transformed_circuit = _run_transformer_on_circuit(
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_api.py", line 425, in _run_transformer_on_circuit
    return func(mutable_circuit if mutable_circuit else circuit, **kwargs)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/optimize_for_target_gateset.py", line 143, in optimize_for_target_gateset
    circuit = transformer(circuit, context=context)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_api.py", line 380, in func_with_logging
    return _transform_and_log(
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_api.py", line 439, in _transform_and_log
    transformed_circuit = _run_transformer_on_circuit(
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_api.py", line 425, in _run_transformer_on_circuit
    return func(mutable_circuit if mutable_circuit else circuit, **kwargs)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/drop_negligible_operations.py", line 54, in drop_negligible_operations
    return transformer_primitives.map_operations(
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_primitives.py", line 143, in map_operations
    return map_moments(
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_primitives.py", line 87, in map_moments
    return _create_target_circuit_type(
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_primitives.py", line 43, in _create_target_circuit_type
    circuits.Circuit(ops)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/circuits/circuit.py", line 1707, in __init__
    self.append(contents, strategy=strategy)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/circuits/circuit.py", line 2254, in append
    self.insert(len(self._moments), moment_or_operation_tree, strategy)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/circuits/circuit.py", line 1969, in insert
    moments_and_operations = list(
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/ops/op_tree.py", line 134, in flatten_to_ops_or_moments
    yield from flatten_to_ops_or_moments(subtree)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/ops/op_tree.py", line 133, in flatten_to_ops_or_moments
    for subtree in root:
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/ops/op_tree.py", line 169, in <genexpr>
    return iter_transformation(
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_primitives.py", line 88, in <genexpr>
    (map_func(mutable_circuit[i], i) for i in range(len(mutable_circuit))), circuit
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_primitives.py", line 145, in <lambda>
    lambda m, i: [circuits.Moment(apply_map(op, i) for op in m.operations)],
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/circuits/moment.py", line 93, in __init__
    self._operations = tuple(op_tree.flatten_to_ops(contents))
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/ops/op_tree.py", line 113, in flatten_to_ops
    yield from flatten_to_ops(subtree)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/ops/op_tree.py", line 112, in flatten_to_ops
    for subtree in root:
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_primitives.py", line 145, in <genexpr>
    lambda m, i: [circuits.Moment(apply_map(op, i) for op in m.operations)],
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/transformer_primitives.py", line 130, in apply_map
    c = circuits.FrozenCircuit(map_func(op, idx))
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/transformers/drop_negligible_operations.py", line 51, in map_func
    op if protocols.is_measurement(op) or protocols.trace_distance_bound(op) > atol else []
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/protocols/trace_distance_bound.py", line 73, in trace_distance_bound
    result = strat(val)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/protocols/trace_distance_bound.py", line 98, in _strat_distance_from_unitary
    u = unitary_protocol.unitary(val, default=None)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/protocols/unitary_protocol.py", line 119, in unitary
    result = strat(val)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/protocols/unitary_protocol.py", line 183, in _strat_unitary_from_decompose
    state = qis.eye_tensor(val_qid_shape, dtype=np.complex128)
  File "~/git/cirq-realnoise-docs/cirq-core/cirq/qis/states.py", line 1076, in eye_tensor
    identity = np.eye(np.prod(half_shape, dtype=np.int64).item(), dtype=dtype)
  File "~/.virtualenvs/cirq/lib/python3.9/site-packages/numpy/lib/twodim_base.py", line 209, in eye
    m = zeros((N, M), dtype=dtype, order=order)
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 256. GiB for an array with shape (131072, 131072) and data type complex128

Cirq version
0.15.0.dev

[tanujkhattar]

This happens because cirq.drop_negligible_operations(circuit) transformer calls cirq.trace_distance_bound(circuit_op) for the circuit operation, which in-turn tries to compute the unitary matrix of the circuit operation which results in OOM error.

qnum = 17
qubits = cirq.LineQubit.range(qnum)
subcircuit = cirq.FrozenCircuit(
    *[
      cirq.CX(qubits[i], qubits[i+1])
      for i in range(0, qnum-1, 2)
    ],
    *[
      cirq.CZ(qubits[i], qubits[i+1])
      for i in range(1, qnum-1, 2)
    ],
)
# Fails with OOM
cirq.unitary(cirq.CircuitOperation(subcircuit).repeat(10))

--------------------------------------------------------------------------- MemoryError Traceback (most recent call last) in 11 ], 12 ) ---> 13 cirq.unitary(cirq.CircuitOperation(subcircuit).repeat(10))

~/quantum/Cirq/cirq-core/cirq/protocols/unitary_protocol.py in unitary(val, default)
117 ]
118 for strat in strats:
--> 119 result = strat(val)
120 if result is None:
121 break

~/quantum/Cirq/cirq-core/cirq/protocols/unitary_protocol.py in _strat_unitary_from_decompose(val)
181
182 # Apply sub-operations' unitary effects to an identity matrix.
--> 183 state = qis.eye_tensor(val_qid_shape, dtype=np.complex128)
184 buffer = np.empty_like(state)
185 result = apply_unitaries(

~/quantum/Cirq/cirq-core/cirq/qis/states.py in eye_tensor(half_shape, dtype)
1074 The created numpy array with shape half_shape + half_shape.
1075 """
-> 1076 identity = np.eye(np.prod(half_shape, dtype=np.int64).item(), dtype=dtype)
1077 identity.shape = half_shape * 2
1078 return identity

~/anaconda3/envs/cirq/lib/python3.9/site-packages/numpy/lib/twodim_base.py in eye(N, M, k, dtype, order, like)
212 if M is None:
213 M = N
--> 214 m = zeros((N, M), dtype=dtype, order=order)
215 if k >= M:
216 return m

MemoryError: Unable to allocate 256. GiB for an array with shape (131072, 131072) and data type complex128

To fix this, we can hardcode a number N s.t. we ignore the operation if it acts on > N qubits. The hardcode can exist in:

• cirq.drop_negligible_operations transformer -- the cirq.trace_distance_bound and cirq.unitary will still fail if called on this operation.
• cirq.trace_distance_bound method -- the cirq.unitary protocol will still fail
• cirq.unitary -- I think this will be more invasive since we'd ideally expect unitary to work if running on a beefier machine where this memory allocation is possible.

My preference would be to simply change the cirq.drop_negligible_operations transformer and add a hardcoded check to ignore all operation acting on > 10 qubits, for example.

I'd be curious to hear what other maintainers think. We can bring this up in the next Cirq sync.

[tanujkhattar]

From cirq sync:

Add a special case to cirq.drop_negligible_operations and leave the other two methods as it is.