Rust中并行迭代器输出至stdout的性能优化问题

问题背景 #

需要计算结构体中向量的两两点积并输出到标准输出。使用Rayon并行实现时，println!因线程等待输出锁导致性能暴跌：

• 并行计算+println!耗时1.25秒（移除println!仅需151毫秒）
• 实际项目中耗时从2.5秒增至26秒，性能下降约1000%
• 串行计算+输出耗时约6秒

尝试预分配缓冲区统一输出，但因format!的内存分配开销，耗时反而增至1.43秒，希望找到既能保持输出顺序与串行一致，又能提升性能的优化方案。

优化方案 #

方案1：并行计算+串行输出（最简单高效） #

将计算与输出分离：并行完成所有点积计算并收集结果（保持顺序），再串行输出结果。这样避免并行任务中的IO锁竞争，实现成本极低。

use std::time::Instant;
use rayon::iter::{IntoParallelIterator, ParallelBridge, ParallelIterator};

struct Object {
    data: Vec<f64>,
}

fn dot(a: &[f64], b: &[f64]) -> f64 {
    a.iter().zip(b).map(|(a, b)| a * b).sum()
}

fn main() {
    let time = Instant::now();

    let mut objects = Vec::with_capacity(100);
    for _ in 0..objects.capacity() - 1 {
        objects.push(Object {
            data: vec![0.001; 1_000_000],
        });
    }

    let pairs = (0..objects.len() - 1).flat_map(|i| (i + 1..objects.len()).map(move |j| (i, j)));

    // 并行计算所有结果，保持顺序
    let results: Vec<_> = pairs
        .par_bridge()
        .into_par_iter()
        .map(|(i, j)| {
            let c = dot(&objects[i].data, &objects[j].data);
            (i, j, c)
        })
        .collect();

    // 串行输出所有结果
    for (i, j, c) in results {
        println!("{:2} {:2} -> {:<.20}", i, j, c);
    }

    let duration = time.elapsed();
    println!("time elapsed: {:.9?}", duration);
}

优势：

• 实现简单，无需复杂的缓冲区操作
• 计算阶段无IO开销，完全利用并行优势
• 输出阶段无锁竞争，串行写入效率更高
• 内存开销极小：100个对象仅需约116KB存储结果

方案2：预分配缓冲区+无分配格式化（极致性能） #

如果需要进一步优化输出速度，可预分配缓冲区，并用无分配的格式化库（itoa、ryu）直接写入缓冲区，避免format!的临时内存分配开销。

首先添加依赖：

[dependencies]
rayon = "1.10.0"
itoa = "1.0"
ryu = "1.0"

修改后的核心代码：

use std::time::Instant;
use std::io::Write;
use rayon::iter::{IntoParallelIterator, ParallelBridge, ParallelIterator};
use itoa;
use ryu;

struct Object {
    data: Vec<f64>,
}

fn dot(a: &[f64], b: &[f64]) -> f64 {
    a.iter().zip(b).map(|(a, b)| a * b).sum()
}

fn main() {
    let time = Instant::now();

    let mut objects = Vec::with_capacity(100);
    for _ in 0..objects.capacity() - 1 {
        objects.push(Object {
            data: vec![0.001; 1_000_000],
        });
    }

    let n = objects.len();
    let comparisons = (n * (n - 1)) / 2;
    // 准确计算每行的字节数："{:2} {:2} -> {:<.20}\n" 固定为30字节
    let mut buffer = vec![0u8; comparisons * 30];

    let pairs = (0..objects.len() - 1).flat_map(|i| (i + 1..objects.len()).map(move |j| (i, j)));

    buffer
        .chunks_mut(30)
        .zip(pairs)
        .par_bridge()
        .into_par_iter()
        .for_each(|(chunk, (i, j))| {
            let c = dot(&objects[i].data, &objects[j].data);
            let mut pos = 0;

            // 写入i（右对齐2位）
            let i_bytes = itoa::write_to_slice(&mut chunk[pos..], i).unwrap();
            pos += i_bytes;
            if i_bytes < 2 {
                chunk[pos] = b' ';
                pos += 1;
            }
            chunk[pos] = b' ';
            pos += 1;

            // 写入j（右对齐2位）
            let j_bytes = itoa::write_to_slice(&mut chunk[pos..], j).unwrap();
            pos += j_bytes;
            if j_bytes < 2 {
                chunk[pos] = b' ';
                pos += 1;
            }
            // 写入固定分隔符
            chunk[pos..pos+4].copy_from_slice(b" -> ");
            pos +=4;

            // 写入浮点数（左对齐20位）
            let float_str = ryu::Buffer::new().format_finite(c);
            let float_bytes = float_str.as_bytes();
            chunk[pos..pos+float_bytes.len()].copy_from_slice(float_bytes);
            pos += float_bytes.len();
            // 补空格到20位宽度
            while pos < 2 + 1 + 2 + 4 + 20 { // 前序部分总长度为9，加上20位浮点数区域
                chunk[pos] = b' ';
                pos +=1;
            }
            // 写入换行符
            chunk[pos] = b'\n';
        });

    // 一次性写入标准输出
    std::io::stdout().write_all(&buffer).unwrap();

    let duration = time.elapsed();
    println!("time elapsed: {:.9?}", duration);
}

优势：

• 完全消除并行任务中的内存分配和IO锁开销
• 一次性写入stdout，减少系统调用次数
• 严格保证输出顺序与串行版本一致

方案3：并行化点积计算（进一步提升计算性能） #

针对大向量的点积计算，可将点积本身也并行化，进一步缩短计算时间：

fn dot(a: &[f64], b: &[f64]) -> f64 {
    a.par_iter().zip(b.par_iter()).map(|(x, y)| x * y).sum()
}

优势：

• 对于百万级元素的向量，并行点积可显著降低计算耗时
• 与前两种方案兼容，可组合使用

总结 #

• 优先选择方案1，以最小的实现成本获得大幅性能提升
• 若追求极致性能，可采用方案2+方案3的组合
• 避免在并行任务中直接调用println!或format!，减少锁竞争和内存分配开销

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