如何用Rayon实现Rust矩阵块填充函数的并行化改造？

首先假设你的单线程实现大致如下（典型的块填充逻辑）：

// 单线程版本示例
fn square_matrix_block_fill(
    big_matrix: &mut Vec<Vec<f64>>,
    blocks: &[(usize, usize, Vec<Vec<f64>>)], // (起始行, 起始列, 子矩阵块)
) {
    for &(start_row, start_col, ref block) in blocks {
        let block_rows = block.len();
        let block_cols = block[0].len();
        for i in 0..block_rows {
            for j in 0..block_cols {
                big_matrix[start_row + i][start_col + j] = block[i][j];
            }
        }
    }
}

无重叠写入区域的并行实现（最优性能） #

如果你的子矩阵块写入大矩阵的区域完全不重叠，可以直接用Rayon并行遍历块，无需额外同步：

use rayon::prelude::*;

fn square_matrix_block_fill_parallel(
    big_matrix: &mut Vec<Vec<f64>>,
    blocks: &[(usize, usize, Vec<Vec<f64>>)],
) {
    // 并行处理每个子矩阵块
    blocks.par_iter().for_each(|&(start_row, start_col, ref block)| {
        let block_rows = block.len();
        let block_cols = block[0].len();
        
        // 遍历块的每一行，对应写入大矩阵的目标行
        for i in 0..block_rows {
            if let Some(target_row) = big_matrix.get_mut(start_row + i) {
                // 用copy_from_slice替换逐元素赋值，提升效率
                target_row[start_col..start_col + block_cols]
                    .copy_from_slice(&block[i]);
            }
        }
    });
}

说明 #

• Rayon的par_iter()会自动将块分配到多个线程执行，由于写入区域无重叠，不会出现数据竞争。
• 使用copy_from_slice比嵌套循环逐元素赋值更高效，底层调用内存拷贝指令。

存在重叠写入区域的并行实现 #

如果多个子矩阵块会写入大矩阵的同一位置，需要处理数据竞争，以下提供两种常见场景的解决方案：

场景1：覆盖写入（保留最后一次写入的值） #

先并行收集所有待写入的位置和最终值，再一次性写入大矩阵：

use rayon::prelude::*;
use std::collections::HashMap;

fn square_matrix_block_fill_parallel_overlap(
    big_matrix: &mut Vec<Vec<f64>>,
    blocks: &[(usize, usize, Vec<Vec<f64>>)],
) {
    // 并行收集所有位置的最终值（HashMap自动覆盖重复键，保留最后插入的值）
    let updates: HashMap<(usize, usize), f64> = blocks
        .par_iter()
        .flat_map(|&(start_row, start_col, ref block)| {
            block.iter().enumerate().flat_map(move |(i, row)| {
                row.iter().enumerate().map(move |(j, &val)| {
                    ((start_row + i, start_col + j), val)
                })
            })
        })
        .collect();
    
    // 单线程写入大矩阵，避免并行写入的竞争
    for ((row, col), val) in updates {
        big_matrix[row][col] = val;
    }
}

场景2：累加写入（将所有块的对应位置值相加） #

用带锁的HashMap并行累加值，再写入大矩阵：

use rayon::prelude::*;
use std::collections::HashMap;
use std::sync::Mutex;

fn square_matrix_block_fill_parallel_sum(
    big_matrix: &mut Vec<Vec<f64>>,
    blocks: &[(usize, usize, Vec<Vec<f64>>)],
) {
    let updates = Mutex::new(HashMap::new());
    
    blocks.par_iter().for_each(|&(start_row, start_col, ref block)| {
        let mut map = updates.lock().unwrap();
        for (i, row) in block.iter().enumerate() {
            for (j, &val) in row.iter().enumerate() {
                *map.entry((start_row + i, start_col + j)).or_insert(0.0) += val;
            }
        }
    });
    
    // 将累加结果写入大矩阵
    for ((row, col), val) in updates.into_inner().unwrap() {
        big_matrix[row][col] += val;
    }
}

说明 #

• 优先使用预收集的方式（如覆盖写入的实现），比直接给大矩阵加锁的性能更高，因为锁的竞争频率更低。
• 如果重叠区域极少，也可以给大矩阵的行分别加锁，但实现复杂度更高，不如预收集方式简洁。

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