Elasticsearch：使用硬件加速的 SIMD 指令实现超快 BBQ

作者：来自 Elastic Chris Hegarty

我们如何使用硬件加速 SIMD（Single Instruction Multiple Data - 单指令多数据）指令优化 BBQ 中的向量比较。

随着我们继续致力于让 Elasticsearch 和 Apache Lucene 成为存储和搜索向量数据的最佳场所，我们在 BBQ（Better Binary Quantization - 更好的二进制量化）中引入了一种新方法。BBQ 是一个巨大的进步，通过压缩存储的向量带来了显著的（32 倍）效率，同时保持了高排名质量，并同时提供了超快的性能。你可以在 BBQ 博客中阅读有关 BBQ 如何将 float32 量化为单个位向量进行存储、在索引速度（量化时间减少 20-30 倍）和查询速度（查询速度提高 2-5 倍）方面优于乘积量化（Product Quantization）等传统方法的信息，并查看 BBQ 对各种数据集实现的出色准确率和召回率。

由于 BBQ 博客中已经介绍了高级性能，因此我们将在此更深入地了解 BBQ 如何实现如此出色的性能。特别是，我们将研究如何通过硬件加速的 SIMD（单指令多数据）指令优化向量比较。这些 SIMD 指令执行数据级并行，因此一条指令一次可对多个向量组件进行操作。我们将看到 Elasticsearch 和 Lucene 如何针对特定的低级 SIMD 指令（如 x64 上的 AVX 的 VPOPCNTQ 和 ARM 上的 NEON 的向量指令）来加速向量比较。

为什么我们如此关心向量比较？

向量比较决定了向量数据库中的执行时间，通常是成本最高的因素。我们在配置文件跟踪中看到了这一点，无论是使用 float32、int8、int4 还是其他量化级别。这并不奇怪，因为向量数据库必须大量比较向量！无论是在索引过程中，例如构建 HNSW 图，还是在搜索过程中，因为图或分区被导航以找到最近的邻居。向量比较中的低级性能改进对整体高级性能确实有影响。

Elasticsearch 和 Lucene 支持许多向量相似性指标，如点积（dot product）、余弦（cosine）和欧几里得（Euclidean），但是我们将只关注点积，因为其他指标可以从他的值里面得出。尽管我们有能力在 Elasticsearch 中编写自定义的本机（原生）向量比较器，但我们倾向于尽可能留在 Java 领域，以便 Lucene 也能更轻松地获得好处。

比较查询和存储向量

为了使我们的距离比较函数运行速度更快，我们需要尽可能简化它的工作，以便将其转换为一组在 CPU 上高效执行的 SIMD 指令。由于我们已经将向量量化为整数值，因此我们可以将更多组件加载到单个寄存器中，同时避免更昂贵的浮点运算 - 这是一个好的开始，但我们需要更多。

BBQ 进行非对称量化；查询向量（query vector）被量化为 int4，而存储向量被高度压缩为单个位值。由于点积是组件值乘积的总和，我们可以立即看到，唯一可以对结果产生积极影响的查询组件值是存储向量为 1 的值。

进一步的观察是，如果我们以将每个组件值的各个位置位（1、2、3、4）组合在一起的方式转换查询向量，那么我们的点积就会简化为一组基本的按位运算；对每个组件进行 AND + 位数运算，然后进行移位以表示查询部分的相应位置，最后进行求和以得到最终结果。有关更详细的解释，请参阅 BBQ 博客，但下图直观地展示了查询翻译的示例。

从逻辑上讲，点积简化为以下内容，其中 d 是存储的向量，q1、q2、q3、q4 是平移后的查询向量的相应位置部分：

(bitCount(d & q1) << 0) + (bitCount(d & q2) << 1)
+ (bitCount(d & q3) << 2) + (bitCount(d & q4) << 3)

为了执行目的，我们通过增加相应的位位置，在内存中连续布局翻译后的查询部分。因此，我们现在有了翻译后的查询向量 q[]，其大小是存储向量 d[] 的四倍。

此点积的标量 Java 实现如下所示：

byte[] d = ... // stored vector bits
byte[] q = ... // query bits

for (int i = 0; i < 4; i++) {
  long subRet = 0;
  for (int j = 0; j < d.length; j++) {
    subRet += Integer.bitCount(q[i*d.length + j] & d[j] & 0xFF);
  }
  ret += subRet << i;
}

虽然语义上正确，但此实现效率不高。让我们继续看看更优化的实现是什么样子，然后我们可以比较每个实现的运行时性能。

性能从何而来？

到目前为止，我们看到的实现是一种简单的标量实现。为了加快速度，我们使用 Panama Vector API 重写了点积，以便明确针对特定的 SIMD 指令。

以下是仅针对查询部分之一的代码的简化片段 - 请记住，我们需要执行四次，每个翻译的 int4 查询部分执行一次。

var sum = LongVector.zero(LongVector.SPECIES_256);
for (int i=0; i < BYTE_SPECIES_256.loopBound(d.length); i += BYTE_SPECIES_256.length()) {
  var vq = ByteVector.fromArray(BYTE_SPECIES_256, q, i).reinterpretAsLongs();
  var vd = ByteVector.fromArray(BYTE_SPECIES_256, d, i).reinterpretAsLongs();
  sum = sum.add(vq.and(vd).lanewise(VectorOperators.BIT_COUNT));
}
long subRet = sum.reduceLanes(VectorOperators.ADD);
// tail processing, if any
ret += subRet << q_part // query part number

这里我们明确针对 AVX，每个循环迭代操作 256 位。首先在 vq 和 vd 之间执行逻辑与，然后对其结果进行位计数，最后将其添加到总和累加器。虽然我们对位数感兴趣，但我们确实将向量中的字节解释为长整型，因为这简化了加法并确保我们不会冒累加器溢出的风险。然后需要最后一步，将累加器向量的通道水平减少为标量结果，然后再按代表性查询部件号进行移位。

在我的 Intel Skylake 上拆解它，我们可以清楚地看到循环的主体。

0x00007cf9bcd57430:   movsxd r14,ecx
0x00007cf9bcd57433:   vmovdqu ymm4,YMMWORD PTR [rsi+r14*1+0x10]
0x00007cf9bcd5743a:   vmovdqu ymm5,YMMWORD PTR [rdx+r14*1+0x10]
0x00007cf9bcd57441:   vpand  ymm4,ymm4,ymm5
0x00007cf9bcd57445:   vpopcntq ymm4,ymm4
0x00007cf9bcd5744b:   vpaddq ymm3,ymm3,ymm4
0x00007cf9bcd5744f:   add    ecx,0x20
0x00007cf9bcd57452:   cmp    ecx,edi
0x00007cf9bcd57454:   jl     0x00007cf9bcd57430

rsi 和 rdx 寄存器保存待比较向量的地址，从中分别将下一个 256 位加载到 ymm4 和 ymm5 寄存器中。加载完值后，我们执行按位逻辑与操作（vpand），将结果存储在 ymm4 中。接下来执行的是 vpopcntq 指令，该指令统计设置为 1 的位数。最后，我们将 0x20（32 字节 = 256 位）加到循环计数器中并继续执行。

为了简化，这里没有展示实际操作，但我们实际上展开了 4 个查询部分，并在每次循环迭代中同时执行它们，从而减少了数据向量的加载量。此外，我们为每个部分使用独立的累加器，最后再进行汇总。

当向量维度不足以每次处理 256 位时，采用 128 位的变体；在 ARM 平台上，这些操作会编译成一系列 Neon 向量指令，包括 AND、CNT 和 UADDLP。当然，对于未对齐的数据部分，我们会以标量方式处理。不过，鉴于大多数流行模型的维度大小，这种情况在实际中并不常见。我们还在继续进行 AVX-512 实验，但到目前为止，由于常见的维度大小限制，按 512 位步进处理此数据布局并未显示出明显优势。

SIMD 能改善多少？

当我们比较标量和向量化点积实现时，我们分别看到从 384 到 1536 的一系列流行维度的吞吐量提高了 8 倍到 30 倍。

通过优化点积，我们大大提高了整体性能，使得向量比较不再是使用 BBQ 搜索和索引时最主要的因素。对于感兴趣的人，这里有一些基准和代码的链接。

总结

BBQ 是一种新技术，它带来了令人难以置信的效率和出色的性能。在这篇博客中，我们研究了如何通过硬件加速的 SIMD 指令在 BBQ 中优化向量距离比较。你可以在 BBQ 博客中阅读有关索引和搜索性能以及准确度和召回率的更多信息。BBQ 现已作为技术预览版在 Elasticsearch 8.16 和 Serverless 中发布！

除了 BBQ 等新技术外，我们还在不断改进向量数据库的低级性能。你可以在其他博客中阅读有关我们已经完成的工作的更多信息：FMA、FFM 和 SIMD。此外，随着我们不断提高 Elasticsearch 的性能，使其成为存储和搜索向量数据的最佳场所，未来还会看到更多像这样的以低级性能为重点的博客。

Elasticsearch 包含许多新功能，可帮助你为你的用例构建最佳搜索解决方案。深入了解我们的示例笔记本以了解更多信息，开始免费云试用，或立即在你的本地机器上试用 Elastic。

原文：Smokin' fast BBQ with hardware accelerated SIMD instructions - Search Labs