一個很酷的 Rust 優(yōu)化故事

這是一篇關(guān)于 Rust 優(yōu)化的文章，文章有點(diǎn)長，很硬核，但讀完，相信會有收獲。

在 Quickwit，我們正在為大數(shù)據(jù)構(gòu)建最具成本效益的搜索引擎。我們的整個搜索引擎^[1]是用 rust 開發(fā)的，搜索的核心是由一個名為 tantivy^[2] 的庫提供的。

人們經(jīng)常問我為什么 tantivy^[3] 在基準(zhǔn)測試中的^[4]表現(xiàn)優(yōu)于 Lucene^[5]。這是一個復(fù)雜的問題。許多人認(rèn)為其中之一是 rust 比 Java 更快。但真相要復(fù)雜得多。

Rust 的真正好處是它為程序員提供了更多的特性，供你玩耍。相比之下，JVM 是工程的寶藏，但優(yōu)化是一種令人沮喪的體驗。雖然 JIT 做了一項出色的一刀切的工作，但它使程序員很難理解和控制生成的代碼。

在這篇博文中，我們將介紹一段特定的性能關(guān)鍵代碼，這些代碼多年來經(jīng)歷了一些引人入勝的變化。

這段代碼是我最喜歡的展示 rustc/LLVM 功能的片段之一。

今天我將成為你的兔子向?qū)?。請跟我進(jìn)我的兔子洞。

問題設(shè)置

tantivy 的基本原理在于接受用戶查詢并返回一個迭代器，遍歷與該查詢匹配的文檔 ID（無符號 32 位整數(shù)）。

你可能知道，整個過程依賴于倒排索引^[6]。讓我們考慮一個用戶搜索hello world默認(rèn)情況下，tantivy 會將其解釋為布爾查詢hello AND world。倒排索引方便地為每個術(shù)語存儲包含該術(shù)語的文檔 ID 列表。我們稱之為 posting list（發(fā)布列表）。發(fā)布列表以及 tantivy 生成的所有文檔 ID 迭代器都已排序。

倒排索引將提供兩個迭代器，每個迭代器分別動態(tài)解碼與 hello 和 world 關(guān)聯(lián)的發(fā)布列表。

Tantivy 的工作包括有效地組合這兩個迭代器以在它們的交集上創(chuàng)建迭代器。由于所有涉及的迭代器都可以方便地排序，tantivy 可以在有限的內(nèi)存量和線性的時間內(nèi)完成此操作。

在實(shí)踐中，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???????doc?=?self.advance();
????}
????doc
??}
}

這里的 seek 操作大大簡化了我們 DocSet 交集的實(shí)現(xiàn)。

impl?DocSet?for?IntersectionDocSet
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(...) 盡可能快地獲得最佳交集性能至關(guān)重要。

事實(shí)上，profiling 告訴我們調(diào)用 seek(...) 運(yùn)行交集查詢時占 73.6% 的時間。

如果交集中的兩個詞具有非常不同的頻率（例如， The Mandalorian），則查找可能一次跳過數(shù)百萬個文檔。這一事實(shí)暗示了一個重要的優(yōu)化機(jī)會。我們能否無痛地掃描這數(shù)百萬份文件的情況下找到我們的目標(biāo)？

Tantivy 的發(fā)布列表被壓縮成 128 個文檔塊。我們通過內(nèi)存映射訪問所有這些數(shù)據(jù)。在搜索時，我們使用非常有效的 SIMD 位打包方案動態(tài)解壓縮這些塊。

為了避免在不需要時訪問和解壓縮這些塊，tantivy 單獨(dú)保留每個塊的最后一個 doc id 的列表。我們稱這個列表為 skip 列表。

在 seek 期間，tantivy 首先使用此列表來識別可能包含目標(biāo)文檔的單個塊。它只是最后一個文檔超過目標(biāo)的第一個塊。然后 Tantivy 對其進(jìn)行解壓縮并在解壓縮的塊中搜索目標(biāo)。

在這最后一步中，我們必須將目標(biāo)文檔搜索到一個由 128 個排序的 DocId 組成的塊中。我們知道這個塊中的最后一個文檔大于或等于我們的目標(biāo)，我們想要找到大于或等于我們的目標(biāo)的第一個元素的索引。

我們的問題歸結(jié)為實(shí)現(xiàn)以下功能，即今天要優(yōu)化的功能。

///?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;

這就是我今天要討論的這個功能的實(shí)現(xiàn)。

01 第一個實(shí)現(xiàn)：標(biāo)準(zhǔn)二分查找

當(dāng)術(shù)語的頻率有些平衡并且傾向于出現(xiàn)在同一文檔中時，在同一塊中 seek 和 find 多個文檔的情況并不少見。

有了這個設(shè)置，每次都從塊的開頭繼續(xù)尋找（seeking）是很愚蠢的。如果在第 62 位找到最后一個文檔，我們可以將搜索限制為&block[63..128]。

對于頻率較低的，候選者很可能會在我們最后一個位置后不久出現(xiàn)。這種情況相當(dāng)頻繁。一個項是小步，而另一個是大步。

出于這個原因，該算法首先運(yùn)行指數(shù)搜索以限制我們的搜索范圍。然后我們將對數(shù)組的剩余部分執(zhí)行簡單的二分搜索。

總體而言，該功能^[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 標(biāo)準(zhǔn)庫 binary_search 實(shí)現(xiàn)。

那時，它剛剛被 Alkis Evlogimenos 改進(jìn)^[9] 為無分支。這個新實(shí)現(xiàn)非常適合我的用例，在這個用例中，我的針（needle）的分布本質(zhì)上是均勻的。

如果您不熟悉無分支算法背后的思想，這里有一個簡而言之的關(guān)鍵思想?，F(xiàn)代 CPU 在長管道中處理指令。為了避免花費(fèi)愚蠢的時間等待分支條件的結(jié)果，CPU 對分支的結(jié)果下注，并圍繞這個假設(shè)組織他們的工作。

分支預(yù)測器負(fù)責(zé)根據(jù)歷史數(shù)據(jù)預(yù)測這個結(jié)果。當(dāng)一個分支被錯誤預(yù)測時，CPU 需要處理它的所有工作并重新組織它的管道。這是我們想避免的非常昂貴的事件。

雖然現(xiàn)代分支預(yù)測器因其預(yù)測準(zhǔn)確性而受到稱贊，但在我們的陣列中的任何地方搜索一根針是 50/50 的賭注。

出于這個原因，扭曲我們的算法以刪除其所有分支是有用的。最常用的工具之一是用條件移動替換我們的分支。條件移動是相當(dāng)于此代碼段的 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] 上檢查這個函數(shù)，它會生成以下代碼：

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

cmove 是個神奇的指令。

現(xiàn)在讓我們看看它在標(biāo)準(zhǔn)庫二分查找中的作用。

在 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

循環(huán)體發(fā)生在 .LBB0_3 并且它包含的唯一分支是在這里檢查我們是否已經(jīng)完成了足夠的迭代。這個唯一的分支是可以預(yù)測的，所以一切都很好。

不幸的是，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 基準(zhǔn)的回歸并在此處報告了該 issue #57935^[11]，但它與 #53823^[12] 重復(fù)。直到今天，這個問題仍然沒有解決。

03 CMOV 或者 not CMOV

這是一個簡短的旁白。我不知道為什么 LLVM 在這種情況下不發(fā)出 CMOV 指令。

發(fā)出 CMOV 是否是一個好主意是一個非常棘手的難題，通常取決于提供給程序的數(shù)據(jù)。

即使對于二分查找，如果已知指針幾乎總是在 0 位置，則分支實(shí)現(xiàn)在任何 CPU 上都將優(yōu)于無分支實(shí)現(xiàn)。

從歷史上看，CMOV 名聲不佳，部分原因是 Pentium 4 在這條指令上的表現(xiàn)非常糟糕。讓我們看看 Linus Torvald^[13] 在 2007 年對 CMOV 的評價：

CMOV（以及更一般地說，任何“謂詞指令”）在嚴(yán)重?zé)o序的 CPU 上通常是一個壞主意。它并不總是很糟糕，但實(shí)際上它很少很好，而且（像往常一樣）在 P4 上它可能真的很糟糕。

在 P4 上，我認(rèn)為 cmov 基本上需要 10 個周期。

但是即使忽略通常的 P4 “我討厭不完全正常的事情”，cmov 實(shí)際上也不是一個好主意。你可以隨時替換它：

j forward
mov ..., %reg
 forward:

并且假設(shè)分支完全是可預(yù)測的（并且所有分支的 95+% 是可預(yù)測的），對于 CPU 來說，分支實(shí)際上會好很多。

為什么？因為分支是可以預(yù)測的，當(dāng)它們被預(yù)測時，它們基本上就會消失。它們也在許多層面消失了。不僅僅是分支本身，而且就代碼的關(guān)鍵路徑而言，分支的條件消失了：CPU 仍然需要計算并檢查它，但從性能角度來看它“不再存在” ，因為它不持有任何東西了。

Pentium 4 在 CMOV 上確實(shí)很爛，但現(xiàn)代編譯器通常不再針對它進(jìn)行優(yōu)化。

事實(shí)上，現(xiàn)在 LLVM 往往更喜歡 CMOV 而不是分支。

例如，這是一個令人驚訝的編譯結(jié)果：

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 結(jié)果，而不是發(fā)出一個分支！

當(dāng)然，這只是因為 LLVM 觀察到兩個分支內(nèi)的工作量足夠輕，可以證明這種權(quán)衡是合理的……但這仍然很令人驚訝，不是嗎？

線性 SIMD 搜索

這種性能回歸非常煩人，而且 rustc 編譯器似乎不太可能很快修復(fù)它。

我決定用一個始終快速執(zhí)行并且對編譯器的變遷不那么敏感的實(shí)現(xiàn)來交換我的指數(shù) + 二進(jìn)制搜索。

鑒于數(shù)組很小，我決定實(shí)現(xiàn)一個簡單的 SIMD 無分支線性搜索。

使線性搜索無分支的技巧是將搜索問題重新表述為計算有多少元素小于針的問題。

這個想法轉(zhuǎn)化為以下標(biāo)量代碼：

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

不幸的是，SSE 的實(shí)現(xiàn)又長又臭：

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?`
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
}

該實(shí)現(xiàn)為我?guī)砹嗽?1.25 之前享受的性能。

04 二分查找的反擊

在閱讀了 dirtyhandscoding 的博客文章后^[14]，我決定再試一次二分查找。

這里的要點(diǎn)是簡化代碼庫。不僅 SIMD 的使用難以閱讀和維護(hù)，而且 SIMD 指令集也是特定于體系結(jié)構(gòu)的，這意味著我還必須維護(hù)算法的標(biāo)量版本。我獲得的性能提升只是蛋糕上的一顆櫻桃。

這次我會一直搜索整個 block。該塊的長度為 128 個元素，這意味著我們應(yīng)該能夠在 7 次比較中確定結(jié)果。這樣，我們就可以不惜一切代價進(jìn)行這 7 個比較。

當(dāng)然，我們也希望我們生成的代碼盡可能高效且完全無分支。

這是我能想出的最慣用的代碼來實(shí)現(xiàn)我們的目標(biāo)。

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]?????????????start?+=?len;
????????}
????}
????start
}

我沒想到用這么簡單的一段代碼就能達(dá)到我想要的效果。

這里的關(guān)鍵部分是我們沒有傳遞切片 ( &[u32]) 作為參數(shù)，而是傳遞了一個數(shù)組 ( &[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 確實(shí)超越了自己。想象一下那里發(fā)生了什么：LLVM 設(shè)法

• 展開 while 循環(huán)
• 證明start + len - 1總是小于 128 并刪除所有邊界檢查
• 在不那么瑣碎的地方發(fā)出 CMOV。
• 為第一個和最后一個迭代案例找到了一個非平凡的優(yōu)化。

讓我們一起分解這個匯編代碼：

在這個函數(shù)中，

• rax: 是我們的返回值，start在 Rust 代碼中。使所有這些有點(diǎn)混亂的一件事是我們還通過eax 和 al 寫入和讀取它。它占 64 位。雖然 rax 是一個 64 位寄存器，eax 和 al分別是指其最低 32 位和它的最低 8 位。
• esi 是我們的針參數(shù)。
• rdi 是我們數(shù)組中第一個元素的地址。

讓我們一步一步地完成匯編。

EAX 歸零

xor     eax, eax

這可能看起來很奇怪，但這只是設(shè)置 eax 為 0 的最常見方式。為什么？機(jī)器碼只需要 2 個字節(jié)。此外，現(xiàn)代 CPU 實(shí)際上不會在這里計算 XOR。它們只是將這條指令視為 “wink”，告訴它們我們想要一個值為 0 的寄存器。

第一次迭代

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

第一次迭代有點(diǎn)奇怪。顯然，LLVM 發(fā)現(xiàn)了一些特定于第一次迭代的優(yōu)化。但是第一次迭代有什么特別之處呢？此時，start 等于 0，其終值只能是 0 或 64。

rdi + 252 只是訪問數(shù)組第 63 個元素的指針運(yùn)算。( 252 = 63 * size_of::())

如果之前比較較低，setb al 指令設(shè)置 al 為 1。

shl 是位移指令。

Rust 代碼如下所示：

let?cmp?=?arr[63].cmp(&needle);
start?=??//
????if?cmp?==?Ordering::Lower?{
????????1
????}?else?{
????????0
????}
start?<<=?6;

因為64 = 2 ^ 6，我們確實(shí)以預(yù)期的 start = 64 或 start = 0 結(jié)束。。。

迭代 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 基準(zhǔn)測試

CPU 模擬器和微基準(zhǔn)測試

CPU 模擬器估計此代碼需要 9.55 個周期^[15]。很優(yōu)秀！請記住，我們正在搜索 128 個整數(shù)塊中的值。

相比之下，我們的 SIMD 線性搜索實(shí)現(xiàn)估計為 26.29 個周期。

我能想到的最好的 C++ 實(shí)現(xiàn)是在 Clang 上超過 12 個周期，在 GCC 上超過 40 個周期。

讓我們看看模擬器告訴我們的內(nèi)容是否真的轉(zhuǎn)化為現(xiàn)實(shí)世界。

Tantivy 有幾個微基準(zhǔn)測試來測量在發(fā)布列表上調(diào)用 seek 的速度。這些基準(zhǔn)測試將大致對應(yīng)于每次調(diào)用 Advance 時跳過的平均文檔數(shù)的倒數(shù)作為參數(shù)。跳躍越短，值越高。

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 (無分支內(nèi)聯(lián)二分查找)

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)

這一點(diǎn)也不差！在這個基準(zhǔn)測試中，Seek 現(xiàn)在大約快 11%。

真實(shí)場景的基準(zhǔn)測試

Tantivy 帶有一個搜索引擎基準(zhǔn)測試^[16]，可以比較不同的搜索引擎實(shí)現(xiàn)。它嘗試將不同類型的現(xiàn)實(shí)世界查詢與不同的數(shù)據(jù)集進(jìn)行比較。

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

由于 block-WAND 算法，使用 top-K 收集器的 OR 查詢也受益于這種優(yōu)化。

06 Conclusion

所以 LLVM 是完美的，看匯編代碼是徒勞的？在這一點(diǎn)上，你們中的一些人可能認(rèn)為這里的教訓(xùn)是 LLVM 在編譯慣用代碼方面做得如此出色，以至于查看匯編、手動展開內(nèi)容等只是浪費(fèi)時間。

我被告知過無數(shù)次，但我不同意。

為了得到這個版本的 rust 代碼，我不得不反復(fù)琢磨并確切地知道我想要什么。例如，這是我的第一個實(shí)現(xiàn)。

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]?????????????mid..range.end
????????}?else?{
????????????range.start..mid
????????};
????}
????range.start
}

它使用范圍而不是(start, len)獨(dú)立操作。LLVM 無法應(yīng)用我們在此代碼中討論的優(yōu)化的任何優(yōu)化。

我會進(jìn)一步優(yōu)化。本博客中的實(shí)現(xiàn)實(shí)際上不是 tantivy 中提供的代碼版本。雖然 rustc 今天在編譯這個函數(shù)方面做得很好，但我不相信未來的rustc版本也會這樣做。

例如，一年前，Rustc 1.41 未能刪除邊界檢查。為了獲得一致的編譯結(jié)果，tantivy 實(shí)際上調(diào)用了不安全的 get_unchecked。我有信心它是安全的嗎？我會在晚上睡覺嗎……我會的。rustc 1.55 生成的代碼提供了它安全的正式證明。

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

參考資料