如何最好地并行化修改同一 Rust 向量的多个切片的代码？

Question

假设我们想要将向量的每个切片中的每个元素加倍（就地），其中切片由一系列对（开始、结束）位置定义。以下代码惯用地表达了意图，但由于并行内部向量的可变借用而无法编译

for_each

:

use rayon::prelude::*;

fn main() {
    let mut data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
    let slice_pairs = vec![(0, 3), (4, 7), (8, 10)];

    slice_pairs.into_par_iter().for_each(|(start, end)| {
        let slice = &mut data[start..end];
        for elem in slice.iter_mut() {
            *elem *= 2;
        }
    });

    println!("{:?}", data);
}

这里确实存在数据竞争的可能性 - 要排除它们，您需要检查切片是否重叠。问题是在 Rust 中实现这一点的最佳方法是什么，要么通过不安全的代码，要么通过安全的 API。以下代码使用

unsafe

来“继续执行此操作”；我的问题是是否有比下面更好的方法（将向量的基指针转换为 i64 并返回到“盲目”借用检查器来解决问题。）

use rayon::prelude::*;
use std::mem;

fn main() {
    let mut data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
    let slice_pairs = vec![(0, 4), (4, 7), (7, 10)];

    let ptr_outer = data.as_mut_ptr();
    let ptr_int : i64 = unsafe { mem::transmute(ptr_outer) };

    slice_pairs.into_par_iter().for_each(|(start, end)| {
        unsafe {
            let ptr : *mut i32 = mem::transmute(ptr_int);
            let slice = std::slice::from_raw_parts_mut(ptr.add(start), end - start);

            for elem in slice.iter_mut() {
                *elem *= 2;
            }
        }
    });

    println!("{:?}", data);
}

Answer 1

我建议首先将

slice_pairs

转换为一系列可变切片，然后并行使用所有这些切片。

可以使用

slice::split_at_mut()

将整个切片细分为多个独立的子切片（从借用检查器的角度来看）。
当然，

slice_pairs

中的索引必须有序且不能重叠，这样子切片才是正确的。

请注意，我尝试使用

.map().collect()

，而不是使用

.push()

进行显式循环，以构建切片序列，但我失败了...
编译器说

FnMut

中的

.map()

闭包无法返回引用；也许有人可以修复我的代码...

use rayon::prelude::*;

fn main() {
    let mut data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
    let slice_pairs = vec![(0, 3), (4, 7), (8, 10)];

    // build a sequence of mutable slices
    let mut slices = Vec::with_capacity(slice_pairs.len());
    let mut remaining = data.as_mut_slice();
    let mut idx = 0;
    for (start, end) in slice_pairs {
        let (_skip, tail) = remaining.split_at_mut(start - idx);
        let (sl, tail) = tail.split_at_mut(end - start);
        remaining = tail;
        idx = end;
        slices.push(sl);
    }
    println!("slices: {:?}", slices);

    // parallel usage of the mutable slices
    slices.into_par_iter().for_each(|sl| {
        for elem in sl.iter_mut() {
            *elem *= 2;
        }
    });

    println!("data: {:?}", data);
}
/*
slices: [[1, 2, 3], [5, 6, 7], [9, 10]]
data: [2, 4, 6, 4, 10, 12, 14, 8, 18, 20]
*/

Answer 2

您可以使用

split_at_mut()

使用安全代码将切片分成多个切片：

fn split_many<'a, T>(mut slice: &'a mut [T], regions: &[(usize, usize)]) -> Vec<&'a mut [T]> {
    let mut regions = regions.to_vec();
    regions.sort_by_key(|&(b, _e)| b);
    let mut ret = vec![];
    let mut offset = 0;
    for (b, e) in regions {
        let (chosen, rest) = slice.split_at_mut(e - offset);
        ret.push(&mut chosen[b - offset..]);
        offset = e;
        slice = rest;
    }
    ret
}

有了该辅助函数，您就可以以“明显”的方式实现就地操作：

fn main() {
    let mut data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
    let slice_pairs = vec![(0, 3), (4, 7), (8, 10)];

    split_many(&mut data, &slice_pairs)
        .into_par_iter()
        .for_each(|region| {
            for elem in region.iter_mut() {
                *elem *= 2;
            }
        });

    println!("{:?}", data);
}

游乐场