Parallel calculation of Array2<f64> values

[MartinRJDagleish]

Hey,
I am trying to speed up my calculations and I am not sure, how to best use .into_par_iter() or some of the Zip:: options in order to "correctly" parallelize the computation to speed up things.

I have tried 2 ways, but I think due to my understanding of Mutex the whole Array gets "locked", which just creates more overhead and the calculations are slower than before.

1. example:

//* Step 2.1: Calculate the overlap matrix S
let S_matr_tmp = Array2::<f64>::zeros((n, n));
let S_matr_mutex = Mutex::new(S_matr_tmp);

(0..n).into_par_iter().for_each(|i| {
    (0..=i).into_par_iter().for_each(|j| {
        let mut s = S_matr_mutex.lock().unwrap();
        if i == j {
            s[(i, j)] = 1.0;
            // s = 1.0;
        } else {
            s[(i, j)] = calc_overlap_int_cgto(
                &self.mol.wfn_total.basis_set_total.basis_set_cgtos[i],
                &self.mol.wfn_total.basis_set_total.basis_set_cgtos[j],
            );
            s[(j, i)] = s[(i, j)];
        }
    })
});

self.mol.wfn_total.HF_Matrices.S_matr = S_matr_mutex.into_inner().unwrap();

2. example:

let D_matr_mutex = Mutex::new(Array2::<f64>::zeros((no_cgtos, no_cgtos)));

//* D^0 matrix */
//* Trying to parallelize it */
Zip::indexed(C_matr_AO_basis.axis_iter(Axis(0))).par_for_each(|mu, row1| {
    Zip::indexed(C_matr_AO_basis.outer_iter()).par_for_each(|nu, row2| {
        let mut d = D_matr_mutex.lock().unwrap();
        let slice1 = row1.slice(s![..no_occ_orb]);
        let slice2 = row2.slice(s![..no_occ_orb]);
        d[(mu, nu)] = slice1.dot(&slice2);
    });
});

let mut D_matr = D_matr_mutex.into_inner().unwrap();

3. example (most important code to parallelize):

let F_matr_mutex = Mutex::new(F_matr);

(0..no_cgtos).into_par_iter().for_each(|mu| {
    (0..no_cgtos).into_par_iter().for_each(|nu| {
        (0..no_cgtos).into_par_iter().for_each(|lambda| {
            (0..no_cgtos).into_par_iter().for_each(|sigma| {
                let mu_nu_lambda_sigma =
                    calc_ijkl_idx(mu + 1, nu + 1, lambda + 1, sigma + 1);
                let mu_lambda_nu_sigma =
                    calc_ijkl_idx(mu + 1, lambda + 1, nu + 1, sigma + 1);
                let mut f = F_matr_mutex.lock().unwrap();
                f[(mu, nu)] += D_matr[(lambda, sigma)]
                    * (2.0
                        * self.mol.wfn_total.HF_Matrices.ERI_arr1[mu_nu_lambda_sigma]
                        - self.mol.wfn_total.HF_Matrices.ERI_arr1[mu_lambda_nu_sigma]);
            });
        });
    });
});

let F_matr = F_matr_mutex.into_inner().unwrap();

Any help on how to parallelize the code "correctly" and more information on how to use Mutex better would be appreciated.

[adamreichold]

Sorry for the very brief answer, but I think your best bet is to use the Zip approach, but use the par_map_collect or par_map_assign_into to gather up the results instead of locking an output matrix.

[MartinRJDagleish]

@adamreichold Thank you for your hint.

After a lot of trying around and not too much documentation, I found a way to do this, but I am convinced that this solution is "awful". There is definetely a better way of doing this, but this seemed to work:

//* New version: par_map_assign_into */
let S_matr_tmp = Array2::<f64>::zeros((n, n));
self.mol.wfn_total.HF_Matrices.S_matr = Array2::<f64>::zeros((n, n));

Zip::indexed(&S_matr_tmp).par_map_assign_into(
    &mut self.mol.wfn_total.HF_Matrices.S_matr,
    |(i, j), x| {
        if i == j {
            1.0
        } else {
            calc_overlap_int_cgto(
                &self.mol.wfn_total.basis_set_total.basis_set_cgtos[i],
                &self.mol.wfn_total.basis_set_total.basis_set_cgtos[j],
            )
        }
    },
);

The main problem is that I need the index to access the Vec of CGTOs for calling my calc_overlap_int_cgto function, but I cannot do Zip::indexed(&S_matr).par_map_assign_into(&mut S_matr, …) of course because you can not have a mutable and a read-only reference to S_matr at the same time and in the same function call.

Any help would be deeply appreciated. :)

[MartinRJDagleish]

@adamreichold Thank you a lot. I first overlooked your &mut self and only had a &self which did not work, but your solution works perfectly for me, so thanks a lot. :)

Just out of curiosity: is there any advantage in using the for_each (or here par_for_each) approach rather than building a way of parallelizing the "normal" for loops?

[adamreichold]

Just out of curiosity: is there any advantage in using the for_each (or here par_for_each) approach rather than building a way of parallelizing the "normal" for loops?

I'd say this is ndarray's way of parallelizing nested for loops.

If you are thinking about OpenMP et al., this has the advantage of being library code which does not need compiler support (and yet is memory safe).

If you are thinking Rayon's parallel iterators, the methods defined on Zip have the potential for more optimizations behind the scenes as they provide a much more limited interface, even though this does not apply in the particular case of par_for_each. But for example, par_map_collect will efficiently collect the results of the parallel computations into a single array without initializing it to zeros first.

[MartinRJDagleish]

@adamreichold Ok. Very interesting. Yes I was thinking of OpenMP.

I do not have a huge background in C++, but I know a little and also my main subject area (computational chemistry) heavily relies upon parallelized code and a lot of "number crunching".

Anyway thank you very much for your help of understanding the way of parallelizing in ndarray.

If I may ask one follow-up question, I am try to reduce computation and only calculate the lower (or upper) triangular matrix and then assign the lower (upper) to the upper (lower), which I have written in this rather 'redundant' way, because I could not get the correct way to work.

My attempt:

self.mol.wfn_total.HF_Matrices.V_ne_matr = Array2::<f64>::zeros((n, n));
Zip::indexed(&mut self.mol.wfn_total.HF_Matrices.V_ne_matr).par_for_each(
    |(idx1, idx2), v_ne_val| {
        if idx1 >= idx2 { // COMPUTING THE LOWER HALF
            for (atom_idx, atom_pos) in self
                .mol
                .geom_obj
                .geom_matr
                .axis_iter(ndarray::Axis(0))
                .enumerate()
            {
                *v_ne_val += (-self.mol.geom_obj.Z_vals[atom_idx] as f64)
                    * calc_nuc_attr_int_cgto(
                        &self.mol.wfn_total.basis_set_total.basis_set_cgtos[idx1],
                        &self.mol.wfn_total.basis_set_total.basis_set_cgtos[idx2],
                        &atom_pos.to_owned(),
                    );
            }
            // self.mol.wfn_total.HF_Matrices.V_ne_matr[(idx1, idx2)] = *v_ne_val;
        } else {
            *v_ne_val = 0.0;
        }
    },
);

//TODO: is there a better way to copy lower triangle to upper triangle in parallel?
let V_matr_lower = self.mol.wfn_total.HF_Matrices.V_ne_matr.t().to_owned();

Zip::indexed(&mut self.mol.wfn_total.HF_Matrices.V_ne_matr)
    .and(&V_matr_lower)
    .par_for_each(|(idx1, idx2), v_ne_val, v_ne_val_t| {
        if idx1 < idx2 {
            *v_ne_val = *v_ne_val_t;
        }
    });

This way I am at least using parallel 'assignment', but the map_assign_into does not work in this way (if I understood that correctly).

[adamreichold]

is there a better way to copy lower triangle to upper triangle in parallel?

To be honest, I'd be highly surprised if this copying which is surely memory bound would measurably benefit from parallelization. I would suspect that trying to use slicing and ArrayBase::assign to copy whole runs of rows/columns at once might be the larger optimization opportunity if can be made to end up calling into the platform's native memcpy-like routines.

[MartinRJDagleish]

@adamreichold Perfect. That's a good idea. After trying around for a bit, my attempt:

for i in (0..n - 1) {
    let slice = self.mol.wfn_total.HF_Matrices.V_ne_matr.slice(s![i + 1..n, i]).to_shared(); // to_shared or to_owned?
    self.mol.wfn_total.HF_Matrices.V_ne_matr.slice_mut(s![i, i + 1..n]).assign(&slice);
}

The offset +1 allows to not reassign elements on the diagonal of the matrix. (I don't know if this is any better)