Rust并行填充矩阵：同步向量转普通二维向量的效率疑问

我正在学习Rust，刚看完《Rust编程圣经》里的「无畏并发（Fearless Concurrency）」章节，试着用并行方式填充矩阵的行。因为矩阵规模大，我把行块分配给每个线程，让线程负责填充块内所有行的列。

初始化矩阵的代码如下：

// `rows` 是一个 nrows x ncols 的矩阵，将并行填充，每个线程处理一块行
let mut rows = vec![];
for _ in 0..nrows {
    // 每一行必须是 `Arc<Mutex<_>>`，因为要并行填充
    let row = Arc::new(Mutex::new(vec![0; ncols]));
    rows.push(row);
}

之后用thread::spawn完成并行处理，一切正常，但最终得到的rows类型是Vec<Arc<Mutex<Vec<usize>>>>，而后续只需要一个拥有所有权的Vec<Vec<usize>>，所以我用下面的代码转换：

// 这里是我的疑问
// 现在 `rows` 的类型是 `Vec<Arc<Mutex<Vec<usize>>>>`
// 我们需要 `Vec<Vec<usize>>`
let mut ret_rows = vec![];
for row in rows {
    ret_rows.push(row.lock().unwrap().clone())
}
// 返回一个拥有所有权的向量
Vec::from_iter(ret_rows)

我的问题是：这种转换方式高效吗？

我怀疑.clone()（甚至Vec::from_iter()）会复制所有数据，这不是我想要的，实际上也完全没必要。

复现代码 #

use std::sync::{Arc, Mutex};
use std::thread;

/// 填充单行所有列的示例函数
fn worker_row(row: &mut Vec<usize>, f: usize) {
    for i in 0..row.len() {
        row[i] = 2 * i * f  // 随便填充一些数据
    }
}

/// 创建并返回一个 nrows x ncols 的矩阵
fn build_matrix(ncols: usize, nrows: usize) -> Vec<Vec<usize>> {

    // `rows` 是一个 nrows x ncols 的矩阵，将并行填充，每个线程处理一块行
    let mut rows = vec![];

    // 用零初始化矩阵
    for _ in 0..nrows {
        // 每一行必须是 `Arc<Mutex<_>>`，因为要并行填充
        let row = Arc::new(Mutex::new(vec![0; ncols]));
        rows.push(row);
    }

    // 每个线程处理一块行
    let nthreads = 8;
    let chunk_size = rows.len() / nthreads; // 忽略剩余行，不影响示例
    let mut handles = vec![];

    // 生成线程
    for i in 0..nthreads {
        let mut rows_i = vec![];
        // 克隆当前线程要处理的行
        for r in i*chunk_size..(i+1)*chunk_size {
            rows_i.push(Arc::clone(&rows[r]));
        }
        let handle = thread::spawn(move || {
            for r in i*chunk_size..(i+1)*chunk_size {
                let j = r - i * chunk_size;
                // 获取锁
                let mut row = rows_i[j].lock().unwrap();
                worker_row(&mut row, r);
            }
        });
        handles.push(handle);
    }
    for handle in handles {
        handle.join().unwrap();
    }

    // 这里是我的疑问
    // 现在 `rows` 的类型是 `Vec<Arc<Mutex<Vec<usize>>>>`
    // 我需要 `Vec<Vec<usize>>`
    // 这种转换高效吗？
    let mut ret_rows = vec![];
    for row in rows {
        ret_rows.push(row.lock().unwrap().clone())
    }
    // 返回一个拥有所有权的向量
    Vec::from_iter(ret_rows)
}


fn main() {
    let rows = build_matrix(10, 24);
    for row in rows {
        println!("{:?}", row);
    }
}

解决方案 #

你的怀疑是对的，当前的转换方式确实会复制所有矩阵数据，这完全是不必要的性能损耗。因为所有线程已经完成工作，Arc的引用计数此时应该只有1（主线程持有的那个），Mutex也不再需要，我们可以直接提取内部的Vec<usize>而无需克隆。

优化后的转换代码

let mut ret_rows = Vec::with_capacity(rows.len());
for row in rows {
    // 尝试将 Arc 转换为内部的 Mutex
    match Arc::try_unwrap(row) {
        Ok(mutex) => {
            // 尝试将 Mutex 转换为内部的 Vec
            match mutex.into_inner() {
                Ok(vec) => ret_rows.push(vec),
                Err(_) => panic!("Mutex is poisoned"),
            }
        }
        Err(_) => panic!("Arc has multiple references"),
    }
}
ret_rows

为什么这样更高效？

• Arc::try_unwrap：当Arc的引用计数为1时，这个方法会直接返回内部的Mutex，不需要复制数据，只是转移所有权。
• Mutex::into_inner：当Mutex没有被锁定时，这个方法会返回内部的Vec，同样是转移所有权，没有复制操作。
• Vec::with_capacity：预先分配足够的容量，避免后续push时的内存重新分配。

额外优化：调整初始化和线程分配逻辑

其实你还可以进一步优化初始化步骤，避免一开始就用Arc<Mutex>包裹每一行。因为你是按行块分配给线程，每个线程只处理自己块内的行，这些行不会被其他线程访问，完全不需要Mutex：

fn build_matrix(ncols: usize, nrows: usize) -> Vec<Vec<usize>> {
    // 预先分配整个矩阵的内存，每个线程拿到自己的行块的可变引用
    let mut matrix = vec![vec![0; ncols]; nrows];
    let nthreads = 8;
    let chunk_size = matrix.len() / nthreads;
    let mut handles = vec![];

    // 将矩阵拆分为多个可变切片，每个线程处理一个切片
    for i in 0..nthreads {
        let start = i * chunk_size;
        let end = (i + 1) * chunk_size;
        // 注意：这里需要用 split_at_mut 来安全地获取可变切片，避免别名问题
        let (_, rest) = matrix.split_at_mut(start);
        let (chunk, _) = rest.split_at_mut(end - start);
        let handle = thread::spawn(move || {
            for (idx, row) in chunk.iter_mut().enumerate() {
                let row_idx = start + idx;
                worker_row(row, row_idx);
            }
        });
        handles.push(handle);
    }

    for handle in handles {
        handle.join().unwrap();
    }

    matrix
}

这个版本的效率更高：

• 不需要Arc和Mutex的开销，减少了同步成本。
• 完全没有数据复制，直接在原始矩阵上修改，最后直接返回所有权。
• 避免了Arc和Mutex的内存分配，减少了内存开销。

内容的提问来源于stack exchange，提问作者Raf