如何使用 OpenMP simd 对此循环进行矢量化？

Question

我有以下代码，它接受输入 anx并填充输出向量y，我想使用 OpenMPsimd指令对其进行向量化：

for (size_t qi = 0; qi < nq; ++qi){
    const auto ii = face_indices[qi*3], jj = face_indices[qi*3+1], kk = face_indices[qi*3+2];
    y[ii] += x[kk] * fpl[ii] * fq[qi] * fpr[kk] - x[jj] * fpl[ii] * fq[qi] * fpr[jj];
    y[kk] += x[ii] * fpl[kk] * fq[qi] * fpr[ii] - x[jj] * fpl[kk] * fq[qi] * fpr[jj]; 
    y[jj] -= x[ii] * fpl[jj] * fq[qi] * fpr[ii] + x[kk] * fpl[jj] * fq[qi] * fpr[kk]; 
}

除了 之外y，所有输入数组 ( fq, fpl, fpr, face_indices) 都是const/ 只读的。就上下文而言，x、y、fpl和 的fpr大小均为np，而fq大小为nq。

在实践中，nq大约是数十万（有时是数百万），并且由于face_indices有大小3*nq，所以我想使用 OpenMP 的simd指令对代码进行矢量化。作为参考，该代码正在实现从特定类型的拉普拉斯算子导出的结构化矩阵向量乘法代码。

虽然我得到了一个通过添加 #pragma 来编译的版本，例如：

#pragma omp simd 
for (size_t qi = 0; qi < nq; ++qi){
    ...
}

问题是生成的矢量化是错误的：它与没有编译指示的编译版本不匹配。这是一个godbolt 链接，给出了该行为的基本示例。此外，在更大的测试用例上，加速相对较小（仅约 2-3 倍）。

所以我有两个问题：

1. 如何修复代码以使 simd 版本与非 simd 版本匹配？
2. 我可以做任何明显的优化来使 simd 版本更快吗？

当行为与非 simd 版本不匹配时，我感到很惊讶，因为循环是完全隔离的/没有向后依赖性：rhs 上的输入都是const，输出y是唯一被修改的东西。此外，可以安全地假设ii、jj、 和kk每次迭代都是不同的。

---

编辑：

我现在意识到，如果 simd-length 为 $k$，对同一位置进行多次赋值确实不安全y——也就是说，不足以确保值ii、jj、 和kk都是不同的，看起来就像我需要确保它们对于每个 $k$ 展开/矢量化迭代都是不同的face_indices。

因此，例如，除非指定，否则使用以下输入face_indices是不安全的：simdlen(1)

face_indices = { 0,1,4,0,2,5,0,3,7,1,3,8,2,3,10,4,7,8,5,7,10 }

相反，以下重新排序适用于simdlen(2)：

face_indices = { 0,1,4,5,7,10,0,2,5,1,3,8,2,3,10,4,7,8,0,3,7 }

因为每组连续的 6 个索引都是不同的

Answer 1

如何修复代码以使 simd 版本与非 simd 版本匹配？

AFAIK隐含的约束omp simd在文档中没有明确说明，因此它取决于目标平台（目标体系结构、编译器、操作系统）。上一篇文章证实了这一点。话虽这么说，一个普遍的假设是SIMD 块的迭代必须完全独立。这意味着不仅所有ii、jj、 和kk在一个块中必须单独不同，而且对于给定的迭代， ii、jj和kk也不能相同。我强烈建议您检查这是否属实，因为这是最可能的解释。只要这个假设被打破，循环就无法使用大多数 OpenMP 实现进行矢量化（否则无论如何都不会高效）。

我可以做任何明显的优化来使 simd 版本更快吗？

主要问题是这段代码实际上并不适合 SIMD。

首先，face_indices[qi*3]由于.​ *3这种访问的效率通常明显低于连续访问。不过，编译器有时可以根据这种特定情况的洗牌生成不太可怕的汇编代码（因为步幅很小并且在编译时已知）。

此外，ii、jj和kk是运行时定义的索引，不能保证是均匀的或线性的。这意味着类似的访问x[kk]也不连续。在这种情况下，编译器不能在这种情况下使用 shuffle 指令。在 SSE (x86) 中，它们只能执行缓慢的标量加载。在 AVX (x86) 中，有一个为此用例设计的收集指令。然而，它通常无法在主流 x86 处理器上有效实现（它比使用标量加载填充 SIMD 寄存器更快，但处理器仍然在内部执行独立的标量加载）。在 Neon (ARM) 上，据我所知没有收集指令。在SVE（ARM）上，有一个收集指令，但迄今为止很少有处理器实现SVE。由于您的代码肯定主要受到内存访问模式的限制，因此我不认为使用 SIMD 指令会更快。

此外，在提供的 godbolt 链接中，指令集未指定。这意味着默认使用 SSE2（因为目标架构是 x86-64）。如今， SSE2 已经相当过时了。它已被 AVX/AVX2 以及最近的 AVX-512 超级种子（仅受最新的 x86-64 处理器或服务器处理器支持 - >=IceLake-Client、>=Skylake-Server 和 >=Zen4）。-march=native您可以使用（不可移植）或 等标志-mavx2。请注意，Clang 在 AVX/AVX2 中不使用收集指令，但在 AVX-512 中使用。

最重要的是，您可以对fpl[ii] * fq[qi]、 、fpl[kk] * fq[qi]和 等运算进行因式分解，fpl[jj] * fq[qi]每个运算都会计算两次。您可能认为编译器应该能够自动执行此操作，但浮点数学不具有关联性，因此它们无法执行此操作*。因此，编译器需要tmp = x[kk] * fpl[ii]先计算，然后计算tmp * fq[qi]。这是次优的，因为x[kk] * fpl[ii]不会重新计算（同样的情况也适用于其他子表达式）。

总而言之，除非您可以修改代码以使ii、jj、kk均匀或线性（至少其中之一），否则该代码无法在主流平台上有效矢量化。更改 so 的布局face_indices以使访问更加 SIMD 友好应该会有所帮助，但我不希望（仅）由此带来很大的性能改进。

* 从技术上讲，OpenMP SIMD 指令足以让目标实现忽略浮点关联性规则，因为无论如何这是循环矢量化所必需的，但我不希望编译器在此用例中这样做。您可以使用编译器标志-ffast-math来避免此问题，假设此标志在您的情况下是安全的（例如，处理 NaN 或 Inf 值时很危险）。请参阅这篇相关文章和这篇文章。

补充笔记

请注意，在 GPU 上，此代码可以变得更加高效（假设迭代实际上是完全独立的）。事实上，一些 GPU（例如 Nvidia 的 GPU）可以快速收集一种称为共享内存的特殊内存（通常每个流式多处理器只有几十 KiB）。不过，只有在没有太多存储体冲突的情况下，内存访问才会很快。