Rust多线程采样应选哪种RNG？并行化正态分布采样遇错求解

问题根源 #

1. 线程RNG不可跨线程传递：rand::thread_rng()返回的是线程本地RNG，未实现Send trait，无法在Rayon的并行迭代中跨线程共享。
2. 可变共享状态冲突：原代码通过可变借用修改全局preds向量，Rayon的并行闭包要求为Fn类型（需满足线程安全），但捕获可变引用的闭包是FnMut，不符合要求。

解决方案：每个并行任务独立生成数据与RNG #

核心思路是让每个并行任务完全独立：自己创建RNG，生成对应分布的采样向量，最后合并所有结果，避免共享可变状态和跨线程传递RNG。

方案1：用线程本地RNG快速实现（推荐）

利用rand::thread_rng()的线程本地特性，每个并行任务在自己的线程中获取专属RNG，无需跨线程传递：

use rand_distr::{Normal, Distribution};
use std::time::Instant;
use rayon::prelude::*;
use rand::Rng;

fn rprednorm(n: i32, means: Vec<f64>, sds: Vec<f64>) -> Vec<Vec<f64>> {
    // 并行迭代每个均值-标准差对，生成对应采样向量
    means.into_par_iter()
        .zip(sds.into_par_iter())
        .map(|(mean, sd)| {
            // 每个任务获取当前线程的本地RNG
            let mut rng = rand::thread_rng();
            let normal = Normal::new(mean, sd).unwrap();
            // 生成n次采样并收集为子向量
            (0..n).map(|_| normal.sample(&mut rng)).collect()
        })
        .collect()
}

fn main() {
    let means = vec![0.0; 67000];
    let sds = vec![1.0; 67000];
    let start = Instant::now();
    let preds = rprednorm(100, means, sds);
    let duration = start.elapsed();
    
    println!("{:?}", duration);
}

方案2：用StdRng实现可复现采样

如果需要固定种子以复现结果，可预先生成独立种子，为每个任务初始化StdRng（StdRng实现了Send，支持跨线程传递种子）：

use rand_distr::{Normal, Distribution};
use std::time::Instant;
use rayon::prelude::*;
use rand::{Rng, SeedableRng, rngs::StdRng};
use rand::os::OsRng;

fn rprednorm(n: i32, means: Vec<f64>, sds: Vec<f64>) -> Vec<Vec<f64>> {
    // 为每个任务生成独立种子
    let seeds: Vec<_> = (0..means.len()).map(|_| OsRng.gen()).collect();
    
    means.into_par_iter()
        .zip(sds.into_par_iter())
        .zip(seeds.into_par_iter())
        .map(|((mean, sd), seed)| {
            let mut rng = StdRng::from_seed(seed);
            let normal = Normal::new(mean, sd).unwrap();
            (0..n).map(|_| normal.sample(&mut rng)).collect()
        })
        .collect()
}

fn main() {
    let means = vec![0.0; 67000];
    let sds = vec![1.0; 67000];
    let start = Instant::now();
    let preds = rprednorm(100, means, sds);
    let duration = start.elapsed();
    
    println!("{:?}", duration);
}

关键改进点 #

• 移除全局共享的RNG和preds向量，每个任务独立生成数据，避免线程安全问题。
• 用into_par_iter()替代into_iter()实现并行，通过map收集结果而非修改共享状态，符合Rayon的闭包要求。

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