Бинарный поиск ускорили в 8 раз, не меняя ни алгоритм, ни язык: 16,6 процента промахов предсказателя ветвлений превратились в ноль

Итамар Тёрнер-Трауринг разобрал шаг из градиентного бустинга в scikit-learn: миллион чисел с плавающей точкой нужно разложить по 255 корзинам, для чего по массиву границ гоняется бинарный поиск. Код уже скомпилированный и уже параллелится по ядрам, алгоритм оптимальный. Статья опубликована 11 июля и обновлена 18-го, работа сделана в рамках Quansight.

Остался запас, который не виден на уровне алгоритма. Современное ядро процессора выполняет несколько инструкций одновременно и угадывает, куда пойдёт ветвление. В бинарном поиске сравнение по определению непредсказуемо: каждая итерация с равной вероятностью идёт влево или вправо, и предсказатель ошибается примерно в половине случаев, а конвейер каждый раз сбрасывается.

🔘 исходная версия: 45 200,4 микросекунды, 16,6 процента неверных предсказаний, 0,7 инструкции за такт, около 27 инструкций ветвления на одно значение;
🔘 первая переделка убирает ветвление, заменяя его условным присваиванием через select_unpredictable: 9 685,2 мкс, промахов ноль, 3,2 инструкции за такт, ветвлений 19 на значение;
🔘 предвычисление половины диапазона и доступ без проверки границ дают 7 280,4 мкс и обрушивают число инструкций ветвления до 6 020 571 на весь прогон;
🔘 финальная версия обрабатывает значения кусками по 16 штук с внешним циклом по шагам поиска: 5 453,4 мкс и 4,9 инструкции за такт;
🔘 любопытно, что она выполняет больше инструкций, чем предыдущая, но идёт быстрее: независимые куски работы позволяют процессору выполнять их одновременно;
🔘 итог — примерно восьмикратное ускорение на том же алгоритме, том же языке и одном ядре.

Автор оговаривает, что это упрощённый пример, а не полная реализация из scikit-learn, и что статья не заменяет учебник по устройству процессора. В обновлении он честно пишет, что раньше ошибся со сравнением чисел с плавающей точкой, и после исправления SIMD перестал давать выигрыш, хотя код от этого стал только быстрее. Дальнейший запас автор видит в параллелизме по ядрам.

Польза для питониста прямая: когда расчёт уже переписан на расширение и всё равно кажется медленным, следующий уровень — не алгоритм, а поведение процессора на этом коде.

Полная статья: https://pythonspeed.com/articles/branchless-binary-search/

@zen_of_python