Rust 中 `ndarray` `Array2<f64>` 值的并行计算

Question

我正在尝试加快我的计算速度，但我不确定如何最好地使用

.into_par_iter()

或

Zip::

板条箱中的一些

ndarray

选项，以便“正确地”并行计算以加快速度。

我尝试了两种方法，但我认为由于我对

Mutex

的理解，整个数组被“锁定”，这只会产生更多的开销并且计算速度比以前慢。

1. 例子：

//* 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. 例子：

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. 示例（最重要的并行化代码）：

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();

任何有关如何“正确”并行化代码的帮助以及有关如何更好地使用

Mutex

的更多信息将不胜感激。

Answer 1

感谢@adamreichold，我能够更多地了解我的问题，然后“正确地”并行化我的代码，从而显着加快我的计算速度。

现在开始解释：

1. 我的问题是，我试图找出一种方法来“按原样”并行化

   for

   循环。

2. ndarray

   可以以“经典”方式与

   for

   循环一起使用，但在函数式编程方式中似乎更有效，因为您正在应用例如

   for_each

   直接在

   object

   上，而不是使用索引循环，然后使用索引访问字段。
3. 功能性似乎是要走的路；包括

   Zip::

   方法的正确使用。

“更好”版本例如 1:

//* New version: par_for_each and indexed */
self.mol.wfn_total.HF_Matrices.S_matr = Array2::<f64>::zeros((n, n));

Zip::indexed(&mut self.mol.wfn_total.HF_Matrices.S_matr).par_for_each(|(i, j), s_val| {
    if i == j {
        *s_val = 1.0;
    } else if i >= j {
        //* Calculate only lower triangle matrix */
        *s_val = 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],
        )
    } else {
        *s_val = 0.0;
    }
});

// * Assign lower triangle to upper triangle with slices -> larger chunks
for i in 0..n - 1 {
    let slice = self
        .mol
        .wfn_total
        .HF_Matrices
        .S_matr
        .slice(s![i + 1..n, i])
        .to_shared();
    self.mol
        .wfn_total
        .HF_Matrices
        .S_matr
        .slice_mut(s![i, i + 1..n])
        .assign(&slice);
}

代码的小解释：

1. Zip::indexed

   需要一个

   &mut arr

   来索引

   arr

   并且

   par_for_each

   使用

   par_iter

   遍历所有元素。

2. (i,j)

   是索引到

   basis_set_cgtos

   所需的 2 个索引（这是一个

   Vec

   ）
3. 节省计算时间，因为只计算下三角矩阵然后复制到上半部分，所以我们基本上是说：

其余示例可以类似地并行化（我的 GitHub 存储库中的代码：SCF_MRJD_RS）