Shared state concurrency help -- rayon

[b-vanstraaten]

I am trying to implement a density matrix simulator of a quantum computer. The quantum gates that act of the qubit density matrix state do so in an embarrassingly parallel fashion. More precisely:

A single qubit gate operation on the density matrix of size 2^N by 2^N, where N is the number of qubits, can be computed through doing a 2 by 2 matric multiplication on the elements of 2^(2N - 2) different slices. The elements of one slice are unique to that slice. There is no sharing between slices.

I wish to use rayons parallel for_each function to iterate over the unique elements in parallel. The density_matrix (an nalgebra matrix) can be Gigabytes in size. Therefore, copying the data is impractical.

The code I have so far is:

fn single_qubit_gate(state: &mut State, qubit: &Qubit, u: Matrix2x2) {
    debug!("density matrix before:\n{}", state.density_matrix);

    let mut density_matrix = Arc::new( &state.density_matrix);

    (0..1 << state.number_of_qubits)
        .into_par_iter()
        .step_by(2)
        .for_each(|i_offset:usize| {

            let swap_qubits: Box<dyn Fn(usize) -> Qubit> = match qubit {
                0 => Box::new(|x: usize| x),
                _ => Box::new(|x: usize| swap_two_qubits(x, (&qubit, &0))),
            };

            let mut rho = Matrix2x2::zeros();
            (0..(1 << state.number_of_qubits))
                .step_by(2)
                .for_each(|j_offset| {

                    for (i, j) in iproduct!(0..2, 0..2) {
                        rho[(i, j)] = density_matrix[(swap_qubits(i_offset + i), swap_qubits(j_offset + j))];
                    }

                    rho = u * rho * u.adjoint();

                    for (i, j) in iproduct!(0..2, 0..2) {
                        density_matrix[(swap_qubits(i_offset + i), swap_qubits(j_offset + j))] = rho[(i, j)]
                    }

                });
    });

However, this errors with the stack trace:

error[E0596]: cannot borrow data in an `Arc` as mutable
  --> src/gates.rs:85:25
   |
85 |                         density_matrix[(swap_qubits(i_offset + i), swap_qubits(j_offset + j))] = rho[(i, j)]
   |                         ^^^^^^^^^^^^^^ cannot borrow as mutable
   |
   = help: trait `DerefMut` is required to modify through a dereference, but it is not implemented for `Arc<&Matrix<Complex<f64>, Dynamic, Dynamic, VecStorage<Complex<f64>, Dynamic, Dynamic>>>`

I would really appreciate a pointer in the right direction, using the Arc was rather my final hope.

[kornel]

You will not be able to use mutable indexing with [ ] in parallel code (like density_matrix[i]).

The [ ] (IndexMut) operator always exclusively borrows the entire object whenever you try to change one element. It isn't even trying to check if the elements don't overlap, and always assumes the most restrictive interpretation.

You will have to split density_matrix using iterator methods like chunks_mut, so that for_each does not have any shared mutable state, and each callback gets an independent object.

Alternatively, you will have to use raw pointers and unsafe code to force it through. In some cases it's possible to write a safe abstraction on top:

https://blog.rom1v.com/2019/04/implementing-tile-encoding-in-rav1e/