使用SIMD加速std::copy_if

# 使用SIMD加速copy_if 来源: https://loonatick-src.github.io/posts/vectorized-copy-if-analysis/ ## 引言 我有一块Zen 4 CPU，带着一堆AVX512特性标志。于是我想——不妨试试用它实现点什么，哪怕是重复造轮子。我设定了以下目标： 1. 实现一个即便用多面体循环模型也无法被优化编译器自动向量化的算法。 2. 系统分析其性能，并回答以下问题：1) 它是否已经达到最快？2) 如果不是，原因是什么？我们该如何改进？ 3. 从简单入手，先让它跑起来。这意味着像 map/transform、reduce、adjacent_difference 这类极为简单的算法排除在外，因为它们很容易被自动向量化（参见 Compiler Explorer 链接）。即便是二维模板计算也不行——看看这个（Compiler Explorer 链接）。于是，我选择了 `std::copy_if`。实现一个SIMD版本并不难，但理解其性能表现却比我预想的复杂得多。 我清楚需要用到哪些工具： 1. Google benchmark（链接）——编写微基准测试。 2. likwid-bench（链接）——确定机器上的性能上限。 3. llvm-mca（链接）——在Zen 4模型上模拟内核。 4. perf-stat（链接）——通过计数事件进行深入性能分析。 根据 cppreference（链接），`std::copy_if` 是一个极其简单的算法： ``` template<InputIt, OutputIt, UnaryPred> OutputIt copy_if(InputIt first, InputIt last, OutputIt d_first, UnaryPred pred) { for (; first != last; ++first) if (pred(*first)) { *d_first = *first; ++d_first; } return d_first; } ``` 生成的代码也非常简洁（Compiler Explorer 链接）。然而，它不容易向量化，因为存在一个循环携带依赖：第 `i+1` 次迭代中 `d_first` 的值依赖于第 `i` 次迭代中 `pred(*first)` 的结果。 在开始向量化之前，先测量一下基线。以下是我们可以进行性能测量的维度： 1. 输入大小（记作 `n`） 2. 谓词函数的选择 3. 输入分布 4. 输入熵 问题规模（1）很容易遍历；改变 `n` 会导致与内存子系统（缓存、硬件预取器、DRAM等）的不同交互。谓词和分布共同决定了输出的密度/稀疏度。例如，谓词 `[](auto x){ return x > 0; }` 配合范围 `(-1000, 1000)` 内的均匀分布，预计会有50%的输入值被复制。熵（entropy）与分布并非完全正交，但值得单独提出来。或许我也得想个更好的名字。它决定了输入的可预测性，因为所有流水线CPU都有分支预测逻辑。例如，如果CPU前端（FE）遇到条件跳转指令，它不会等待操作数就绪，而是推测性地跳转到一个目标地址。推测错误会导致巨大的惩罚，需要完全清空流水线并重新开始执行。上面同样的谓词和分布组合，可能会使得大多数分支预测器难以保持低分支误测率，从而严重影响吞吐量。 为简洁起见，我们固定谓词（`x > 0`）和分布（`(-1000, 1000)` 内的均匀分布），仅遍历问题规模。这里使用的性能分析方法可以很好地推广到针对其他维度的输入进行实现调优。 使用 likwid-bench 中的 `copy` 和 `copy_avx512` 基准测试达到的速度（MB/s） **图 1.** likwid-bench 中 `copy` 和 `copy_avx512` 基准测试达到的速度（MB/s） 可通过以下命令复现： ``` $ for size in 16kB 64kB 256kB 1MB 4MB 16MB 64MB 256MB 1GB 4GB; do bw=$(likwid-pin -c 1 likwid-bench -t copy_avx512 -w S0:${size}:1 2>/dev/null | grep "MByte/s" | awk '{print $NF}'); echo "$size $bw"; done $ for size in 16kB 64kB 256kB 1MB 4MB 16MB 64MB 256MB 1GB 4GB; do bw=$(likwid-pin -c 1 likwid-bench -t copy -w S0:${size}:1 2>/dev/null | grep "MByte/s" | awk '{print $NF}'); echo "$size $bw"; done ``` ## 首次SIMD尝试 循环体包含三个部分： 1. 从 `&input[i]` 加载数据 2. 计算谓词以得到 `bool` 值 3. 根据上一步的结果，有条件地将加载的值存储到目标位置，并更新输出计数器/指针 1和2在大多数SIMD实现中都很直接。设 `N` 为SIMD寄存器的宽度。例如，在AVX-512中加载32位值时，`N = 512 / 32 = 16`。 1. 从 `&input[i]` 加载到SIMD寄存器 —— `_mm512_loadu_epi32` 等。 2. 对SIMD寄存器计算谓词，得到SIMD掩码值。 - 对于我们的谓词（`> 0`），使用：`const auto zero = _mm512_setzero_epi32(); return _mm512_cmpgt_epi32_mask(a, zero);` 3. 将对应掩码位为1的SIMD数据通道连续存储。 - 恰好有一条指令可以实现此功能：`_mm512_mask_compressstoreu_epi32`¹（链接）。 - 目标指针可以通过对第二步得到的掩码值进行population count来更新。 当然，还需要另一个循环来处理剩余少于SIMD宽度的元素。 将所有内容整合在一起，并为了乐趣进行一次循环展开，我们得到以下代码： ``` namespace ck { template<typename Arg> struct positive { bool operator()(const Arg& arg) const& { return arg > 0; } bool operator()(const Arg&& arg) const& { return arg > 0; } }; template<> struct positive<__m512i> { __mmask16 operator()(const __m512i& a) const& { const auto zero = _mm512_setzero_epi32(); return _mm512_cmpgt_epi32_mask(a, zero); } __mmask16 operator()(const __m512i&& a) const& { const auto zero = _mm512_setzero_epi32(); return _mm512_cmpgt_epi32_mask(a, zero); } }; template<> ```