Static search trees: 40x faster than binary search

Table of Contents1 Introduction1.1 Problem statement1.2 Motivation1.3 Recommended reading1.4 Binary search and Eytzinger layout1.5 Hugepages1.6 A note on benchmarking1.7 Cache lines1.8 S-trees and B-trees2 Optimizing find2.1 Linear2.2 Auto-vectorization2.3 Trailing zeros2.4 Popcount2.5 Manual SIMD3 Optimizing the search3.1 Batching3.2 Prefetching3.3 Pointer arithmetic3.3.1 Up-front splat3.3.2 Byte-based pointers3.3.3 The final version3.4 Skip prefetch3.5 Interleave4 Optimizing the tree layout4.1 Left-tree4.2 Memory layouts4.3 Node size \(B=15\)4.3.1 Data structure size4.4 Summary5 Prefix partitioning5.1 Full layout5.2 Compact subtrees5.3 The best of both: compact first level5.4 Overlapping trees5.5 Human data5.6 Prefix map5.7 Summary6 Multi-threaded comparison7 Conclusion7.1 Future work7.1.1 Branchy search7.1.2 Interpolation search7.1.3 Packing data smaller7.1.4 Returning indices in original data7.1.5 Range queries7.1.6 Sorting queries7.1.7 Suffix array searchingIn this post, we will implement a static search tree (S+ tree) for high-throughput searching of sorted data, as introduced on Algorithmica. We’ll mostly take the code presented there as a starting point, and optimize it to its limits. For a large part, I’m simply taking the ‘future work’ ideas of that post and implementing them. And then there will be a bunch of looking at assembly code to shave off all the instructions we can. Lastly, there will be one big addition to optimize throughput: batching.All source code, including benchmarks and plotting code, is at github:RagnarGrootKoerkamp/static-search-tree.Discuss on r/programming, hacker news, twitter, bsky, or youtube.1 Introduction Link to heading1.1 Problem statement Link to headingInput.

A sorted list of \(n\) 32bit unsigned integers vals: Vec<u32>.Output. A data structure that supports queries \(q\), returning the smallest element of vals that is at least \(q\), or u32::MAX if no such element exists. Optionally, the index of this element may also be returned.Metric. We optimize throughput. That is, the number of (independent) queries that can be answered per second. The typical case is where we have a sufficiently long queries: &[u32] as input, and return a corresponding answers: Vec<u32>.1Note that we’ll usually report reciprocal throughput as ns/query (or just ns), instead of queries/s. You can think of this as amortized (not average) time spent per query.Benchmarking setup. For now, we will assume that both the input and queries are simply uniform random sampled 31bit integers2.Code. In code, this can be modelled like this:1 2 3 4 5 6 7 8 9 trait SearchIndex { /// Two functions with default implementations in terms of each other. fn query_one(&self, query: u32) -> u32 { Self::query(&[query])[0] } fn query(&self, queries: &[u32]) -> Vec<u32> { queries.iter().map(|&q| Self::query_one(q)).collect() } } Code Snippet 1: Trait that our solution should implement.1.2 Motivation Link to headingAside from doing this project just for the fun of it, there is some higher goal. One of the big goals of bioinformatics is to make efficient datastructures to index DNA, say a single human genome (3 billion basepairs/characters) or even a bunch of them. One such datastructure is the suffix array (wikipedia), that sorts the suffixes of the input string. Classically, one can then find the locations where a string occurs by binary searching the suffix array.This project is a first step towards speeding up the suffix array search.Also note that we indeed assume that the input data is static, since usually we use a fixed reference genome.1.3 Recommended reading Link to headingThe classical solution to this problem is binary search, which we will briefly visit in the next section.

A great paper on this and other search layouts is “Array Layouts for Comparison-Based Searching” by Khuong and Morin (2017). Algorithmica also has a case study based on that paper.This post will focus on S+ trees, as introduced on Algorithmica in the followup post, static B-trees. In the interest of my time, I will mostly assume that you are familiar with that post.I also recommend reading my work-in-progress introduction to CPU performance, which contains some benchmarks pushing the CPU to its limits. We will use the metrics obtained there as baseline to understand our optimization attempts.Also helpful is the Intel Intrinsics Guide when looking into SIMD instructions. Note that we’ll only be using AVX2 instructions here, as in, we’re assuming intel. And we’re not assuming less available AVX512 instructions (in particular, since my laptop doesn’t have them).1.4 Binary search and Eytzinger layout Link to headingAs a baseline, we will use the Rust standard library binary search implementation. 1 2 3 4 5 6 7 8 9 10 pub struct SortedVec { vals: Vec<u32>, } impl SortedVec { pub fn binary_search_std(&self, q: u32) -> u32 { let idx = self.vals.binary_search(&q).unwrap_or_else(|i| i); self.vals[idx] } } Code Snippet 2: The binary search in the Rust standard library.The main conclusion of the array layouts paper (Khuong and Morin 2017) is that the Eytzinger layout is one of the best in practice. This layout reorders the values in memory: the binary search effectively is a binary search tree on the data, with as root the middle node, then the nodes at positions \(\frac 14 n\) and \(\frac 34 n\), then \(\frac 18n, \frac 38n, \frac 58n, \frac 78n\), and so on.

The main benefit of the Eytzinger layout is that all values needed for the first steps of the binary search are close together, so they can be cached efficiently: we put the root at index \(1\) and the two children of the node at index \(i\) are at \(2i\) and \(2i+1\). This means that we can effectively prefetch the next cache line, before knowing whether we need index \(2i\) or \(2i+1\). This can be taken a step further and we can prefetch the cache line containing indices \(16i\) to \(16i+15\), which are exactly the values needed 4 iterations from now. For a large part, this can quite effectively hide the latency associated with the traversal of the tree. 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 pub struct Eytzinger { /// The root of the tree is at index 1. vals: Vec<u32>, } impl Eytzinger { /// L: number of levels ahead to prefetch. pub fn search_prefetch<const L: usize>(&self, q: u32) -> u32 { let mut idx = 1; while (1 << L) * idx < self.vals.len() { idx = 2 * idx + (q > self.get(idx)) as usize; prefetch_index(&self.vals, (1 << L) * idx); } // The last few iterations don't need prefetching anymore. while idx < self.vals.len() { idx = 2 * idx + (q > self.get(idx)) as usize; } let zeros = idx.trailing_ones() + 1; let idx = idx >> zeros; self.get(idx) } } Code Snippet 3: Implementation of searching the Eytzinger layout, with \(L=4\) levels of prefetching.

If we plot these two, we see that Eytzinger layout performs as good as binary search when the array fits in L2 cache (256kB for me, the middle red line), but starts to be much better than binary search as the array grows to be much larger than the L3 cache (12MB). In the end, Eytzinger search is around 4 times faster, which nicely corresponds to being able to prefetch 4 iterations of cache lines from memory at a time.Figure 1: Query throughput of binary search and Eytzinger layout as the size of the input increases. At 1GB input, binary search needs around 1150ns/query, while Eytzinger is 6x faster at 200ns/query. (I’m sorry for the two blue lines here. The top one is binary search and the bottom one Eytzinger. Getting all the plots to work took long enough, and customizing all the colours is annoying, so I’m just using cycling through the default colours. Unfortunately that ended up with two equal colours here. At least colours will be consistent from one plot to the next.)1.5 Hugepages Link to headingFor all experiments, we’ll make sure to allocate the tree using 2MB hugepages by default, instead of the usual 4kB pages. This reduces pressure on the translation lookaside buffer (TLB) that translates virtual memory addresses to hardware memory addresses, since its internal table of pages is much smaller when using hugepages, and hence can be cached better.With transparent hugepages enabled, they are automatically given out whenever allocating an exact multiple of 2MB, and so we always round up the allocation for the tree to the next multiple of 2MB. However, it turns out that small allocations below 32MB still go on the program’s heap, rather than asking the kernel for new memory pages, causing them to not actually be hugepages. Thus, all allocations we do are actually rounded up to the next multiple of 32MB instead.All together, hugepages sometimes makes a small difference when the dataset is indeed between 1MB and 32MB in size. Smaller data structures don’t really need hugepages anyway.

Enabling them for the Eytzinger layout as in the plot above also gives a significant speedup for larger sizes.1.6 A note on benchmarking Link to headingThe plots have the size of the input data on the logarithmic (bottom) x-axis. On the top, they show the corresponding number of elements in the vector, which is 4 times less, since each element is a u32 spanning 4 bytes. Measurements are taken at values \(2^i\), \(1.25 \cdot 2^i\), \(1.5\cdot 2^i\), and \(1.75\cdot 2^i\).The y-axis shows measured time per query. In the plot above, it says latency, since it is benchmarked as for q in queries { index.query(q); }. Even then, the pipelining and out-of-order execution of the CPU will make it execute multiple iterations in parallel. Specifically, while it is waiting for the last cache lines of iteration \(i\), it can already start executing the first instructions of the next query. To measure the true latency, we would have to introduce a loop carried dependency by making query \(i+1\) dependent on the result of query \(i\). However, the main goal of this post is to optimize for throughput, so we won’t bother with that.Thus, all plots will show the throughput of doing index.query(all_queries).For the benchmarks, I’m using my laptop’s i7-10750H CPU, with the frequency fixed to 2.6GHz using Code Snippet 4.31 sudo cpupower frequency-set -g powersave -d 2.6GHz -u 2.6GHz Code Snippet 4: Pinning the CPU frequency to 2.6GHz.Also relevant are the sizes of the caches: 32KiB L1 cache per core, 256KiB L2 cache per core, and 12MiB L3 cache shared between the physical 6 cores. Furthermore, hyper-threading is disabled.All measurements are done 5 times. The line follows the median, and we show the spread of the 2nd to 4th value (i.e., after discarding the minimum and maximum). Observe that in most of the plot above, the spread is barely visible!