一个很酷的 Rust 优化故事

大家好，我是螃蟹哥。

这是一篇关于 Rust 优化的文章，文章有点长，很硬核，但读完，相信会有收获。

在 Quickwit，我们正在为大数据构建最具成本效益的搜索引擎。我们的整个搜索引擎^[1]是用 rust 开发的，搜索的核心是由一个名为 tantivy^[2] 的库提供的。

人们经常问我为什么 tantivy^[3] 在基准测试中的^[4]表现优于 Lucene^[5]。这是一个复杂的问题。许多人认为其中之一是 rust 比 Java 更快。但真相要复杂得多。

Rust 的真正好处是它为程序员提供了更多的特性，供你玩耍。相比之下，JVM 是工程的宝藏，但优化是一种令人沮丧的体验。虽然 JIT 做了一项出色的一刀切的工作，但它使程序员很难理解和控制生成的代码。

在这篇博文中，我们将介绍一段特定的性能关键代码，这些代码多年来经历了一些引人入胜的变化。

这段代码是我最喜欢的展示 rustc/LLVM 功能的片段之一。

今天我将成为你的兔子向导。请跟我进我的兔子洞。

##01 问题设置

tantivy 的基本原理在于接受用户查询并返回一个迭代器，遍历与该查询匹配的文档 ID（无符号 32 位整数）。

你可能知道，整个过程依赖于倒排索引^[6]。让我们考虑一个用户搜索hello world默认情况下，tantivy 会将其解释为布尔查询hello AND world。倒排索引方便地为每个术语存储包含该术语的文档 ID 列表。我们称之为 posting list（发布列表）。发布列表以及 tantivy 生成的所有文档 ID 迭代器都已排序。

倒排索引将提供两个迭代器，每个迭代器分别动态解码与 hello 和 world 关联的发布列表。

Tantivy 的工作包括有效地组合这两个迭代器以在它们的交集上创建迭代器。由于所有涉及的迭代器都可以方便地排序，tantivy 可以在有限的内存量和线性的时间内完成此操作。

在实践中，tantivy 不依赖于 rust 的Iteratortrait，而是依赖于如下所示的 trait DocSet^[7]：

/// Represents an iterable set of sorted doc ids.
pub trait DocSet: Send {
  /// Goes to the next element.
  ///
  /// The DocId of the next element is returned.
  /// In other words we should always have :
  /// ```ignore
  /// let doc = docset.advance();
  /// assert_eq!(doc, docset.doc());
  /// ```
  ///
  /// If we reached the end of the DocSet, TERMINATED
  /// should be returned.
  ///
  /// Calling `.advance()` on a terminated DocSet
  /// should be supported, and TERMINATED should
  /// be returned.
  fn advance(&mut self) -> DocId;

  /// Returns the current document
  /// Right after creating a new DocSet, the docset
  /// points to the first document.
  ///
  /// If the DocSet is empty, .doc() should return
  ///`TERMINATED`.
  fn doc(&self) -> DocId;

  /// Advances the DocSet forward until reaching the
  /// target, or going to the lowest DocId greater than
  /// the target.
  ///
  /// If the `DocSet` is already at a `doc == target`,
  /// then it is not advanced, and `self.doc()` is
  /// simply returned.
  ///
  /// Calling seek with a target lower than the current
  /// `document` is illegal and may panic.
  ///
  /// If the end of the DocSet is reached, TERMINATED
  /// is returned.
  ///
  /// Calling `.seek(target)` on a terminated DocSet is
  /// legal. Implementation of DocSet should support it.
  ///
  /// Calling `seek(TERMINATED)` is also legal and is
  /// the normal way to consume a DocSet.
  fn seek(&mut self, target: DocId) -> DocId {
    // This is just a default implementation
    // that `DocSet` implementation should override.
    let mut doc = self.doc();
    debug_assert!(doc <= target);
    while doc < target {
      doc = self.advance();
    }
    doc
  }
}

这里的 seek 操作大大简化了我们 DocSet 交集的实现。

impl<Lhs, Rhs> DocSet for IntersectionDocSet<Lhs, Rhs>
where
    Lhs: DocSet,
    Rhs: DocSet {

  fn advance(&mut self) -> DocId {
    let mut candidate = self.left.advance();
    loop {
      let right_doc = self.right.seek(candidate);
      candidate = self.left.seek(right_doc);
      if candidate == right_doc {
        return candidate;
      }
    }
  }

  /* ... */
}

很明显，使得 seek(...) 尽可能快地获得最佳交集性能至关重要。

事实上，profiling 告诉我们调用 seek(...) 运行交集查询时占 73.6% 的时间。

如果交集中的两个词具有非常不同的频率（例如， The Mandalorian），则查找可能一次跳过数百万个文档。这一事实暗示了一个重要的优化机会。我们能否无痛地扫描这数百万份文件的情况下找到我们的目标？

Tantivy 的发布列表被压缩成 128 个文档块。我们通过内存映射访问所有这些数据。在搜索时，我们使用非常有效的 SIMD 位打包方案动态解压缩这些块。

为了避免在不需要时访问和解压缩这些块，tantivy 单独保留每个块的最后一个 doc id 的列表。我们称这个列表为 skip 列表。

在 seek 期间，tantivy 首先使用此列表来识别可能包含目标文档的单个块。它只是最后一个文档超过目标的第一个块。然后 Tantivy 对其进行解压缩并在解压缩的块中搜索目标。

在这最后一步中，我们必须将目标文档搜索到一个由 128 个排序的 DocId 组成的块中。我们知道这个块中的最后一个文档大于或等于我们的目标，我们想要找到大于或等于我们的目标的第一个元素的索引。

我们的问题归结为实现以下功能，即今天要优化的功能。

/// Returns the position of the first document that is
/// greater or equal to the target in the sorted array
/// `arr`.
fn search_first_greater_or_equal(
    arr: &[u32],
    needle: u32) -> usize;

这就是我今天要讨论的这个功能的实现。

01 第一个实现：标准二分查找

当术语的频率有些平衡并且倾向于出现在同一文档中时，在同一块中 seek 和 find 多个文档的情况并不少见。

有了这个设置，每次都从块的开头继续寻找（seeking）是很愚蠢的。如果在第 62 位找到最后一个文档，我们可以将搜索限制为&block[63..128]。

对于频率较低的，候选者很可能会在我们最后一个位置后不久出现。这种情况相当频繁。一个项是小步，而另一个是大步。

出于这个原因，该算法首先运行指数搜索以限制我们的搜索范围。然后我们将对数组的剩余部分执行简单的二分搜索。

总体而言，该功能^[8]如下所示：

/// Search the first index containing an element greater or equal to the needle.
///
/// # Assumption
///
/// The array is assumed non empty.
/// The target is assumed greater or equal to the first element.
/// The target is assumed smaller or equal to the last element.
fn search_within_block(arr: &[u32], needle: u32) -> usize {
    let (start, end) =
        exponential_search(needle, arr);
    start + arr[start..end]
        .binary_search(&needle)
        .unwrap_or_else(|e| e),
}

fn exponential_search(arr: &[u32], needle: u32) -> Range<usize> {
    let end = arr.len();
    let mut begin = 0;
    for &pivot in &[1, 3, 7, 15, 31, 63] {
        if pivot >= end {
            break;
        }
        if arr[pivot] > needle {
            return begin..pivot;
        }
        begin = pivot;
    }
    begin..end
}

02 Rust 1.25 中的性能回归

请注意，尽管听起来很普通，但 tantivy 只是依赖于 rust 标准库 binary_search 实现。

那时，它刚刚被 Alkis Evlogimenos 改进^[9] 为无分支。这个新实现非常适合我的用例，在这个用例中，我的针（needle）的分布本质上是均匀的。

如果您不熟悉无分支算法背后的思想，这里有一个简而言之的关键思想。现代 CPU 在长管道中处理指令。为了避免花费愚蠢的时间等待分支条件的结果，CPU 对分支的结果下注，并围绕这个假设组织他们的工作。

分支预测器负责根据历史数据预测这个结果。当一个分支被错误预测时，CPU 需要处理它的所有工作并重新组织它的管道。这是我们想避免的非常昂贵的事件。

虽然现代分支预测器因其预测准确性而受到称赞，但在我们的阵列中的任何地方搜索一根针是 50/50 的赌注。

出于这个原因，扭曲我们的算法以删除其所有分支是有用的。最常用的工具之一是用条件移动替换我们的分支。条件移动是相当于此代码段的 CPU 指令。

fn conditional_mov(
  cond: bool,
  val_if_true: usize,
  val_if_false: usize) ->  usize {
  if cond {
    val_if_true
  } else {
    val_if_false
  }
}

如果你在 Godbolt^[10] 上检查这个函数，它会生成以下代码：

mov     rax, rsi
test    edi, edi
cmove   rax, rdx
ret

cmove 是个神奇的指令。

现在让我们看看它在标准库二分查找中的作用。

在 rustc 1.24 中，以下代码段：

pub fn binary_search(sorted_arr: &[u32], needle: u32) ->  usize {
  sorted_arr
    .binary_search(&needle)
    .unwrap_or_else(|e| e)
}

编译成：

  push    rbp
  mov     rbp, rsp
  xor     eax, eax
  test    rsi, rsi
  je      .LBB0_5
  cmp     rsi, 1
  je      .LBB0_4
  xor     r8d, r8d
.LBB0_3:
  mov     rcx, rsi
  shr     rcx
  lea     rax, [rcx + r8]
  cmp     dword ptr [rdi + 4*rax], edx
  cmova   rax, r8
  sub     rsi, rcx
  mov     r8, rax
  cmp     rsi, 1
  ja      .LBB0_3
.LBB0_4:
  cmp     dword ptr [rdi + 4*rax], edx
  adc     rax, 0
.LBB0_5:
  pop     rbp
  ret

循环体发生在 .LBB0_3 并且它包含的唯一分支是在这里检查我们是否已经完成了足够的迭代。这个唯一的分支是可以预测的，所以一切都很好。

不幸的是，rustc 1.55 生成的代码非常不同。

  xor     eax, eax
  test    rsi, rsi
  je      .LBB0_8
  mov     rcx, rsi
  jmp     .LBB0_2
.LBB0_5:
  inc     rsi
  mov     rax, rsi
  mov     rsi, rcx
.LBB0_6:
  sub     rsi, rax
  jbe     .LBB0_8
.LBB0_2:
  shr     rsi
  add     rsi, rax
  cmp     dword ptr [rdi + 4*rsi], edx
  jb      .LBB0_5
  je      .LBB0_7
  mov     rcx, rsi
  jmp     .LBB0_6
.LBB0_7:
  mov     rax, rsi
.LBB0_8:
  ret

从 rustc 1.25 开始，二分查找不再是无分支的！

我观察了 tantivy 基准的回归并在此处报告了该 issue #57935^[11]，但它与 #53823^[12] 重复。直到今天，这个问题仍然没有解决。

03 CMOV 或者 not CMOV

这是一个简短的旁白。我不知道为什么 LLVM 在这种情况下不发出 CMOV 指令。

发出 CMOV 是否是一个好主意是一个非常棘手的难题，通常取决于提供给程序的数据。

即使对于二分查找，如果已知指针几乎总是在 0 位置，则分支实现在任何 CPU 上都将优于无分支实现。

从历史上看，CMOV 名声不佳，部分原因是 Pentium 4 在这条指令上的表现非常糟糕。让我们看看 Linus Torvald^[13] 在 2007 年对 CMOV 的评价：

CMOV（以及更一般地说，任何“谓词指令”）在严重无序的 CPU 上通常是一个坏主意。它并不总是很糟糕，但实际上它很少很好，而且（像往常一样）在 P4 上它可能真的很糟糕。

在 P4 上，我认为 cmov 基本上需要 10 个周期。

但是即使忽略通常的 P4 “我讨厌不完全正常的事情”，cmov 实际上也不是一个好主意。你可以随时替换它：

j<negated condition> forward
mov ..., %reg
 forward:

并且假设分支完全是可预测的（并且所有分支的 95+% 是可预测的），对于 CPU 来说，分支实际上会好很多。

为什么？因为分支是可以预测的，当它们被预测时，它们基本上就会消失。它们也在许多层面消失了。不仅仅是分支本身，而且就代码的关键路径而言，分支的条件消失了：CPU 仍然需要计算并检查它，但从性能角度来看它“不再存在” ，因为它不持有任何东西了。

Pentium 4 在 CMOV 上确实很烂，但现代编译器通常不再针对它进行优化。

事实上，现在 LLVM 往往更喜欢 CMOV 而不是分支。

例如，这是一个令人惊讶的编译结果：

pub fn work_twice_or_take_a_bet(cond: bool, val: usize) ->  usize {
  if cond {
    val * 73 + 9
  } else {
    val * 17 + 3
  }
}

编译成：

lea     rax, [rsi + 8*rsi]
lea     rcx, [rsi + 8*rax]
add     rcx, 9
mov     rax, rsi
shl     rax, 4
add     rax, rsi
add     rax, 3
test    edi, edi
cmovne  rax, rcx
ret

在这里，LLVM 更喜欢计算两个分支和 CMOV 结果，而不是发出一个分支！

当然，这只是因为 LLVM 观察到两个分支内的工作量足够轻，可以证明这种权衡是合理的……但这仍然很令人惊讶，不是吗？

线性 SIMD 搜索

这种性能回归非常烦人，而且 rustc 编译器似乎不太可能很快修复它。

我决定用一个始终快速执行并且对编译器的变迁不那么敏感的实现来交换我的指数 + 二进制搜索。

鉴于数组很小，我决定实现一个简单的 SIMD 无分支线性搜索。

使线性搜索无分支的技巧是将搜索问题重新表述为计算有多少元素小于针的问题。

这个想法转化为以下标量代码：

fn branchless_linear_search(arr: &[u32; 128], needle: u32) -> usize {
  arr
    .iter()
    .cloned()
    .map(|el| {
      if el < needle { 1 } else { 0 }
    })
    .sum()
}

不幸的是，SSE 的实现又长又臭：

use std::arch::x86_64::__m128i as DataType;
use std::arch::x86_64::_mm_add_epi32 as op_add;
use std::arch::x86_64::_mm_cmplt_epi32 as op_lt;
use std::arch::x86_64::_mm_load_si128 as op_load;
use std::arch::x86_64::_mm_set1_epi32 as set1;
use std::arch::x86_64::_mm_setzero_si128 as set0;
use std::arch::x86_64::_mm_sub_epi32 as op_sub;
use std::arch::x86_64::{_mm_cvtsi128_si32, _mm_shuffle_epi32};

const MASK1: i32 = 78;
const MASK2: i32 = 177;

/// Performs an exhaustive linear search over the
///
/// There is no early exit here. We simply count the
/// number of elements that are `< needle`.
unsafe fn linear_search_sse2_128(
    arr: &[u32; 128],
    needle: u32) -> usize {
    let ptr = arr as *const DataType;
    let vkey = set1(needle as i32);
    let mut cnt = set0();
    // We work over 4 `__m128i` at a time.
    // A single `__m128i` actual contains 4 `u32`.
    for i in 0..8 {
        let cmp1 = op_lt(op_load(ptr.offset(i * 4)), vkey);
        let cmp2 = op_lt(op_load(ptr.offset(i * 4 + 1)), vkey);
        let cmp3 = op_lt(op_load(ptr.offset(i * 4 + 2)), vkey);
        let cmp4 = op_lt(op_load(ptr.offset(i * 4 + 3)), vkey);
        let sum = op_add(op_add(cmp1, cmp2), op_add(cmp3, cmp4));
        cnt = op_sub(cnt, sum);
    }
    cnt = op_add(cnt, _mm_shuffle_epi32(cnt, MASK1));
    cnt = op_add(cnt, _mm_shuffle_epi32(cnt, MASK2));
    _mm_cvtsi128_si32(cnt) as usize
}

该实现为我带来了在 1.25 之前享受的性能。

04 二分查找的反击

在阅读了 dirtyhandscoding 的博客文章后^[14]，我决定再试一次二分查找。

这里的要点是简化代码库。不仅 SIMD 的使用难以阅读和维护，而且 SIMD 指令集也是特定于体系结构的，这意味着我还必须维护算法的标量版本。我获得的性能提升只是蛋糕上的一颗樱桃。

这次我会一直搜索整个 block。该块的长度为 128 个元素，这意味着我们应该能够在 7 次比较中确定结果。这样，我们就可以不惜一切代价进行这 7 个比较。

当然，我们也希望我们生成的代码尽可能高效且完全无分支。

这是我能想出的最惯用的代码来实现我们的目标。

pub fn branchless_binary_search(arr: &[u32; 128], needle: u32) -> usize {
    let mut start = 0;
    let mut len = arr.len();
    while len > 1 {
        len /= 2;
        if arr[start + len - 1] < needle {
            start += len;
        }
    }
    start
}

我没想到用这么简单的一段代码就能达到我想要的效果。

这里的关键部分是我们没有传递切片 ( &[u32]) 作为参数，而是传递了一个数组 ( &[u32; 128])。这样，LLVM 在编译时就知道我们的块正好有 128 个文档 ID。

生成的程序汇编代码如下所示：

; Idiom to set eax to 0.
xor     eax, eax

; Iteration 1 (len=64)
cmp     dword ptr [rdi + 252], esi
setb    al
shl     rax, 6

; Iteration 2 (len=32)
lea     rcx, [rax + 32]
cmp     dword ptr [rdi + 4*rax + 124], esi
cmovae  rcx, rax

; Iteration 3 (len=16)
lea     rax, [rcx + 16]
cmp     dword ptr [rdi + 4*rcx + 60], esi
cmovae  rax, rcx

; Iteration 4 (len=8)
lea     rcx, [rax + 8]
cmp     dword ptr [rdi + 4*rax + 28], esi
cmovae  rcx, rax

; Iteration 5 (len=4)
lea     rdx, [rcx + 4]
cmp     dword ptr [rdi + 4*rcx + 12], esi
cmovae  rdx, rcx

; Iteration 6 (len=2)
lea     rax, [rdx + 2]
cmp     dword ptr [rdi + 4*rdx + 4], esi
cmovae  rax, rdx

; Iteration 7
cmp     dword ptr [rdi + 4*rax], esi
adc     rax, 0
ret

LLVM 确实超越了自己。想象一下那里发生了什么：LLVM 设法

• 展开 while 循环
• 证明start + len - 1总是小于 128 并删除所有边界检查
• 在不那么琐碎的地方发出 CMOV。
• 为第一个和最后一个迭代案例找到了一个非平凡的优化。

让我们一起分解这个汇编代码：

在这个函数中，

• rax: 是我们的返回值，start在 Rust 代码中。使所有这些有点混乱的一件事是我们还通过eax 和 al 写入和读取它。它占 64 位。虽然 rax 是一个 64 位寄存器，eax 和 al分别是指其最低 32 位和它的最低 8 位。
• esi 是我们的针参数。
• rdi 是我们数组中第一个元素的地址。

让我们一步一步地完成汇编。

EAX 归零

xor     eax, eax

这可能看起来很奇怪，但这只是设置 eax 为 0 的最常见方式。为什么？机器码只需要 2 个字节。此外，现代 CPU 实际上不会在这里计算 XOR。它们只是将这条指令视为 “wink”，告诉它们我们想要一个值为 0 的寄存器。

第一次迭代

// Iteration 1 (len=64)
cmp     dword ptr [rdi + 252], esi
setb    al
shl     rax, 6

第一次迭代有点奇怪。显然，LLVM 发现了一些特定于第一次迭代的优化。但是第一次迭代有什么特别之处呢？此时，start 等于 0，其终值只能是 0 或 64。

rdi + 252 只是访问数组第 63 个元素的指针运算。( 252 = 63 * size_of::<u32>())

如果之前比较较低，setb al 指令设置 al 为 1。

shl 是位移指令。

Rust 代码如下所示：

let cmp = arr[63].cmp(&needle);
start =  //<  well actually we only set the lowest 8 bits.
    if cmp == Ordering::Lower {
        1
    } else {
        0
    }
start <<= 6;

因为64 = 2 ^ 6，我们确实以预期的 start = 64 或 start = 0 结束。。。

迭代 2-6

迭代 2-6 是类似的，不包含任何诡秘。例如，让我们看一下迭代 2：

lea     rcx, [rax + 32]
cmp     dword ptr [rdi + 4*rax + 124], esi
cmovae  rcx, rax

如果比较将我们带到右半部分，我们使用 rcx 存储我们想要的 start 值。因此等效的 Rust 代码是：

// lea     rcx, [rax + 32]
let start_if_right_of_pivot: usize = start + 32;
// cmp     dword ptr [rdi + 4*rdx + 124], esi
let pivot = arr[start + 31];
let pivot_needle_cmp: std::cmp::Ordering = pivot.cmp(target);
// cmovb  rax, rcx
let start =
    if pivot_needle_cmp_order == Ordering::Lower {
        start_if_right_of_pivot
    } else {
        start
    };

哦，等一下！我只是撒了谎。 我们拥有的代码不是 cmovb rax, rcx，它是 cmovae rcx, rax。

出于某种原因，LLVM 最终轮换了寄存器的角色。注意 rax 和的角色 rcx 是如何一次又一次地交换的。它在性能方面没有任何好处，所以我们忽略它。

最后一次迭代

最后一次迭代似乎也很特别。这里有趣的是 shift 值只是1，所以我们可以直接添加我们比较的输出到 start。

等效的 rust 代码如下所示：

// cmp     dword ptr [rdi + 4*rax], esi
let cmp = arr[start].cmp(&target)
// adc     rax, 0
if cmp == std::cmp::Lower {
    start += 1;
}

05 基准测试

CPU 模拟器和微基准测试

CPU 模拟器估计此代码需要 9.55 个周期^[15]。很优秀！请记住，我们正在搜索 128 个整数块中的值。

相比之下，我们的 SIMD 线性搜索实现估计为 26.29 个周期。

我能想到的最好的 C++ 实现是在 Clang 上超过 12 个周期，在 GCC 上超过 40 个周期。

让我们看看模拟器告诉我们的内容是否真的转化为现实世界。

Tantivy 有几个微基准测试来测量在发布列表上调用 seek 的速度。这些基准测试将大致对应于每次调用 Advance 时跳过的平均文档数的倒数作为参数。跳跃越短，值越高。

Before (无分支 SSE2 线性查找)

bench_skip_next_p01  58,585 ns/iter (+/- 796)
bench_skip_next_p1   160,872 ns/iter (+/- 5,164)
bench_skip_next_p10  615,229 ns/iter (+/- 25,108)
bench_skip_next_p90  1,120,509 ns/iter (+/- 22,271)

After (无分支内联二分查找)

bench_skip_next_p01  44,785 ns/iter (+/- 1,054)
bench_skip_next_p1   178,507 ns/iter (+/- 1,588)
bench_skip_next_p10  512,942 ns/iter (+/- 11,090)
bench_skip_next_p90  733,268 ns/iter (+/- 12,529)

这一点也不差！在这个基准测试中，Seek 现在大约快 11%。

真实场景的基准测试

Tantivy 带有一个搜索引擎基准测试^[16]，可以比较不同的搜索引擎实现。它尝试将不同类型的现实世界查询与不同的数据集进行比较。

以下是 AND 查询的输出示例。Tantivy 在交集查询上平均快 10%。

由于 block-WAND 算法，使用 top-K 收集器的 OR 查询也受益于这种优化。

06 Conclusion

所以 LLVM 是完美的，看汇编代码是徒劳的？在这一点上，你们中的一些人可能认为这里的教训是 LLVM 在编译惯用代码方面做得如此出色，以至于查看汇编、手动展开内容等只是浪费时间。

我被告知过无数次，但我不同意。

为了得到这个版本的 rust 代码，我不得不反复琢磨并确切地知道我想要什么。例如，这是我的第一个实现。

pub fn branchless_binary_search(arr: &[u32; 128], target: u32) -> usize {
    let mut range = 0..arr.len();
    while range.len() > 1 {
        let mid = range.start + range.len() / 2;
        range = if arr[mid - 1] < target {
            mid..range.end
        } else {
            range.start..mid
        };
    }
    range.start
}

它使用范围而不是(start, len)独立操作。LLVM 无法应用我们在此代码中讨论的优化的任何优化。

我会进一步优化。本博客中的实现实际上不是 tantivy 中提供的代码版本。虽然 rustc 今天在编译这个函数方面做得很好，但我不相信未来的rustc版本也会这样做。

例如，一年前，Rustc 1.41 未能删除边界检查。为了获得一致的编译结果，tantivy 实际上调用了不安全的 get_unchecked。我有信心它是安全的吗？我会在晚上睡觉吗……我会的。rustc 1.55 生成的代码提供了它安全的正式证明。

原文链接：https://quickwit.io/blog/search-a-sorted-block/

参考资料