硬件加速 SIMD 指令：如何在 BBQ 中实现极速向量比较优化

本文介绍了利用 SIMD（单指令多数据）指令优化 Elasticsearch 和 Apache Lucene 中的向量比较的新方法——BBQ（更好的二进制量化）。通过压缩存储向量数据，BBQ 提升了效率（高达32倍），并保持高排名质量。详细探讨了 SIMD 指令在 BBQ 中的应用，特别是 x64 上的 AVX 和 ARM 上的 NEON 指令。性能测试显示，SIMD 优化的点积实现比传统标量方法快 8 到 30 倍，显著提升了整体性能。BBQ 已在 Elasticsearch 8.16 中作为技术预览版发布，并可在 Serverless 中使用。

我们如何利用硬件加速 SIMD（单指令多数据）指令优化 BBQ 中的向量比较

为了让 Elasticsearch 和 Apache Lucene 成为存储和搜索向量数据的最佳选择，我们引入了一种新的方法——BBQ（更好的二进制量化）。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 支持多种向量相似度指标，如点积、余弦和欧几里得距离，但我们将重点放在点积上，因为其他指标可以从点积中推导出来。尽管我们可以在 Elasticsearch 中编写自定义的原生向量比较器，但我们更倾向于尽可能在 Java 环境中操作，以便 Lucene 也能更容易地受益。

比较查询和存储的向量

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

BBQ 采用非对称量化；查询向量被量化为 int4，而存储的向量则被进一步压缩到仅一个比特值。由于点积是组件值乘积的总和，我们可以立即看到，只有存储向量为 1 的查询组件值才能对结果产生积极贡献。

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

bbq in elastic

逻辑上，点积简化为以下公式，其中 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 = ... // 存储向量比特
byte[] q = ... // 查询比特
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);
// 尾部处理（如果有的话）
ret += subRet << q_part; // 查询部分编号

这里我们明确针对 AVX，每次循环操作 256 位。首先在 vq 和 vd 之间执行逻辑与操作，然后在结果上执行比特计数，最后将其添加到 sum 累加器中。虽然我们关注比特计数，但我们将向量中的字节解释为长整型，因为这样简化了加法操作，确保不会溢出累加器。最后一步是将累加器向量的各个通道水平地简化为标量结果，然后根据查询部分编号进行移位。

在我的 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 x 8 位 = 256 位）添加到循环计数器并继续。

这里为了简化，我们没有展示实际的实现，但我们实际上会展开 4 个查询部分，在每次循环迭代中同时执行它们，这样可以减少数据向量的加载。我们还为每个部分使用独立的累加器，最后进行简化。

对于那些向量维度不适合每次 256 位步进的情况，或者在 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。