← На главную

Разработчик из Quansight ускорил бинарный поиск scikit-learn в 6-8 раз

12.07.2026 07:00 · hackernews

В scikit-learn есть шаг, где массив чисел с плавающей точкой нужно равномерно разбросать по 255 корзинам (0–254). Оригинальная реализация брала отсортированные границы и для каждого значения делала бинарный поиск — на Cython, с параллельностью на нескольких ядрах. Автор (из Quansight, вдохновлённый постами Paul Khuong) ускорил этот код в шесть раз, не меняя алгоритм, а просто перестав бороться с CPU.

Классический бинарный поиск — ветвящийся: if boundaries[middle] < value то туда, иначе сюда; цикл while может закончиться на любой итерации. Это убивает предсказание ветвлений процессора — в тесте на миллионе случайных значений 16,6% веток были mispredicted, а IPC (инструкций за такт) всего 0,7. Первый фикс — сделать поиск безусловным: фиксированное число итераций (log2(количество границ)), а выбор левой границы — через std::hint::select_unpredictable() из Rust, который компилятор превращает в условное перемещение без ветвления. Пропали misprediction (0%), IPC подскочил до 3,1, время упало с 44,6 мс до 9,8 мс.

Дальше — убрали лишние граничные проверки: Rust-компилятор добавляет checks на каждый доступ к массиву. Через unsafe и get_unchecked() их отключили. Заодно предрассчитали halves (сдвиги середины) — раньше они считались заново для каждого значения, хотя одинаковы. Время снизилось до 7,3 мс, инструкций стало вдвое меньше.

Финальный трюк — перестановка циклов. В исходном коде внутренний цикл по halves был зависим: left на каждой итерации зависел от предыдущей. Теперь внешний цикл — по halves, а внутренний — по чанкам из 16 значений. Так CPU видит независимые операции над разными данными и может выполнять их параллельно. IPC вырос до 4,9, время — 5,45 мс. Итоговое ускорение в ~8 раз относительно классической версии.

Автор отмечает, что можно добавить ещё параллелизм на нескольких ядрах, но цель статьи — показать, что понимание работы процессора (механическая симпатия) даёт огромный выигрыш. Тем, кто хочет разобраться глубже, он рекомендует свою новую книгу «Practices of Performance», а также классические книги по оптимизации низкоуровневого кода.

Читать оригинал →