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

查看原文 / View Original 编辑 / Edit

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

6× faster binary search: from compiled code to mechanical sympathy by Itamar Turner-Trauring Last updated 11 Jul 2026, originally created 11 Jul 2026

6 倍速的二分查找：从编译代码到“机械同理心” 作者：Itamar Turner-Trauring 最后更新：2026 年 7 月 11 日，最初创建于 2026 年 7 月 11 日

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?

你该如何加速计算密集型的 Python 代码？一个常见且有效的起点是：选择一个好的算法。使用编译语言编写 Python 扩展。或许还可以增加并行处理，以便利用多个 CPU 核心。但如果你需要更快的速度呢？

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.

考虑以下实际问题，它是 scikit-learn 的梯度直方图提升（gradient histogram boosting）算法中的一个步骤：你有一个巨大的浮点数数组，想要将它们均匀地分配到 0-254 的整数范围内。scikit-learn 的实现方式是将浮点值的完整范围划分为 255 个桶，创建一个已排序的桶边界数组，然后使用二分查找为每个值选择合适的桶。该二分查找是用编译语言实现的，并且可以在多个核心上并行运行。

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.

最近，作为我在 Quansight 工作的一部分，并受到 Paul Khuong 两篇文章的启发，我显著提升了该实现的性能。如何做到的？通过确保代码不会与 CPU 的工作方式“作对”。在本文中，我将通过一个简化示例带你了解这一加速过程。随后，我将演示一系列额外的优化，最终版本的运行速度将比原始版本快 6 倍。

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.

值得注意的是，我将快速提及许多不同的底层硬件主题：指令级并行、分支（预测）错误、内存缓存、SIMD 等。这只是一篇文章，只能简要介绍可能实现的目标，无法作为深入的教程。因此，我将在文章末尾讨论如何进一步学习这些主题。

起点：标准二分查找

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:

原始的 scikit-learn 代码是用 Cython 实现的，但在本文中我将使用 Rust。以下是一个非常标准的二分查找实现（基于 NumPy 中的实现），专为给定边界数组查找桶这一用例而设计：

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
}

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.

我该如何加速这个二分查找实现？它已经使用了可扩展的算法和编译语言。并行化当然是一个选择，但我将采用另一种方法：机械同理心（mechanical sympathy），即更好地理解 CPU 的工作原理。我将首先快速回顾现代 CPU 如何在单个核心内并行运行代码。

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.

对于 Python 代码，一个合理的思维模型是代码一次执行一条指令。算术运算量增加一倍，代码运行速度就会慢一倍。一旦切换到编译语言，操作有时会映射为一到两条 CPU 指令，这种思维模型就不再准确了。现代 CPU 有时可以在单个 CPU 核心上同时运行多个独立的 CPU 指令（“指令级并行”），从而实现更快的执行速度。

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?

然而，由 if/while/for 表达式产生的代码分支带来了一个问题：你的代码可能走这条路，也可能走那条路。面对两种选择，CPU 应该尝试并行执行哪一组可能的后续指令呢？

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 ve…

为了确保快速执行，CPU 拥有一个分支预测器，它会启发式地选择并行执行哪个分支。如果猜对了，你的代码就会更快。如果猜错了，CPU 最终会察觉到，撤销错误的工作，然后去执行正确的分支……这意味着你的代码变慢了。在某些情况下，会慢得多。上述二分查找算法不幸地非常……

← 返回新闻列表 / Back to news list