Rust 中交错 RGB 通道（从 [R, R, G, G, B, B] 到 [R, G, B, R, G, B]）的最快方法

Question

我一直致力于寻找一种更快的方法来将向量的内容与 [R,R,R,...,G,G,G,...,B,B,B 形式的数据交错,...] 到 [R,G,B,R,G,B,R,G,B,...]。目前我有这段代码，对于大小为 (1024*1024*3) 的向量运行大约需要 60-200ms，但如果可能的话，我需要它在 10ms 范围内。我尝试过一些函数方法，但除非我做了一些非常错误的事情，否则它们会将时间增加到 500-700 毫秒。

static TILE_LENGTH: usize = 1024 * 1024;

let channels = vec![0; TILE_LENGTH * 3];

let mut tile = Vec::with_capacity(TILE_LENGTH * 3);
for i in 0..TILE_LENGTH {
    tile.push(channels[i]);
    tile.push(channels[i + (TILE_LENGTH)]);
    tile.push(channels[i + (TILE_LENGTH * 2)]);
}

Answer 1

我做的第一件事就是重新创建你的代码。

pub fn plain_index<'a>(channels: &[u8], tile: &'a mut Vec<u8>) -> &'a [u8] {
    tile.clear();
    tile.reserve(channels.len());

    let len_each = channels.len() / 3;

    for i in 0..len_each {
        tile.push(channels[i]);
        tile.push(channels[i + len_each]);
        tile.push(channels[i + len_each * 2]);
    }

    tile
}

我已将这两个Vec分配移至函数之外。如果你的程序只需要调用这个函数一次，那也没关系，但是如果你需要调用它很多次，你可以重用这个缓冲区，这样可以节省很多时间。这个运行时间为32 毫秒。--release评判表现时记得要跑。

然后，我制作了一些替代版本，最终得到了这个。

pub fn zip_extend<'a>(channels: &[u8], tile: &'a mut Vec<u8>) -> &'a [u8] {
    tile.clear();
    tile.reserve(channels.len());

    let len_each = channels.len() / 3;
    let rs = &channels[..len_each];
    let gs = &channels[len_each..len_each * 2];
    let bs = &channels[len_each * 2..];

    tile.extend(
        rs.iter()
            .zip(gs)
            .zip(bs)
            .flat_map(|((&r, &g), &b)| [r, g, b]),
    );
    
    tile
}

首先，为什么它应该快：

• 通道立即被分割。编译器将能够知道这些是独立前进的。
• extend可以在更新 的长度之前写入所有项目Vec。
• 切片上的迭代器往往可以很好地优化。

确实如此！它的运行时间为5.2 毫秒。

我也尝试过 Rayon，但它不是最干净的。这个任务不太适合并行化，因为它已经太快了。我也找不到现有的 Rayon API 来使这项工作在安全代码中进行，因此我不安全地编辑了Vec.

use rayon::prelude::*;

pub fn par_zip_spare_arr<'a>(channels: &[u8], tile: &'a mut Vec<u8>) -> &'a [u8] {
    tile.clear();
    tile.reserve(channels.len());
    let spare = tile.spare_capacity_mut();

    let len_each = channels.len() / 3;
    let rs = &channels[..len_each];
    let gs = &channels[len_each..len_each * 2];
    let bs = &channels[len_each * 2..];

    (spare.par_chunks_exact_mut(3), rs, gs, bs)
        .into_par_iter()
        .for_each(|(spare, &r, &g, &b)| {
            let spare: &mut [MaybeUninit<u8>; 3] = spare.try_into().unwrap();
            *spare = [r, g, b].map(MaybeUninit::new);
        });

    // SAFETY: all the elements were written in the previous loop
    unsafe { tile.set_len(len_each * 3) }
    tile
}

它的运行时间为1.9ms，但这取决于你有多少个核心，当然，如果你已经用其他东西占用了核心，这将无济于事。

我做的最后一件事是更改构建配置。如果您添加-C target-cpu=native编译器标志，而不是生成适用于具有您的 CPU 架构的任何 CPU 的程序集，它将生成适合您的特定 CPU 的最佳程序集。在本例中，我的 CPU 支持 AVX2，这对于较新的 x86_64 CPU 来说很常见。我重新运行基准测试并发现我的第一个函数 ，zip_extend现在运行时间为1.5ms。其余的没有受到明显影响。

您可以在Performance Book中尝试很多其他的事情。