Provide a way to avoid copying state vector from Cirq

[mhucka]

What is your request or suggestion?

An internal user reported getting the following error with a 40 qubit simulation:

Traceback (most recent call last):
  File "/tmp/test/exact_sim.py", line 23, in <module>
    psi = qsim.simulate(circuit, qubit_order=qubits).state_vector()
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venv/lib64/python3.12/site-packages/cirq/sim/simulator.py", line 496, in simulate
    return self.simulate_sweep(
           ^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venv/lib64/python3.12/site-packages/cirq/sim/simulator.py", line 511, in simulate_sweep
    return list(self.simulate_sweep_iter(program, params, qubit_order, initial_state))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venv/lib64/python3.12/site-packages/qsimcirq/qsim_simulator.py", line 560, in simulate_sweep_iter
    final_state = cirq.StateVectorSimulationState(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venv/lib64/python3.12/site-packages/cirq/sim/state_vector_simulation_state.py", line 351, in __init__
    state = _BufferedStateVector.create(
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venv/lib64/python3.12/site-packages/cirq/sim/state_vector_simulation_state.py", line 84, in create
    state_vector = state_vector.copy()
                   ^^^^^^^^^^^^^^^^^^^
numpy._core._exceptions._ArrayMemoryError: Unable to allocate 8.00 TiB for an array with shape (2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2) and data type complex64

They were able to work around the problem by calling QSimSimulator()._simulate_impl() directly.

What solution or approach do you envision?

We should investigate whether the copy can be handled in a more user-friendly way.

How urgent is this for you?

P2 – needed within two quarters

[jkalsi1]

Hi @mhucka . In cirq-core/cirq/sim/state_vector_simulation_state.py.__BufferedStateVector.create(), the following condition evaluates truthy, which forces the copy of the array to be made, causing the above error:

if np.may_share_memory(state_vector, initial_state):
        state_vector = state_vector.copy()

It appears to me that this is to protect the integrity of the original initial_state argument passed to _BufferedStateVector.create() against unwanted mutations.

Could we create a flag in cirq-core/cirq/sim/state_vector_simulation_state.py.StateVectorSimulationState that allows the caller to choose whether they care about the integrity of the initial_state array or not. If they don't care, then skip the copy condition.

Something like:

In StateVectorSimulationState:

    def __init__(
        self,
        *,
        available_buffer: np.ndarray | None = None,
        prng: np.random.RandomState | None = None,
        qubits: Sequence[cirq.Qid] | None = None,
        initial_state: np.ndarray | cirq.STATE_VECTOR_LIKE = 0,
        dtype: type[np.complexfloating] = np.complex64,
        classical_data: cirq.ClassicalDataStore | None = None,
        # should_preserve_initial_state: bool = True,
    ):

Then in _BufferedStateVector.create():

    @classmethod
    def create(
        cls,
        *,
        initial_state: np.ndarray | cirq.STATE_VECTOR_LIKE = 0,
        qid_shape: tuple[int, ...] | None = None,
        dtype: type[np.complexfloating] | None = None,
        buffer: np.ndarray | None = None,
        # should_preserve_initial_state: bool = True,
    ):

  # some code here 
 # ....
if np.may_share_memory(state_vector, initial_state) and should_preserve_initial_state: # listen to the flag
        state_vector = state_vector.copy()
# ....

These code changes would appear in Cirq, not qsim, but I believe may resolve the issue that the internal user was experiencing.