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:

1. Pick a good algorithm.
2. Use a compiled language to write a Python extension.
3. 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.

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:

➘ 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:

➘ Lower numbers are better, ➚ Higher numbers are better

Notice that:

1. The new version has fewer branches: 19 per value, instead of 27.
2. Branch mispredictions are completely gone.
3. 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.

Getting rid of branches and wasted work

Having improved the code’s mechanical sympathy, I now have the opportunity to return to other sources of speed: tuning compilation, and making the code more algorithmically efficient by doing less work.

First, the branchless code still has 19 branches per value in the input array. 7 or 8 of those are the for _ in 0..n_iterations: deciding if the for loop should continue or finish requires a branch, and n_iterations is log2 of 255 or so. But where are the rest coming from?

These additional branches exist because the Rust compiler adds bounds checks to vector and array access. Every indexed read and write has a check to see if the index is out of bounds, in order to ensure memory safety: writing to index 100 of a vector of length 10 can lead to crashes or memory corruption. Bounds checks are good, insofar as they catch bugs, but they’re bad insofar as they add more computation. So I’m going to use Rust unsafe code to call APIs like get_unchecked() that bypass the bounds checks, meaning it’s now my responsibility to ensure that the code is still correct.

Second, notice the calculation for half/remaining_size is exactly the same for every single binary search. That’s wasted work. So I’m going to try to precalculate those values, doing less computational work.

Here’s what the resulting code looks like:

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

    // Precalculate the halving values. Create a Vec of size
    // n_iterations, filled with zeros:
    let mut halves = vec![0usize; n_iterations];
    let mut remaining_size = size;
    for i in 0..n_iterations {
        let half = remaining_size / 2;
        halves[i] = half;
        remaining_size -= half;
    }
    // Tell Rust (and reader) that halves is not mutable
    // anymore:
    let halves = halves;

    // Safety check for invariants that later unsafe
    // operations depend on: sum of halves is less than the
    // length of boundaries.
    assert!(
        halves.iter().copied().sum::<usize>() <
            boundaries.len()
    );

    let mut result = Vec::with_capacity(arr.len());
    for value in arr {
        let mut left = 0;

        for half in &halves {
            let middle = left + half;
            left = select_unpredictable(
                // SAFETY: `left` and `middle` can be at
                // most the sum of `halves`, and the code
                // above asserts that this sum is smaller
                // than the length of `boundaries`.
                less_than(
                    unsafe {
                        *boundaries.get_unchecked(middle)
                    },
                    value,
                ),
                middle,
                left,
            );
        }

        left = select_unpredictable(
            // SAFETY: See comment above.
            less_than(
                unsafe {
                    *boundaries.get_unchecked(left)
                },
                value,
            ),
            left + 1,
            left,
        );
        result.push(left);
    }
    result
}

And here’s its performance:

➘ Lower numbers are better, ➚ Higher numbers are better

As expected, there are now fewer branches. Overall, there are fewer instructions (good), and a modest decline in IPC (bad), and most importantly the code is faster.

Auto-vectorization and SIMD

Finally, I will return to mechanical sympathy. It’s worth noticing that different values in the input array get basically the same operation done on them. In our original version, the number of inner iterations differed, but now the same number of iterations is happening for each.

If the same operation is happening on multiple values in memory, specialized SIMD CPU instructions can help: “Same Instruction, Multiple Data” or SIMD for short, they are custom built to speed up your code by processing data in parallel. You can use them manually, or you can have the Rust compiler “auto-vectorize” your code, i.e. automatically generate SIMD instructions if it decides it’s useful.

To achieve this I’m going to do two things.

First, I tell Rust that’s it’s not 2004 anymore, and specifically that it can generate CPU instructions that require modern hardware, x86-64 machines from the past 10 years or so. This mean access to a much wider range of SIMD CPU instructions, at the cost of preventing the resulting machine code from running on old computers. Specifically, I do:

$ export RUSTFLAGS="-C target-cpu=x86-64-v3"

Then, I restructure the code so the iteration over halves is an outer loop, and iterating over values in arr is an inner loop; previously it was the other way around. That way there are much more obviously identical operations happening on adjacent data, hopefully allowing the compiler to figure out some way to convert this code to SIMD instructions:

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

    let mut halves = vec![0usize; n_iterations];
    let mut remaining_size = size;
    for i in 0..n_iterations {
        let half = remaining_size >> 1;
        halves[i] = half;
        remaining_size -= half;
    }
    let halves = halves;

    assert!(
        halves.iter().copied().sum::<usize>() <
            boundaries.len()
    );

    let mut result = Vec::with_capacity(arr.len());
    // Iterate in chunks of 16:
    for chunk in arr.chunks(16) {
        // Create an array (like a Vec, but fixed length and
        // on the stack) of length 16 and filled with zeros:
        let mut lefts = [0; 16];

        // Instead of iterating over values first, iterate
        // over halves first, and then iterate over a chunk
        // of values. Iterating over all values once per
        // `half` would worsen memory cache locality, thus
        // the use of chunks.
        for half in &halves {
            for (i, value) in chunk.iter().enumerate() {
                let left = unsafe { lefts.get_unchecked(i) };
                let middle = left + half;
                *unsafe { lefts.get_unchecked_mut(i) } =
                    select_unpredictable(
                        less_than(
                            unsafe {
                                *boundaries.get_unchecked(middle)
                            },
                            value,
                        ),
                        middle,
                        *left,
                    );
            }
        }

        for (value, mut left) in chunk.iter().zip(lefts) {
            left = select_unpredictable(
                less_than(
                    unsafe { *boundaries.get_unchecked(left) },
                    value,
                ),
                left + 1,
                left,
            );
            result.push(left);
        }
    }
    result
}

And here’s a performance comparison of all the variations; the latest one is the fastest:

➘ Lower numbers are better, ➚ Higher numbers are better

This final version uses fewer instruction. When I compiled normally, without a target architecture of x86-64-v3—not shown here—it used far more instructions and wasn’t quite as fast, though it was still the fastest of the four, suggesting the auto-vectorization was successful. Custom SIMD usage might allow for even more speedups.

The one downside is that this compiled code won’t run on older computers. So one solution is to ship the final function with two copies, one compiled with SIMD for x86-64-v3 and one compiled normally with full backwards compatibility, and at runtime pick the one that can run on the current CPU.

You can speed up your code too

There are more opportunities for speed here. First, and most significantly, I could use parallelism to take advantage of multiple CPU cores. And there are additional potential ways to speed up the single-threaded code too; I’ve successfully implemented one, just to see if it would work.

But in the end this is just an article with some code examples. The real value here is not me demonstrating another 14% speedup, it’s you going and speeding up your code. This is something you can learn, just like I did: I can now write far faster software than I could when I started, even if there’s still plenty more for me to learn.

Some books that can help you build mechanical sympathy include:

• Computer Systems: A Programmer’s Perspective, 3rd edition, by Bryant and O’Halloran. This textbook covers how CPUs work, but unlike most it focuses only on the parts that are relevant to programmers.
• Performance Analysis and Tuning on Modern CPUs, 2nd Edition, by Denis Bakhalov. This free book talks about critical tools for measuring performance, and how to interpret their results; I paid for the printed version. He also has a free online course.
• Algorithms for Modern Hardware, by Sergey Slotin, is a free online book.

These books mostly focus on mechanical sympathy, rather than on other sources of performance. They assume you know some amount of C. And because these are general purpose books, you will need to translate what you’ve learned into the kind of code you’re writing.

So if you are specifically speeding up data science or scientific computing software used from Python, you might want to start with reading my new book on Python performance, now in early access. Besides its focus on this particular domain of software, the book:

• Acts as a broader introduction to computational performance, covering different approaches to speed: algorithmic efficiency, compilation, mechanical sympathy, and parallelism.
• Only assumes you know Python.
• Introduces you to all of the concepts I discussed in this article, and many others besides.

You can learn more about the Practices of Performance here.