6x faster batched binary search in Rust via branch prediction fixes
scikit-learn's binary search bucketization gets a 6x Rust speedup by eliminating CPU branch mispredictions via branchless code.
A Quansight engineer rewrote the binary search step used in scikit-learn's gradient histogram boosting algorithm in Rust, achieving a 6x speedup. The core idea is 'mechanical sympathy': writing code that works with the CPU's branch predictor instead of against it.
The original binary search bucketizes floating-point values into 255 evenly spread buckets, but because the input is uniformly distributed, the CPU can't reliably predict which branch each comparison will take. Hardware performance counters show that on one million values, roughly 27 million branch instructions are executed, with 16.6% mispredicted—a costly overhead.
The fix replaces the variable-length while loop with a fixed number of iterations (log2 of the bucket count) and converts the branching if/else logic into branchless code. The approach illustrates how instruction-level parallelism and low-level hardware behavior can yield real performance gains, offering a concrete technique for engineers optimizing compiled Python extensions.