Rust中向量元素成对比较更新的多线程性能问题

问题描述 #

我尝试为需要成对比较并调整元素的向量实现多线程处理，没有用嵌套循环，而是生成不重复的usize类型ID对向量，确保同一ID不会分配到不同线程。

为绕过编译器安全检查，我使用了unsafe代码和指针操作——我能保证不同线程不会操作主向量的同一元素，但编译器仍判定向不同线程传递可变向量的部分而报错。

代码在启用和禁用scoped_threadpool时都能编译，但启用多线程后性能反而下降50倍，而非提升。

场景说明 #

我有一个包含N个结构体的向量（N是2的幂，即N=1<<m），每个结构体包含一个Vec<f64>（共10个浮点数，代表3D粒子的位置、速度、加速度及一个额外比较值）。需要通过成对比较位置（第0-2位浮点数）来更新粒子间的加速度（第6-8位浮点数）。

我实现了get_ijpairs_vec函数生成唯一、无重叠的元素ID对向量，共需循环N-1次这类向量。但使用scoped_threadpool时仍无法通过安全检查，于是将坐标向量指针转为usize传入线程池，再在内部转回指针。

相关代码如下：

get_ijpairs_vec 函数

pub fn get_ijpairs_vec ( //in impl GasBox
  &self, //GasBox
  segment: usize, //there are N-1 possible vectors, pick one.
  result: &mut Vec<usize>, //where to build the vector.
) {
  let maxblock = self.nmol/2 ;
  let shift: usize = segment/2 + 1 ;
  let block: usize = 1<<(shift.trailing_zeros()) as usize ;
  let nblock = maxblock/block ;
  let skip = if segment.trailing_zeros() > 0 { 0 } else { block } ;
  let needed: usize = (self.nmol as isize - result.capacity() as isize) as usize ;
  if needed > 0 { result.extend(vec![0usize; needed]) } ;
  let mut k: usize = 0 ;
  let mut i: usize = skip ;
  for _b in 0..nblock{
    for _c in 0..block {
      let mut j: usize = i + shift ;
      if j >= self.nmol {
        j -= self.nmol ;
        result[2*k] = j ;
        result[2*k+1] = i ;
      } else {
        result[2*k] = i ;
        result[2*k+1] = j ;
      }
      i += 1 ;
      k += 1 ;
      if i >= self.nmol { i -= self.nmol ; }
    }
    i += block ;
    if i >= self.nmol { i -= self.nmol ; }
  }
}

update_accelerations 函数

pub fn update_accelerations( //in impl GasBox
  &mut self, //GasBox
) {
  let mut pairs: Vec<usize> = Vec::new() ;
  for i in (0..).take(self.nmol) {
    self.molecules[i].zero_accel() ; //sets elements 6-8 of phase to 0.0.
  }
  for k in (0..).take(self.nmol-1) { //N-1 possible vectors of pairs
    self.get_ijpairs_vec(k, &mut pairs) ;
//      self.pool.scoped(|scope|{
        for chunk in pairs.chunks_exact_mut(2) {
          let iptr: usize = self.molecules[chunk[0]].phase.as_ptr() as usize ;
          let jptr: usize = self.molecules[chunk[1]].phase.as_ptr() as usize ;
//          scope.execute(move || {
              GasBox::fix_accel(iptr, jptr) ;
//          } ) ;
        }
//    } ) ;
  }
}

fix_accel 函数

pub fn fix_accel(
  i: usize,
  j: usize,
) {
  let mut acc: Vec<f64> = vec![0.0f64; 3] ;
  unsafe {
    let iptr: *mut f64 = i as *mut f64 ;
    let jptr: *mut f64 = j as *mut f64 ;
    let mut sep2: f64 = 0.0 ;
    for k in (0..).take(3) {
      let tmp: f64 = *iptr.offset(k) - *jptr.offset(k) ;
      sep2 += tmp*tmp ;
    }
    //do some math to set up acc based on sep2.
    for k in (0..).take(3) {
      let tmp = *iptr.offset(k) - *jptr.offset(k) ;
      *iptr.offset(k+6) -= tmp*acc[k] ;
      *jptr.offset(k+6) += tmp*acc[k] ;
    }
  } //end the unsafe block
}

当前代码运行正常，但多线程版本在8核机器上慢50倍。我认为该配对系统类似线程安全的chunks_mut变体，想请教：

1. 当前实现是否仍存在不安全情况（如竞态条件）？
2. 若实现安全，如何通过多线程获得性能提升？

问题解答 #

1. 当前实现是否存在不安全情况？ #

你的实现确实存在潜在的不安全问题，主要体现在以下几点：

• 指针有效性未保障：你将phase.as_ptr()转为usize传入线程，但无法保证线程执行期间原Vec<f64>不会重新分配（比如后续修改zero_accel或其他操作改变phase长度）。一旦原Vec扩容，指针就会失效，访问会触发未定义行为。
• 内存可见性缺失：多线程下修改内存需要保证可见性，直接用裸指针操作没有使用原子操作或编译器栅栏，可能导致线程间修改无法及时同步，甚至出现指令重排序问题。
• 生命周期绕过风险：通过usize传递指针跳过了Rust的生命周期检查，无法确保指针使用期间原数据仍存活。如果GasBox在线程执行前被销毁，指针会变成悬垂指针。

不过就当前代码逻辑而言，若你严格保证：

1. update_accelerations执行期间，self.molecules和每个phase Vec都不会被重新分配或销毁；
2. 生成的ID对完全无重叠，不会有两个线程操作同一个粒子的加速度；
   那么不会出现竞态条件，因为每个粒子的加速度只会被一个线程（或成对线程仅修改各自粒子）操作。但上述unsafe风险依然存在。

2. 如何通过多线程获得性能提升？ #

多线程版本变慢50倍的核心原因是线程调度开销远大于计算任务的开销。每个fix_accel仅做少量浮点数运算，而scope.execute的任务调度、线程切换成本远超过计算收益，直接拖垮整体性能。可以从以下方向优化：

（1）批量打包任务，减少调度开销

不要为每一对粒子创建单独任务，而是将多对粒子打包成大任务交给线程池。比如把pairs分成包含几十上百对的批次，每个线程处理一个批次：

self.pool.scoped(|scope| {
    // 按每200对为一个批次拆分
    for batch in pairs.chunks(200) {
        let batch_clone = batch.to_vec();
        // 传递molecules的指针（需保证整个函数执行期间molecules存活）
        let molecules_ptr = self.molecules.as_ptr() as usize;
        scope.execute(move || {
            for chunk in batch_clone.chunks_exact(2) {
                unsafe {
                    let molecules = &*(molecules_ptr as *mut [Molecule]);
                    let iptr = molecules[chunk[0]].phase.as_mut_ptr();
                    let jptr = molecules[chunk[1]].phase.as_mut_ptr();
                    GasBox::fix_accel_raw(iptr, jptr);
                }
            }
        });
    }
});

同时修改fix_accel直接接收指针，避免反复转换：

pub fn fix_accel_raw(
    iptr: *mut f64,
    jptr: *mut f64,
) {
    let mut acc: [f64; 3] = [0.0; 3]; // 用数组代替Vec，消除堆分配开销
    unsafe {
        let mut sep2: f64 = 0.0;
        for k in 0..3 {
            let tmp = *iptr.offset(k) - *jptr.offset(k);
            sep2 += tmp * tmp;
        }
        // 补充计算acc的逻辑
        for k in 0..3 {
            let tmp = *iptr.offset(k) - *jptr.offset(k);
            *iptr.offset(k + 6) -= tmp * acc[k];
            *jptr.offset(k + 6) += tmp * acc[k];
        }
    }
}

（2）用rayon替代scoped_threadpool，优化任务调度

rayon库提供了高效的并行迭代器，能自动处理任务拆分和调度，比手动使用scoped_threadpool更高效，代码也更简洁：
首先在Cargo.toml添加依赖：

rayon = "1.7"

然后修改update_accelerations：

use rayon::prelude::*;

pub fn update_accelerations(&mut self) {
    // 并行重置加速度，效率更高
    self.molecules.par_iter_mut().for_each(|mol| mol.zero_accel());

    for k in 0..self.nmol-1 {
        let mut pairs: Vec<usize> = Vec::with_capacity(self.nmol);
        self.get_ijpairs_vec(k, &mut pairs);
        
        // rayon自动拆分任务，避免过度调度
        pairs.par_chunks_exact(2).for_each(|chunk| {
            let i = chunk[0];
            let j = chunk[1];
            unsafe {
                let iptr = self.molecules[i].phase.as_mut_ptr();
                let jptr = self.molecules[j].phase.as_mut_ptr();
                GasBox::fix_accel_raw(iptr, jptr);
            }
        });
    }
}

（3）减少不必要的内存分配

• 提前为pairs预分配足够容量（比如pairs.reserve(self.nmol)），避免反复扩容的开销；
• 将fix_accel中的Vec<f64>换成[f64;3]数组，消除堆分配——多线程下堆分配的锁竞争也是性能下降的原因之一。

（4）优化缓存局部性

当前粒子的phase Vec分散在堆上，访问时缓存命中率低。可以将粒子数据改为结构数组存储：

struct Phase {
    pos: [f64; 3],
    vel: [f64; 3],
    accel: [f64; 3],
    extra: f64,
}

// GasBox中用Vec<Phase>替代Vec<{ phase: Vec<f64> }>

连续的粒子数据会被加载到CPU缓存中，大幅提升访问速度，多线程下的缓存利用率也会更高。

内容的提问来源于stack exchange，提问作者C. Bartels