利用SIMD指令加速向量搜索

Chris Hegarty

多年来，Java 平台上运行的代码一直受益于自动向量化——HotSpot C2 编译器中的superword优化，将多个标量操作打包到 SIMD（单指令多数据）向量指令中。这很好，但是这些类型的优化有些脆弱，具有天然的复杂性限制，并且受到 Java 平台规范的约束（例如，浮点运算的严格排序）。这并不是说这样的优化不再有价值，只是在某些情况下，明确代码的形状可以获得明显更好的性能。Lucene 中支持向量搜索的低级底层操作就是这样一种情况。

本文将介绍 Lucene 向量搜索中使用的底层基本操作，它们如何在运行时可靠地编译为 SIMD 指令（例如x64 上的AVX指令和 AArch64 上的 NEON 指令），以及这对性能有何影响。

底层基本操作

Lucene 向量搜索实现的核心在于查找两个向量之间的相似性时使用的三个基本操作：点积、平方和余弦距离。这些操作都有浮点和二进制变体。为了简洁起见，我们只看其中一个基本操作——点积。界面很简单，如下所示：

到目前为止，这些原始操作已得到标量实现的支持，并通过使用现有的已知技术（例如手动展开循环）来考虑性能。如果您正在编写 Java 代码，这已经无法优化了 — 其余的我们留给 HotSpot 即时编译器来尽其所能（例如，自动向量化）。这是一个简化的标量点积实现，已经去除了展开，（真正的实现可以在这里看到）：

最近的变化是JDK现在提供了一种API，可以在运行时可靠地编译为SIMD指令的计算。这就是OpenJDK的Panama项目的向量API。当然，实际在运行时生成的指令取决于底层平台支持的内容（例如AVX2或AVX 512），但该API的结构考虑了这一点。这里再次给出一个简化版本的点积代码，但这次使用了Panama向量API：

虽然代码可能会变得冗长，但是如果它符合惯用语法并且易于理解如何映射到运行时的硬件，那么它会更易于维护。

您可以看到代码有点冗长，但它符合惯用语法并且很容易推理出它在运行时如何映射到硬件，因为您可以在代码中看到向量运算。VectorSpecies包含元素类型（在我们的例子中为浮点数）和形状（向量的位大小）。首选的species是平台上最大位大小的species。首先，有一个循环遍历输入，每次乘以SPECIES::length个元素并累加。然后，将累加器向量缩减。最后，一个标量循环处理任何剩余的“尾部”元素。

当我们在支持 AVX 512 的 CPU 上运行此代码时，我们看到 HotSpot C2 编译器发出 AVX 512 指令。高级矢量扩展 (AVX) 已广泛使用，例如基于英特尔 Ice Lake 微架构的 CPU 和基于此类架构的云计算实例（例如GCP或AWS）。AVX 512 指令一次跨过点积计算 16 个值；512 位大小 / 每个值 32 位 = 每次循环迭代 16 个值。当在支持 AVX2 的 CPU 上运行时，同一代码的一次循环迭代每次迭代都会跨过 8 个值。同样，NEON（128 位）每次循环迭代将跨过 4 个值。

要看到这一点，我们需要查看生成的代码。让乐趣开始！

原生代码

以下是dotProduct的 HotSpot C2 编译器在支持 AVX 512 的 Rocket Lake 上运行时的反汇编。

下面的代码片段包含主循环体，其中rcx和rdx寄存器保存指向第一个和第二个浮点数组的地址。

• 首先，我们看到一条vmovdqu32指令，该指令将 512 位打包双字值从内存位置加载到zmm0寄存器中，即从偏移量到第一个 float[] 的 16 个值。
• 其次，我们看到一条vmulps指令，它将先前加载到zmm0中的打包单精度浮点值与内存位置中相应的打包双字值相乘 - 这是第二个 float[] 中偏移量的 16 个值，并存储生成的浮点值- zmm0中的点值。
• 第三，我们看到vaddps将zmm0中的 16 个打包单精度浮点值与zmm4相加，并将打包单精度浮点结果存储在zmm4中- zmm4是我们的循环累加器。
• 最后，有一个小的计算来递增并检查循环计数器。

不要太担心反汇编的具体细节——提供它们是为了让人们更多地了解“幕后”发生的事情，而不是需要我们对其深入了解。事实上，上面的内容有些简化，因为实际发生的是 C2 展开循环，一次跨过 4 次迭代。

我们在每次迭代中使用更多的寄存器和指令。这很好！更重要的是，我们的 Lucene代码手动也展开了循环，又是 4 倍（嗯...展开量很大）。

那么，性能提升了吗？

为了评估重写此类低级操作的影响，我们求助于 JMH，这是对此类 Java 代码执行微基准测试的普遍接受的方法。在这里，我们使用了 Robert Muir 非常好用且方便的一组基准, 让我们能够快速比较代码变体之前和之后的区别。

请记住，SIMD 提供数据并行性，因此我们处理的数据越多，潜在的好处就越大。在我们的例子中，这与向量的维度大小直接相关——我们期望看到更大的维度大小带来的更大好处。让我们看看在支持 AVX 512 的 CPU （比如，Intel Core i9-11900F @ 2.50GHz）上运行时，具有 1024 维向量的点积的浮点变体：

该基准测试每微秒的操作次数，因此越大越好。在这里，我们看到新点积的执行速度比旧点积快大约八倍。我们看到不同的低级基元操作（无论是浮点数还是二进制）都有类似的性能提升：

我们看到所有原始操作变体以及各种从小到大尺寸的显着改进（此处未显示，但可以在 Lucene PR中看到）。这一切都很棒，但这与更高级别的工作负载有何关系？

从宏观上看

微基准测试对于了解低级基元操作的执行情况非常重要，但这对宏观层面有何影响？为此，我们可以查看 Elasticsearch ®的向量搜索基准，即 SO Vector 和 Dense Vector。两个基准测试都显示出显着的改进，但让我们看看 SO Vector，因为它更有趣，因为它比 Dense Vector 具有更高的向量维度。

SO Vector 基准测试使用 200 万个 768 维向量和带过滤的 kNN 来测试向量搜索性能。这些向量基于从 StackOverflow 帖子转储中导出的数据集。Elasticsearch 使用单节点集群在 GCP 上运行，在具有 8 个 vCPU、16GB RAM 和 1x300GiB SSD 磁盘的自定义 n2 实例上运行。该节点有一个固定到 Ice Lake 的 CPU 平台。

基准测试中存在很多预先存在的差异，但总的来说，我们看到了积极的改进：

• 索引吞吐量提高了约 30%。

• 合并时间减少了约 40%。

• 查询延迟显着改善。

但Panama vector API 不是正在孵化吗？

JDK Vector API 是在 Panama 项目中开发的，目前已经孵化了一段时间。孵化状态并不反映其质量，而更多的是依赖 OpenJDK 中其他令人兴奋的工作（即值类型）的结果。Lucene在这方面走在了“前沿”，并且有一种新颖的方法来利用JDK中的非最终API——通过构建包含 JDK 版本特定 API 的“apijar”——感谢 Uwe。这是一种务实的做法，我们不会掉以轻心。与生活中的大多数事情一样，这是一个权衡。

只有当潜在的好处超过维护成本时，我们才考虑采用非最终的JDK API。Lucene 仍然保留了这些低级底层操作的标量实现版本。实现的版本可在启动时选择（请参见Lucene更改日志）。更快的Panama实现可在JDK 20和即将推出的JDK 21上使用，而对于旧的JDK或其他情况下不可用的情况，我们会回退到标量实现。再次强调，仅支持最新的JDK版本是在平衡潜在好处和维护成本时做出的实用选择。

总结

现在，我们可以使用 Panama vector API 编写可靠地利用硬件加速 SIMD 指令的 Java 代码。在 Lucene 9.7.0 中，我们添加了更快地实现矢量搜索所使用的低级底层操作的功能。Elasticsearch 8.9.0默认情况下开箱即用地启用了这种更快的实现，因此您无需执行任何操作即可获得改进的性能优势。我们在矢量搜索基准测试中看到了显着的性能改进，并完全期望这能够转化为用户工作负载。

SIMD 指令并不新鲜，并且已经存在很长时间了。与往常一样，您需要进行自己的性能基准测试，以了解这对您的环境产生的影响。AVX 512 很酷，但它可能会遭受可怕的“降频”的困扰。总体而言，我们看到这一变化带来了全面的积极改进。

最后，我想感谢 Lucene PMC 成员Robert Muir和Uwe Schindler，感谢他们愉快而富有成效的合作，促成了这一改进，没有他们，这一切就不会发生。