如何让Rust编译器生成AVX2指令？

Question

我正在尝试编写一个简单的例程来使用 AVX2 指令。

作为示例，给出同一函数的以下两个版本：

fn mul1(xs: &[i32], ys: &[i32]) -> Vec<i32> {
    // Assuming xs.len() == ys.len()
    xs.iter().zip(ys.iter()).map(|(x, y)| x * y).collect()
}

fn mul2(xs: &[i32], ys: &[i32]) -> Vec<i32> {
    const CHUNK_SIZE: usize = 8;
    // Assuming xs.len() == ys.len() and xs.len() % CHUNK_SIZE == 0
    let chunks = xs.chunks_exact(CHUNK_SIZE).zip(ys.chunks_exact(CHUNK_SIZE));
    let result_chunks = chunks.map(|(c1, c2)| c1.into_iter().zip(c2.into_iter()).map(|(x, y)| x * y));
    result_chunks.flatten().collect()
}

以及相关的汇编，我并没有真正看到第二个中进行了广泛的乘法。

我做错了什么？

Answer 1

我相信原因可以归结为

flatten()

。

第一个版本始终可以准确且精确地估计它将产生的物品（

xs

和

ys

的长度中较小的一个）。用技术术语来说，它意味着

TrustedLen

。

collect()

中的专业化会检测到这一点，并提前分配所有空间，而无需检查每个项目是否需要增长集合。因此，LLVM 可以对函数进行向量化，因为没有增长检查。

但是，

flatten()

不能是

TrustedLen

，因为每个项目的尺寸可能不同。从技术上讲，在这种情况下可以，因为

chunks_exact()

产生相等的块，但它没有实现。因此

collect()

需要插入增长检查，因此 LLVM 无法向量化该函数。