作者来自 Elastic Lorenzo Dematte, Florian Bernd 及 Chris Hegarty深入剖析四项优化级联展开、批量预取、维度轴展开以及结构性重构这些优化通过顺应 CPU 的工作方式而非与之对抗将 Elasticsearch simdvec 的向量吞吐量提升至 2 倍。通过这个面向 Search AI 的自主实践学习课程亲自体验向量搜索。你现在就可以开始免费云试用或者在你的本地机器上试用 Elastic。Elasticsearch simdvec相比串行代码可实现高达 50 倍更快的向量距离计算。要达到这一目标需要依次解决四个硬件瓶颈而每解决一个瓶颈都会暴露出下一个瓶颈。本文将介绍级联展开、批量预取、维度轴展开以及一次结构性重构 —— 事实证明这项重构带来了最大的收益。所有这些优化共同实现了高达 2 倍的吞吐量提升。本文是《我们如何构建 Elasticsearch simdvec使向量搜索成为全球最快之一》的配套文章。Elasticsearch 中的每一次向量搜索查询无论是分层可导航小世界HNSW遍历、倒排文件IVF扫描还是重排序阶段本质上都归结为同一个问题在每次查询过程中对数百万个向量执行距离计算。Elasticsearch simdvec是 Elasticsearch 中所有向量距离计算背后的引擎。从指令执行角度来看向量距离计算本身是一项非常简单的操作例如点积仅仅是加法和乘法的组合。然而要让这些计算变得足够快 —— 甚至快到极致 —— 则需要深入理解现代 CPU 的工作机制、不同 指令集架构Instruction Set Architectures - ISAs所提供的能力以及它们之间的共性与特性差异。在本文中我们将深入探讨 simdvec 如何针对内存访问进行优化。针对 x86 和 ARM 手工调优的单指令多数据Single Instruction, Multiple Data- SIMD内核仅需几个 CPU 周期即可完成向量距离计算对于 SIMD 内核而言其性能瓶颈往往并非能够执行多少运算而是在每个 CPU 周期内能够获取并消耗多少数据。例如一个 1024 维的 float32嵌入向量在计算一次点积时需要执行 1024 次乘加运算。AVX-512 处理器可以在每个 512 位寄存器中容纳 16 个浮点数并且每个周期可发出两条融合乘加FMA指令。在持续运行的情况下其吞吐量相当于仅需 32 个周期即可完成一次点积计算对于 4GHz 的 CPU 而言每个向量仅需 8 纳秒。搜索 100 万个候选向量意味着需要执行该内核 100 万次并将大约 4GB 的向量数据流经 CPU。从计算能力角度来看芯片只需总计 8 毫秒即可完成这些数学运算真正的问题在于如何及时提供这 4GB 的数据这是一个不可能完全实现的目标但我们究竟能接近到什么程度呢本文其余部分将介绍我们如何尽可能多地让向量穿过这块芯片。这是一场走钢丝般的优化过程每一个让我们更接近峰值吞吐量的步骤都会对下一步施加更严格的约束。以下是我们实施的四项优化按照应用顺序排列级联展开Cascade Unrolling—— 使用独立累加器链最大化 FMA 端口利用率批处理与预取Batch Processing and Prefetching—— 通过提前预取数据来隐藏内存延迟维度轴展开Dim-Axis Unrolling—— 规避 2 的幂维度下的 L1d 缓存别名问题查询加载提升Query Load Hoisting—— 消除每个文档重复执行的查询操作端口、流水线、延迟、吞吐量……现代 CPU 如何执行 SIMD 向量运算现代 CPU 能够在每个周期发出多条操作指令这是因为其芯片内部实现了多个执行单元这些操作会通过称为端口Port的接口进行分派在 x86 架构中如此称呼而在 ARM 架构中则称为执行流水线Execution Pipeline简称流水线Pipe。端口能够并行处理不同类型的工作有些负责内存加载和存储操作有些负责整数运算还有些负责浮点数学计算。任何一个端口能够处理的操作都有两个重要属性延迟Latency和吞吐量Throughput。延迟是指单个操作从开始执行到产生结果所需的 CPU 周期数吞吐量则是指每个周期能够启动多少个此类操作。吞吐量与某种操作可用的端口数量密切相关。例如如果某个 CPU 拥有两个能够执行 FMA融合乘加运算的端口那么在理想条件下它每个周期最多可以发出两个彼此独立的新 FMA 指令即达到每周期两次 FMA 的峰值吞吐量。以 AVX-512 为例。现代 Intel CPU 上的大多数 FMA 指令通常具有约四个周期的延迟并且可以在两个支持 FMA 的端口中的任意一个上执行。从冷启动状态开始第一个结果会在四个周期后产生但当流水线被填满之后只要这些指令之间不存在依赖关系每个周期都可以启动两条新的 FMA 指令。我们的第一步优化目标就是尽可能提高端口利用率在考虑延迟因素的同时平衡各端口的负载。级联展开最大化 FMA 端口利用率延续前面的例子如果一个 FMA 指令的延迟为四个周期而吞吐量为每周期两次那么 CPU 在任意时刻大约可以同时维持八个处于飞行状态In Flight的 FMA 指令 —— 也就是已经发出但尚未完成的指令。当然这只有在存在八个彼此独立的操作时才能实现。如果将这些操作串联起来那么每个 FMA 都必须等待前一个 FMA 的结果。这样一来CPU 的执行速度将受限于延迟每四个周期执行一次 FMA而不是发挥吞吐量能力每周期执行两次 FMA。其速度可能比硬件能够提供的理论性能慢多达 8 倍。这种依赖链很容易在无意间产生。例如一个朴素的向量点积实现可能如下所示foreach (i) { acc acc x[i]*y[i] }由于只使用了一个累加器accumulator每次循环迭代都依赖于上一次迭代的结果。一种很自然的改进思路是循环展开Loop Unrolling如果我们需要同时维持 N 条处于飞行状态的指令那么就将同一条指令发出 N 次。编译器甚至为此提供了专门的指令例如#pragma unroll。循环展开在 simdvec 代码中被广泛使用以充分利用现代 CPU 内部的并行能力。然而#pragma unroll的问题在于它只是给编译器提供的一个提示Hint而不是强制指令。此外它的效果还依赖于编译选项以及编译器自身的启发式策略因此编译器可能决定不展开循环或者执行不够理想的展开方式。例如当我们检查编译器为该循环生成的汇编代码时发现循环确实被展开了但依赖链依然存在。因此如果希望获得精确控制或保持跨平台一致性仍然需要手动展开循环。然而手工展开的代码不仅难以阅读而且几乎无法维护。C 模板与元编程C 模板Template允许你编写带有占位类型或占位值的通用代码而这些占位符会在编译期间由编译器进行替换。函数模板只需编写一次编译器便会针对每一种实际使用的参数组合生成专门化版本。例如占位符可以是一个类型如float与int或者寄存器类型__m512i与uint8x16_t也可以是一个函数、一个编译期整数等等。最后一种形式是我们使用最频繁的通过一个以整数 N 为参数的模板我们能够生成 N 个并行累加器或者生成内层循环体的 N 份副本。元编程Metaprogramming本质上是“编写能够生成代码的代码”。它利用编译器在编译期间执行计算从而不会产生任何运行时开销。我们的主要工具是apply_indexedN。这是一个编译期函数在展开过程中会生成 N 条语句template int N, typename F, int I 0 static inline void apply_indexed(F f) { if constexpr (I N) { f(std::integral_constantint, I{}); apply_indexedN, F, I 1(std::forwardF(f)); } }if constexpr是一种编译期分支机制它使apply_indexed成为一个编译期递归函数编译器会在编译期间解析constexpr条件并实例化模板的下一次迭代。整个过程完全由编译器处理不会生成任何运行时分支指令。我们利用apply_indexedN实现了级联展开Cascade Unrollingapply_indexedN([](auto I) { fma(acc[I], x[i I*stride], y[i I*stride]); });我们将循环展开为级联结构首先使用 N4 条彼此独立的累加器链对于剩余长度减半的尾部数据则降为 2 条累加器链最后对于剩余的标量尾部部分则使用 1 条累加器链。与#pragma unroll相比这种方式在各种内核和 CPU 上带来了约 11%13% 的性能提升完整细节和全部测试数据可参见上文链接的 PR。编译器能够提供的优化终究有限而通过 C 模板实例化实现的泛型编程则让我们能够在不同内核和不同 ISA 架构之间保持极高的执行效率同时维持代码的紧凑性和可维护性。批处理与预取隐藏内存延迟循环展开解决的是单次向量计算内部的指令级并行Instruction-Level ParallelismILP问题但它并没有利用批量处理带来的优势。Elasticsearch 并不是将一个向量与一个查询进行匹配评分仅 HNSW 遍历过程中每个查询就需要对数百个邻居节点进行评分。批量评分即一个查询同时匹配多个文档既带来了新的问题也为我们提供了解决这些问题的工具。需要评分的向量通常分散存储在内存中的不同位置从而形成不规则的访问模式。这种模式很难被 CPU 缓存以及硬件预取器准确预测。其结果是CPU 所需的数据更有可能不存在于高速的 L1d 缓存中而必须从更远的层级加载即发生缓存未命中Cache Miss。一次数据访问的典型代价可能从L1d 缓存命中时的大约 5 个周期一直增加到数据必须从主内存RAM读取时的 200 多个周期这种差距可达到数十倍。因此当 SIMD 计算本身只需几十个周期即可完成时真正的瓶颈往往已经不再是计算而是等待数据从内存系统到达 CPU。如果这些 load 端口即使被最大化利用但仍然在等待数据那么计算端口也会处于空闲状态。通过级联展开所饱和的 FMA 吞吐量也会被白白浪费。把数据从内存取到 L1d cache 是一个非常耗时的过程尤其是在需要遍历整个内存层级、最终访问 RAM 的情况下。幸运的是既然我们知道接下来会对多个向量进行打分就可以提前“预热” CPU cache把下一个或多个向量的数据预取进来从而有效降低或隐藏内存访问延迟。批处理同样也能帮助降低指令延迟其原理与级联展开类似N 条独立的向量流意味着 CPU 可以维护 N 条独立的累加器链从而交错执行计算并隐藏前面提到的 FMA 流水线延迟。这也是我们在 bulk scoring 中引入batch的原因最早从 int7 开始后来扩展到所有数据类型。其机制与级联展开本质相同但作用在“向量之间”而不是“向量内部”。我们不是一次处理一个向量而是同时处理 N 个向量在处理这些向量的同时还会为下一批 N 个向量预加载prefetch数据。这种方式理论上应该同时改善预取效率和指令延迟而在许多情况下也确实如此。例如在 int7 场景中我们相对于未展开的 bulk 实现立即获得了 20%50% 的性能提升完整细节和 JMH benchmark 可见对应 PR。但当我们尝试将其扩展到所有 bulk 函数时我们发现它引入了新的问题与约束。好事过头为什么突发预取会溢出 line-fill buffer只有当预取的 cache line 在内层循环需要它之前已经进入 L1d 时预取才是有效的。我们最初的实现是在 batch 边界一次性发出所有预取请求形成一个约 28100 条软件预取指令的“突发”。处理器每个核心的line-fill bufferLFB是有限的它直接决定了 CPU 能同时跟踪多少个未完成的 cache miss 请求。例如 Sapphire Rapids 的 LFB 只有 16 个条目。这种规模的突发请求会溢出 LFB超出的预取请求会被静默丢弃。结果就是我们以为已经在路上的 cache line 实际并没有被加载内层循环仍然不得不等待 cache miss。解决方法是把预取请求分散到整个内层循环中。在 batch 边界只发出少量 “头部突发”head burst用于覆盖最先会被消费的 cache line然后在后续迭代中逐步均匀地继续预取每次迭代只负责拉取下一部分数据。这样做虽然总预取次数不变但 LFB 的峰值占用下降了一个数量级。cache line 能在被需求加载前约一个 outer iteration 就到达从而隐藏 L2 到 L1 的延迟同时 L2 stream prefetcher 也更适应这种稳定的访问节奏而不是边界突发模式。这种 “head spread” 策略最早在 int8 中落地随后扩展到 int7 以及其他 kernel最终带来了最高约 30% 的性能提升。自己踩到自己的脚为什么在 2 的幂维度上做 batching 会拖慢性能理论上正确数量的预取应该能带来很高的吞吐量通过隐藏大部分内存延迟来提升性能。确实在大多数情况下它是有效的 —— 尤其是在向量数据以稀疏、随机方式访问时效果很好。但当我们尝试在 batch 中并行处理连续的、顺序排列的文档尤其是在维度为 2 的幂时性能却突然出现了断崖式下降。为什么 N 路组相联缓存会在 2 的幂维度上产生冲突CPU cache 由 cacheline和 cacheset组成。cacheline是内存层级中传输数据的基本单位在 ARM 和 x86 上通常都是 64 字节。每个 cache line 会映射到唯一一个 cacheset而每个 set 可以容纳固定数量的 cache line。这种结构被称为N 路组相联缓存N-way associative cache。可以把它类比成一个哈希表每个 bucket 有 N 个槽位。多个内存地址可能映射到同一个 bucketset但当槽位填满后再插入新元素就必须驱逐已有条目。我们来看一个具体例子。Sapphire Rapids 的 L1d cache 大小为 48KiB12 路组相联。按 64 字节 cache line 计算共有 768 条 cache line被组织成 64 个 set。一个 cache line 属于哪个 set由地址的 set index 决定也就是“哈希键”该 index 由地址的 [11:6] 位决定——换句话说就是(address / 64) % 64。假设我们有一个 1024 维的 float32 嵌入向量并且它在内存中是连续存储的。每个向量占用dims * sizeof(float32) 4096 字节也就是正好 64 条 cache line。因此相邻两个向量之间的步长stride就是 4096 字节。由于 4096 恰好对应 64 个 set index 的完整循环范围set index 会发生完美“回绕”每个向量中的第 i 条 cache line都会映射到完全相同的 cache set。常见的 2 的幂向量尺寸会产生一个在字节层面同样是 2 的幂的 stride这个 stride 会均匀整除 64 个 cache set从而以一种“病态”的方式与 cache 结构发生冲突。处理 N 个文档的 batch 时这个问题会被进一步放大由于它们全部落在相同的 L1d cache set 中它们会相互冲突导致 cache 抖动cache thrashing。通过降低 batch 并行度修复 cache 别名问题我们在为 bf16 数据类型实现 kernel 时首次深入分析了这一现象。当我们尝试不同实现以及不同batches参数进行基准测试时结果验证了我们的假设在 2 的幂维度下连续向量的 stride 会映射到相同的 cache set而多个 load stream 交错访问时会导致频繁 eviction。因此我们立刻采用了一个简单的修复方案在具有顺序访问特征的函数*_bulk中将batches1以避免 L1d cache set aliasing。不过我们也清楚这只是一个权宜之计Band-Aid。bulk 处理和更高程度的并行仍然是有价值的例如它们可以帮助隐藏延迟因此我们希望在不重新引入 cache 冲突的前提下保留它们。沿另一条轴展开在不产生 aliasing 的情况下填满 FMA 流水线batches是在文档维度across documents上做并行化但这并不是唯一的并行维度。我们还可以沿着向量的维度轴dimension axis进行展开。也就是说我们不再并行处理多个向量而是并行处理同一对向量中多个独立的分块chunks。因此我们不要重复自己通过提升查询负载以消除冗余端口压力来自unroll_dim在 欧几里得距离 上的微小收益表明在 bulk loop 中仍然隐藏着另一个瓶颈。我们发现的是 bulk function template 中隐藏的一种结构性低效现有的 bulk scorer 会为每个文档调用 single-pair scorer在每个 outer step 中重新加载查询元素 N (4) 次。对于某些函数我们也不必要地多次重复与查询元素相关的操作。例如 int8 欧几里得距离 kernelsqri8在每个 outer step 中会调用 vpmovsxbw sign-extension 指令四次。将查询加载和操作从 per-document loop 中提升出去使查询元素的 L1D bandwidth 降低了 4x对于sqri8它在每个 outer step 中移除了四个vpmovsxbw符号扩展指令中的三个。需要注意的是吞吐量取决于端口可用性vpmovsxbw只能在单个端口上执行Sapphire Rapids 上的 port 5因此每个 step 发出四个副本会完全压满该端口仅符号扩展本身就是瓶颈。不要重复自己通过提升查询负载以消除冗余端口压力来自unroll_dim在 欧几里得距离 上的微小收益表明在 bulk loop 中仍然隐藏着另一个瓶颈。我们发现 bulk function template 中隐藏的一种结构性低效现有的 bulk scorer 会为每个 document 调用 single-pair scorer在每个 outer step 中重新加载 query elements N (4) 次。对于某些 functions我们也不必要地多次重复与 query elements 相关的操作。例如 int8 欧几里得距离 kernelsqri8在每个 outer step 中会调用 sign-extension instruction 四次。将 query load 和 operation 从 per-document loop 中提升出去使 query elements 的 L1D bandwidth 降低了 4x对于sqri8它在每个 outer step 中移除了四个vpmovsxbwsign-extension instruction 中的三个。需要注意的是 throughput 取决于 port availabilityvpmovsxbw只能在单个 port 上执行Sapphire Rapids 上的 port 5因此每个 step 发出四个 copies 会完全压满该 port仅 sign extensions 本身就是瓶颈。即使没有 query-specific operation 可以提升这个改变仍然很重要。对于doti8_mm512_dpbusd_epi32 在 Sapphire Rapids 上会在两个 portsport 0 和 port 5执行latency 为 5 cycles因此我们需要大约 10 个 independent operations 在 in flight 中才能达到 peak throughput。通过提升 query loadinner loop 变成了依赖于每个 batch element 的单一 accumulator chain。冗余工作被移除后unroll_dim2可以通过在 dim axis 上增加 independent chains 来填充 latency window。这种 structural refactoring 带来了不错的 speedupdot product 提升 19–22%Euclidean distance 提升 44–51%所有细节和完整数字可以在链接的 PR 中找到。为了完美落点后退一步并不是所有 optimization 都能存活。经过引入unroll_dim的所有工作之后benchmarking 显示它并没有统一收益对于某些 kernels 和 access patterns额外的 register pressure 和 code complexity 没有带来任何可测提升。我们本可以保留它并设为unroll_dim1功能上等价 no-op但无用的 scaffolding 是技术债务会让下一次修改更难理解。因此我们在它不划算的地方将其移除保持代码干净。在钢丝上行走时有时正确的动作是后退一步。关键要点优化 vector search 低层 memory access本文中的每一次优化都遵循同样的模式解决一个 bottleneck又暴露下一个。cascade unrolling saturate 了 FMA ports从而暴露 memory latency。batching 和 prefetch 隐藏了 latency又暴露 L1d set aliasing。dim-axis unrolling 绕开 aliasing又暴露冗余 query work。而移除这些冗余最终让整个 pipeline 得以“呼吸”。在 simdvec kernels 的性能优化中没有单一优化可以让系统变快。每一次改进都只是改变 bottleneck而不是消除它并且每一步都可能暴露新的未预期约束。在这个层面上memory abstractions 是一种幻觉性能取决于你是否理解 CPU 实际在做什么而不仅仅是模型告诉你的东西。唯一的前进方式是测量、理解、再平衡。优化收益级联展开对比 #pragma unroll11-13%批处理 / 预取int720-50%头部 分散预取最高 30%维度轴展开点积35-65%查询负载提升点积19-22%查询负载提升欧几里得距离44-51%这是关于 Elasticsearch simdvec 的一系列深度解析中的第一篇。下一篇我们将研究代数重写如何让我们完全绕过 CPU 限制。这篇内容有多大帮助原文Elasticsearch simdvec: 2x vector throughput with SIMD - Elasticsearch Labs