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.

[jkalsi1]

See quantumlib/Cirq#8041

[sergeisakov]

Thank you for the proposed solution. In addition to copying the initial state, _BufferedStateVector creates a buffer in __init__. The buffer creation should also be toggled by this flag.

[pavoljuhas]

@sergeisakov - are simulations of 40-qubit state vectors even feasible in qsim?

On my cloudtop with 94GB of memory I am not even able to allocate such state vector once as np.zeros(2**40, dtype=np.int8) fails with the same _ArrayMemoryError as in the description. I can do x = np.arange(2**30, dtype=float), but then x.dot(x) takes about 160ms. For a 40-qubit state vector the dot product can be projected to take 1000 times more, i.e., almost 3 minutes. That sounds painstakingly slow as I'd expect there would be more full-state vector operations in a simulation run.

@mhucka - can you please check with the user for an exact example of that simulation or - if confidential - for a similar sized textbook case? I'd be curious if it is even feasible to handle 40 qubits with a C++ without an extra overhead for Python / numpy interfaces.