simdvec 深度解析：Elasticsearch 如何利用神经网络和视频编解码 CPU 指令实现向量搜索

Elasticsearch 作为向量数据库，为您提供了一种高效的方式来大规模创建、存储和搜索向量嵌入（embeddings）。其核心构建于 simdvec 之上，该引擎提供比串行代码快约 50 倍，比 Java 的 Panama Vector API 快 2-3 倍的向量搜索速度。simdvec 的速度来源于现代 CPU 中实现的单指令多数据（SIMD）指令，其中许多指令最初并非为向量搜索而设计。

现代 CPU 包含惊人的高速执行单元，但它们只加速有限且经过精心挑选的操作集，而您所需的操作往往并不存在。这里的挑战不在于找到最快的指令，而是如何重新构造计算，使其能够用硬件高效执行的指令来表达。

本文将通过四个实际案例，深入探讨这种计算重构：

• 量化到 int7，以适应神经网络的无符号*有符号乘加指令。
• 对全 int8 向量评分进行代数偏差重写，利用视频编码的绝对差求和指令。
• 将 bf16 欧几里得距离表达为三个点积，充分利用有限的 bf16 指令。
• 将二进制向量点积重构为多元素（和多向量）位操作，在可用时使用 SIMD popcount，否则使用字节混洗和查找。

编译器无法识别的位：int7 如何实现最快的整数乘加操作

在我们的 x86 字节向量核中，我们希望利用 vpdpbusd，这是现代 x86 CPU 上可用的最快整数点积指令。但只有一个问题：它最初是为神经网络推理设计的，而非向量搜索。vpdpbusd 的一个操作数必须是无符号的，另一个必须是有符号的，这与向量嵌入的常规表示方式不符。

x86 SIMD 指令集包含多个为整数字节数据设计的乘加指令：除了更通用的 maddubs（所有 AVX2 和 AVX-512 处理器均可用）之外，较新的处理器还实现了更快的 vpdpbusd（支持 Vector Neural Network Instructions [VNNI] 的处理器可用），它将 maddubs、madd 和累加（add）融合为一条指令。

对于这两条指令，一个操作数被视为无符号字节 [0, 255]，另一个被视为有符号字节 [-128, 127]。这种不对称性源于它们最初的深度学习推理目的，即 ReLU 后的激活是无符号的，而权重是有符号的。我们不能直接将这些指令用于两个有符号操作数；那样结果将是错误的。例如，如果您将一个值为 -1（位模式 0xFF）的有符号字节输入到无符号操作数中，指令会悄无声息地将其解释为 255。

那么，我们为什么要费心使用 vpdpbusd 呢？除了其固有的速度（它以与单条指令相同的吞吐量执行多个操作）之外，替代方案是使用 16 位宽度的指令：加载半个寄存器的字节数据，将每个字节扩展到 16 位，然后以 16 位宽度进行乘法运算。这将使每个寄存器和每次加载的元素数量减半；而像 vpdpbusd 这样的字节宽度乘加操作，可以在相同的内存带宽下处理 2 倍的元素。正如我们在上一篇文章中看到的，这些核的性能瓶颈通常在于数据加载速度，而不是计算速度。使用 vpdpbusd 意味着直接获得 2 倍的优势。

Elasticsearch 的优化标量量化（OSQ）支持多种位宽，包括 7 位模式。在 7 位模式下，每个量化值都在 [0, 127] 范围内。最高位始终为零，这意味着该范围内每个数字的有符号和无符号解释是相同的。这意味着我们可以使用无符号*有符号乘加指令，因为它们将在不进行任何转换的情况下产生正确的结果。

这是一个协同设计决策：通过量化到 7 位而不是 8 位，我们牺牲了一位精度，以直接访问硬件上可用的最快整数乘加指令。这个权衡非常值得：基准测试显示速度提升约 6 倍，而召回率损失可忽略不计。 （所有详细信息和 JMH 基准测试都可以在链接的 PR 中找到。）

直觉和设计选择推动了这一提升：编译器永远无法做出这个判断。对于编译器而言，数据是 int8_t，一个范围为 [-128, 127] 的有符号类型。它无法证明这些值是非负的，也无法推断 vpdpbusd 对非负数使用是安全的，因此它永远不会发出 vpdpbusd。这个决策需要关于量化方案的知识，而这些知识超出了语言层面的类型信息。

当数值不匹配时：我们如何用无符号操作处理有符号 int8

完整的 int8 向量仍然受到这个问题的影响：两个操作数都可能包含负值，而 vpdpbusd 要求一个操作数是无符号的。

Elasticsearch 支持字节格式的输入向量数据；对于这些数据集，我们无法控制输入范围，因此我们必须支持跨越整个 [-128, 127] 范围的操作数。我们最初的核使用了我们前面描述的朴素加宽路径：将字节进行符号扩展到 16 位，然后以双倍宽度进行乘法运算。这是正确的，但它将吞吐量减半。

然而，有一种方法可以让 vpdpbusd 在有符号数据上工作。关键是一个代数恒等式：

a \cdot b = (a + 128) \cdot b - 128 \sum b

将 a 的每个元素加 128 会将范围从 [-128, 127] 移到 [0, 255]，从而允许 a 作为 vpdpbusd 所期望的无符号字节操作数。实际上，这种移位可以通过对每个寄存器执行一次 0x80 的 XOR 操作来实现（翻转符号位与在二进制补码中加 128 相同）。代价是一个校正项 128 * sum(b)，必须从结果中减去。这种技术在深度学习推理领域是众所周知的：Intel 的 oneDNN 库使用相同的恒等式来处理带 vpdpbusd 的有符号激活。

校正项使用了另一条为他人的问题而存在的指令：vpsadbw 计算无符号字节之间的绝对差之和，最初是为视频编码中的运动估计算法设计的。当一个操作数是零向量时，这些绝对差就简化为字节值本身，因此 vpsadbw 实际上计算了 b 中字节的部分和。这些部分和可以累加到校正项中。一条为完全不同目的设计的指令，通过一点代数重构，竟然变得非常有用。

这正是上一篇文章中介绍的批量评分架构再次发挥作用的地方。如果我们将 b 视为我们的查询向量，那么校正项是一个只依赖于查询而不依赖于任何文档的标量。在批量路径中，我们可以在评分循环之前预计算校正项，这样每个文档向量只需要 XOR + vpdpbusd，而无需任何每文档校正开销。相比之下，单对路径必须为每次调用计算校正项，并且开销会完全抵消 vpdpbusd 的加速效果。

我们对此进行了诚实的测量：在 AMD Turin（AWS c8a.xlarge，Zen 5 核心）上运行的 JMH 基准测试显示，批量评分性能提升约 20%，但在单对评分上没有优势。 （详细信息可在链接的 PR 中找到。）因此，simdvec 使用两种路径：vpdpbusd 结合偏差技巧用于批量处理，加宽路径用于单对处理。我们以基准测试结果为指导，而非技巧的优雅程度。

三条指令胜过一条：我们如何利用点积计算欧几里得距离

这本应是最简单的一个：CPU 已经有专门的 bf16 点积指令。不幸的是，并没有所谓的“bf16 欧几里得距离指令”。

BFloat16 (bf16) 正成为向量嵌入的流行格式：它以降低的精度覆盖与 float32 相同的数值范围，将存储空间减半，这意味着每个向量的内存带宽也减半。硬件供应商已经通过专用指令（AVX-512 上的 vdpbf16ps 和 ARM Scalable Vector Extension [SVE] 上的 bfdot）来响应，这些指令直接将 bf16 点积计算到 32 位浮点累加器中。在 AVX-512 上，这意味着每个 512 位寄存器处理 32 个 bf16 元素，无需转换步骤。

乍一看，这里似乎没什么好说的：指令第一次完全符合我们的需求。实际上，bf16 点积核直接映射到这些指令。但 Elasticsearch 还需要平方欧几里得距离： |a - b|^2 。

当我们开始AVX-512 bf16 支持的工作时，最初的平方距离实现遵循了显而易见的路径：将两个 bf16 向量都转换为 float32，然后相减并平方：

__m512 diff = _mm512_sub_ps(bf16_to_f32(a), bf16_to_f32(b));
sum = _mm512_fmadd_ps(diff, diff, sum);

512 位寄存器可以容纳 32 个 bf16 值，但只能容纳 16 个 float32 值。转换后，每个后续操作的元素密度减半。而且 bf16_to_f32 转换函数使用整数 ALU 操作（将 16 位零扩展到 32 位，然后左移 16 位）。这意味着每步处理的元素数量减半（16 个）；此外，转换操作会占用执行资源，从而影响实际的数学计算（它们会争夺 ALU 端口；参见我们上一篇文章中的“端口、管道、延迟、吞吐量”）。

当查询已经以 float32 格式存储时，这是正确的路径：无论如何，核都必须以 32 位宽度工作。但当两个向量都是 bf16 时，我们可以做得更好，利用一个代数等价式。

平方欧几里得距离展开为：

|a - b|^2 = a \cdot a + b \cdot b - 2(a \cdot b)

每个项都是一个点积，这意味着每个项都直接映射到 CPU 已经加速的硬件原语（vdpbf16ps/bfdot）。新的核使用三个原生的 bf16 点积，无需转换，并以完整的 32 元素寄存器宽度操作：

sum_self  = _mm512_dpbf16_ps(sum_self, av, av);   // a·a
sum_self  = _mm512_dpbf16_ps(sum_self, bv, bv);   // b·b
sum_cross = _mm512_dpbf16_ps(sum_cross, av, bv);  // a·b
// 循环结束后：
result = reduce(sum_self) - 2.0f * reduce(sum_cross);

我们正在执行更多的算术指令，但每条指令都以全密度操作 bf16 数据。基线执行较少的算术指令，但操作宽度减半。此外，整个内循环是纯浮点运算：没有整数 ALU 操作，这导致资源争用减少。相同的恒等式在 ARM SVE 上与 svbfdot_f32 同样有效。

开销取决于 vdpbf16ps 在目标微架构上的昂贵程度，换句话说，取决于 CPU 供应商决定在这条指令上投入多少硅片。在 AMD Zen 5 上，vdpbf16ps 经过了高度优化：每条指令都是一个微操作（uop），吞吐量为每周期两条指令；每步三次调用以 3 个 uop 处理 32 个元素。相比之下，转换路径每步处理 16 个元素，指令链更长（每次转换 2 个整数操作，1 个减法，1 个融合乘加 [FMA]）。这使得三点积路径在处理 bf16 欧几里得距离时，比现有 AVX2 实现快 2.1-2.4 倍。

在 Intel Emerald Rapids 上，vdpbf16ps 是4 个 uop，吞吐量为 0.5，即每两个周期一条指令。优势较小，但点积公式仍然胜出，为我们带来了1.3-1.4 倍的加速，因为它在整个循环中保持了 bf16 密度。所有细节和 JMH 基准测试都可以在链接的 PR 中找到。

批量评分：将仅查询项提升到每向量循环之外

与 int8 评分一样，bf16 欧几里得距离也有一个仅查询项，可以为每个查询预先计算一次，而不是为每个文档计算一次。如果我们将 b 视为查询向量，那么 b·b 仅依赖于查询，而不依赖于任何文档。批量核将 b·b 从每文档循环中提升出来，并将其累加到单个共享流中，从而将内循环从每个文档的三个点积减少到两个（a·a 和 a·b）。

批量 bf16 欧几里得核比单对评分快 1.25-1.8 倍。 （像往常一样，所有细节和 JMH 基准测试都可以在链接的 PR 中找到。）部分优势归因于通常的批量优势（更好的指令级并行性和预取），但在我们的测量中，提升的 b·b 占增益的 10-20%，具体取决于向量布局。

未曾选择的道路：我们尝试过（并放弃了）的另一种 bf16 距离方法

三点积分解并非是唯一将平方距离表达为 vdpbf16ps 的方法。我们还探索了另一种方式：使用 unpacklo/unpackhi 将两个 bf16 向量交错成 [a₀, b₀, a₁, b₁, ...] 对，然后将每个对与一个常数向量 [1, −1, 1, −1, ...] 进行点积。硬件计算 aᵢ × 1 + bᵢ × (−1) = aᵢ − bᵢ，直接在 float32 中产生逐元素差值，因此无需转换。随后的 FMA 对结果进行平方并累加。

这是对同一指令的另一种创新用法：一种方法将平方距离分解为三个点积；另一种方法则重新利用点积来实现“bf16 减法”。两者都完全避免了 bf16_to_f32 转换。从理论上看，交错路径的算术运算似乎更精简，每步一个 dpbf16 加上一个 FMA，而三点积路径则需要三次 dpbf16 调用。

实际上，交错 a 和 b 所需的混洗指令增加了更多开销，并且由此产生的循环具有较少的独立累加器，限制了指令级并行性。三点积版本，凭借其清晰的加载-计算-累加结构，在我们的基准测试中明显胜出，因此我们撤销了该实验。

这并非一“显而易见”的方法和一“技巧”：两者都是相同恒等式的有效重写。我们探索了硬件所能提供的限制范围内的各种方法，然后进行了测量：直觉和对指令计数的推理只是一个起点，而不是最终答案。

点积即位逻辑：二进制向量、AND 和 SIMD popcount

在理想的指令集架构中，我们应该有专门用于二进制量化向量的点积指令。现代 CPU 并没有提供这样的指令。幸运的是，一旦向量被压缩成单个位，算术本身就开始“消失”了。

到目前为止，所有的重写都处理了完整的整数和浮点数。但 Elasticsearch 也支持二进制量化向量（Better Binary Quantization [BBQ]），其中每个维度都被压缩成一个比特。乍一看，计算一个 1 位文档向量与一个多位查询向量之间的点积似乎需要将位解包成整数，然后相乘并求和。然而，由于值仅限于 0 和 1，乘法和加法会简化为更简单的操作：两个比特的乘积是它们的 AND（仅当两个输入都为 1 时才为 1），而这些乘积的和是结果中 1 比特的数量，即人口计数（population count）。

对所有维度求和：dot(a, b) = popcount(a AND b)。这个恒等式并非新事物；几十年来，它一直是位操作算法的支柱，并广泛应用于信息检索、密码学和机器学习。其优雅之处在于它如何自然地映射到硬件。

使用 AVX-512 popcount 计算二进制向量点积

一个标量 popcnt 指令一次处理一个 64 位字，或每条指令处理 64 个二进制维度。但是，带有 VPOPCNTDQ 扩展的 AVX-512 处理器可以在单个 512 位寄存器中并行处理八个 64 位字。结合位 AND 操作，1 位文档对一个查询位平面的内循环变为两条指令：

__m512i res = _mm512_popcnt_epi64(
    _mm512_and_si512(value, query_plane)
);
acc = _mm512_add_epi64(acc, res);

这会在每次迭代中计算一个 512 维的二进制点积块。一个 4 位查询被视为四个独立的位平面，每个位位置一个。每个平面贡献不同的权重（1、2、4 或 8），因此最后一步会对 popcount 进行移位和累加，以重建完整的点积。整个内循环是少数几条指令，每遍处理数千个维度。

构建缺失的指令：在 AVX2 硬件上合成 popcount

VPOPCNTDQ 在 Intel Ice Lake 和 AMD Zen 4 及更高版本的 AVX-512 处理器上可用。但 simdvec 也需要在没有向量 popcount 指令的旧 AVX2 硬件上运行。解决方案来自 Muła, Kurz, and Lemire (2018)：使用 vpshufb 字节混洗指令作为并行 4 位查找表。其思想是使用 SIMD 寄存器作为 popcount 函数的查找表：预计算每个可能的 4 位值（0 到 15）的 popcount，并将这 16 个结果存储在一个 SIMD 寄存器中。然后，对于每个输入的字节，将其分成两个半字节，并使用 vpshufb 同时查找这两个计数值：

我们如何使用 vpshufb 作为查找指令（左），以及它如何在 32 个元素上并行操作（右）

由于 vpshufb 并行执行 32 次独立的表查找，因此两次混洗、两次掩码和一次加法即可计算 32 字节（256 个二进制维度）的 popcount，而无需专用的 popcount 指令：

const __m256i lookup = _mm256_setr_epi8(
    0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4,
    0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4
);
// 分别查找低位和高位半字节的 popcount
const __m256i lo  = _mm256_and_si256(vec, low_mask);
const __m256i hi  = _mm256_and_si256(_mm256_srli_epi16(vec, 4), low_mask);
const __m256i popcnt1 = _mm256_shuffle_epi8(lookup, lo);
const __m256i popcnt2 = _mm256_shuffle_epi8(lookup, hi);
local = _mm256_add_epi8(local, _mm256_add_epi8(popcnt1, popcnt2));

在这种情况下，指令集架构（ISA）没有提供我们所需的操作，所以我们从现有资源中构建了它。在支持 VPOPCNTDQ 的 AVX-512 上，硬件直接提供了原语；在 AVX2 上，我们可以通过字节混洗和查找表来合成它。我们的原生 AVX2 实现比 Java 标量评分快约 4 倍。升级到支持原生 VPOPCNTDQ 的 AVX-512 又在此基础上提升了 1.3-2.1 倍，因为一条指令取代了整个 popcount 序列。（像往常一样，更多详细信息和 JMH 基准测试可以在链接的 PR 中找到。）

单指令，多向量：将多个文档向量打包到一个 SIMD 寄存器中，用于低维嵌入

通过 VPOPCNTDQ，每个 popcnt 一次处理整个 512 位寄存器，即 512 个二进制维度。然而，具有 128 或 256 维度的嵌入模型非常流行；每个维度 1 位，128 维向量仅为 16 字节，是 512 位寄存器的四分之一。标准循环每次迭代加载一个文档向量，导致四分之三的寄存器未被使用。

多亏了批量处理，我们又可以做得更好：如果向量足够小，我们可以将多个文档向量压缩到同一个寄存器中。例如，对于 128 维（16 字节），四个文档向量可以放入一个 512 位寄存器。我们将查询位平面广播四次到寄存器中，然后一个 AND + POPCNT 就能一次性处理所有四个文档：

// 将查询位平面广播 4 次到 512 位寄存器中
__m512i b = _mm512_broadcast_i32x4(_mm_maskz_expandloadu_epi8(mask, query));
// 在一次 512 位加载中加载 4 个连续的文档向量
__m512i docs = _mm512_maskz_expandloadu_epi8(dataByteMask, ...);
// AND + POPCNT：一条指令对处理四个点积
__m512i res = _mm512_popcnt_epi64(_mm512_and_si512(docs, b));

通过一次加载、一次 AND 和一次 POPCNT，我们对四个文档进行了评分。对于 256 维向量（每个 32 字节），我们可以在每个寄存器中放入两个文档。

JMH 基准测试揭示了一个有趣的故事。在 Intel Granite Rapids 上，这个新内核的扩展性完美：在 128 维度下比 AVX2 快 4.4 倍（与单向量实现相比快 3.88 倍），在 256 维度下快 2.4 倍（与单向量实现相比快 2.79 倍）。但在 AMD Zen 5 上，增益要温和得多；我们测试了不同的实现，发现多向量实现在此 CPU 上相对于单向量实现没有任何优势。

这只是一个理论，但 Zen 5 内部可能使用 128 位 uop 实现 VPOPCNTDQ，以四个顺序块处理 512 位寄存器。这与我们的基准测试结果一致：如果硬件单独处理每个四分之一，那么将四个 128 位向量打包到一个寄存器中并没有帮助。无论如何，这是优化硅片空间的不同实现选择的另一个例子。

在更高维度下，两个平台都受益于 AVX-512 更宽的寄存器（我们前面提到的 1.3-2.1 倍增益）。

后续的更改通过掩码加载（expandloadu）将打包方法扩展到非 2 的幂维度（96、192），使其与 128 和 256 维度的性能持平。

Elasticsearch simdvec 经验：让 CPU 指令适应向量搜索

在本文中，我们探讨了各种优化：使用 int7 适应无符号*有符号硬件，为有符号 int8 应用代数偏差重写，将欧几里得距离表达为三个 bf16 点积，将二进制向量点积重构为位操作，以及通过字节混洗合成 popcount。

硅片空间稀缺，因此 CPU 只加速有限的一组操作。这些操作通常是为非常不同的用例设计的：神经网络推理、视频编码、密码学或图形。本文中的每一个优化都遵循相同的模式：从硬件高效执行的指令入手，然后重构计算，使其能够用这些指令来表达。