6× faster binary search: from compiled code to mechanical sympathy

How do you speed up computational Python code? A common, and useful, starting point is:

Pick a good algorithm. Use a compiled language to write a Python extension. Maybe add parallelism so you can use multiple CPU cores.

But what if you need more speed? Consider the following real problem, one of the steps in scikit-learn’s gradient histogram boosting algorithm:

You have a large array of floating point numbers. You want to assign them to the integer range 0-254, spread out evenly.

scikit-learn implements this by splitting up the full range of float values into 255 buckets, creating a sorted array of bucket boundaries, and then using binary search to choose the appropriate bucket for each value. The binary search is implemented in a compiled language, and it can run in parallel on multiple cores.

Recently, as part of my work at Quansight, and inspired by two posts by Paul Khuong, I sped up this implementation significantly. How? By making sure the code wasn’t fighting against the CPU.

In this article I’m going to walk you through that speed-up, demonstrated on a simplified example. Then I’m going to demonstrate a series of additional optimizations, with the final version running 6× faster than the original one.

It’s worth knowing that I will be speeding through mentions of many different low-level hardware topics: instruction-level parallelism, branch (mis)prediction, memory caches, SIMD, and more. This is only one article, it can only briefly introduce you to what’s possible, it can’t function as an in-depth tutorial. So I’ll talk about how you can learn more about these topics at the end of the article.

The starting point: Standard binary search

The original scikit-learn code was implemented in Cython, but I’m going to use Rust for this article. Here’s a pretty standard implementation of binary search (based on the one in NumPy), designed for this use case of finding a bucket given an array of boundaries:

use std::cmp::Ordering;

/// Rust doesn't let you compare floats with the normal < /// operator (because NaN makes comparison results /// inconsistent), so implement a custom function to do so.

fn less_than(a: f64, b: &f64) -> bool { a.total_cmp(b) == Ordering::Less }

/// Convert floating points values into integer values by /// finding which bucket they fit in, given bucket /// boundaries. fn bucketize_classic_impl( arr: &[f64], boundaries: &[f64], ) -> Vec<usize> { // A Vec or vector is Rust's equivalent of a Python // list. Here I create an empty Vec with enough memory // allocated to store `arr.len()` values: let mut result = Vec::with_capacity(arr.len()); for value in arr { // Standard binary search algorithm: let mut min_idx = 0; let mut max_idx = boundaries.len(); while min_idx < max_idx { let middle = min_idx + ((max_idx - min_idx) / 2); if less_than(boundaries[middle], value) { min_idx = middle + 1; } else { max_idx = middle; } } // This is equivalent to a_list.append() in Python: result.push(min_idx); } // Return the result: result }

For completeness, here’s how you can hook it up to Python, so it takes and returns NumPy arrays; this is boilerplate, so I’m not going to show it for later functions. You can just skip it if you’re not familiar with Rust and PyO3, or don’t particularly care; it’s not relevant to the rest of the article.

Click here to see the code

use pyo3::prelude::*; use numpy::ndarray::Array; use numpy::{PyArray1, PyReadonlyArray1};

#[pyfunction] fn bucketize_classic<'py>( py: Python<'py>, // These two arguments are 1-dimensional arrays of // floats: arr: PyReadonlyArray1<f64>, boundaries: PyReadonlyArray1<f64>, ) -> PyResult<Bound<'py, PyArray1<usize>>> { let result =

bucketize_classic_impl( arr.as_slice().unwrap(), boundaries.as_slice().unwrap(), ); // Convert a Rust vector into a 1D NumPy array we can // pass back to Python: let result = PyArray1::from_owned_array(py, Array::from_vec(result)); Ok(result) }

Branch mispredictions slow your code down

How can I speed up this implementation of binary search? It’s already using a scalable algorithm, and a compiled language. Parallelism is certainly an option, but I’m going to use a different approach: mechanical sympathy, a better understanding of how the CPU works. I’ll start with a very quick review of how modern CPUs run code in parallel within a single core.

A reasonable mental model of Python code is that the code is executed one instruction at a time. Do twice as many arithmetic operations, and the code will run twice as slow. Once you switch to a compiled language, with operations sometimes mapping to one or two CPU instructions, that mental model is no longer correct. Modern CPUs can sometimes run multiple independent CPU instructions at once (“instruction-level parallelism”) on a single CPU core, resulting in faster execution.

fn two_adds(a: i64, b: i64, c: i64, d: i64) -> i64 { // Your CPU can probably run these two adds in parallel, // on a single core, automatically: let t1 = a + b; let t2 = c + d; // Return the result: t1 + t2 }

However, branches in the code created by if/while/for expressions pose a problem: your code might go one way, or the other. Given two choices, which set of possible future instructions should the CPU try to execute in parallel?

fn maybe_add( a: i64, b: i64, c: i64, d: i64, add: bool ) -> i64 { let t1 = a + b; // Which of these two branches should the CPU run in // parallel with `a + b`?

let t2 = if add { c + d } else { c * d }; t1 + t2 }

To ensure fast execution, the CPU has a branch predictor that heuristically chooses which branch to execute in parallel. If the guess is correct, your code is faster. If the guess is wrong, the CPU will eventually notice, undo the incorrect work, and then go execute the correct branch… which means your code is slower. In some cases, much slower.

The binary search algorithm above is unfortunately very unpredictable given the input data. Recall that the bucket boundaries are chosen so that the input values are spread evenly across all buckets. That means that choosing whether to go left or right in the binary search is not at all predictable:

// The CPU can't reliably guess which path will be taken: if less_than(boundaries[middle], value) { min_idx = middle + 1; } else { max_idx = middle; }

Similarly, the number of iterations may vary:

// How many times will this loop continue before stopping? // It's impossible to know given the particular data being // used. while min_idx < max_idx { // ... }

To validate this hypothesis, I can use the CPU’s hardware counters (exposed to Python via py-perf-event) to measure how many branches the running code executes, and how many of these branches get mispredicted. Here are the inputs I’ll be using:

import numpy as np from numba import jit

# Values between 0 and 1: DATA = np.random.random(1_000_000)

# Bucket boundaries, evenly spaced out: BOUNDARIES = np.linspace(0.0, 1.0, 255)[1:-1]

And here’s the result of running the code:

Code ➘ Elapsed µ-seconds

➘ Branch instructions ➘ Branch misprediction % bucketize_classic(DATA, BOUNDARIES) 45,870.2   26,997,038 16.6%

➘ Lower numbers are better

16% of branches being mispredicted isn’t great, and there’s also quite a lot of branches per value. DATA has 1,000,000 values, so with 27 million branches total that’s 27 branches per value.

Switching to branchless execution

I’m going to get rid of both sources of unpredictable branches. As far as the number of while loop iterations, I’m instead going to iterate a fixed number of times, log2 of the number of buckets. For some value this might involve a bit more work, if previously the bucket would be found in an iteration or two, but the saving in speed from avoiding branch mispredictions will make it worth it.

For the if expression, the current code sometimes set one variable, sometimes another. I’m going to replace that with code that always sets the same variable, even if it’s unchanged. Then, I’m going to use Rust’s std::hint::select_unpredictable(), which tells the compiler to avoid emitting branches, if possible. Often CPUs have special instructions for conditionally choosing between two values, without a branch.

Here’s what the new, branchless version looks like:

use std::hint::select_unpredictable;

fn bucketize_branchless_impl( arr: &[f64], boundaries: &[f64], ) -> Vec<usize> { let size = boundaries.len(); let n_iterations = (size as f64).log2().ceil() as usize; let mut result = Vec::with_capacity(arr.len());

for value in arr { let mut left = 0; let mut remaining_size = size;

// Branchless binary search: keep cutting the search // area in half.

for _ in 0..n_iterations { let half = remaining_size / 2; let middle = left + half; // Conditional operation: if bool, then a, else // b. Hopefully generated by the compiler // without an actual branch. left = select_unpredictable( less_than(boundaries[middle], value), middle, left, ); remaining_size -= half; }

// Fix an off-by-one difference from the original // algorithm. left = select_unpredictable( less_than(boundaries[left], value), left + 1, left, ); result.push(left); } result }

Here’s a comparison of the performance of the two versions:

Code ➘ Elapsed µ-seconds   ➘ CPU instructions ➘ Branch instructions ➘ Branch misprediction % ➚ IPC bucketize_classic(DATA, BOUNDARIES) 45,810.2   184,909,902 26,997,064 16.6% 1.1 bucketize_branchless(DATA, BOUNDARIES) 13,197.8 🏆 188,101,254 19,020,569 0.0% 4.0

➘ Lower numbers are better, ➚ Higher numbers are better

Notice that:

The new version has fewer branches: 19 per value, instead of 27. Branch mispredictions are completely gone. IPC (instructions per cycle) is much higher. This measures the CPU’s ability to run multiple instructions in parallel, and higher is better: with fewer and more predictable branches, the CPU can now use more instruction parallelism.

The result is an implementation that is far faster, even though it uses slightly more CPU instructions.