Недавно мне довелось работать над новым алгоритмом приближённого поиска ближайших соседей, который называется RaBitQ. Автор этого алгоритма уже предоставил достаточно скоростную реализацию на C++. Я попытался переписать этот алгоритм на Rust (ещё один случай «а почему бы не переписать на Rust»). Однако я обнаружил, что моя реализация гораздо медленнее оригинальной. Далее я расскажу, как шаг за шагом доработал её производительность.

Подготовка среды

Датасеты

В данном случае максимально важно подобрать адекватные датасеты. Поскольку в вышеупомянутой статье, описывающей реализацию на C++, уже продемонстрированы некоторые результаты по датасетам sift_dim128_1m_l2 и gist_dim960_1m_l2, варианты с 128 и 960 измерениями считаем типичными, а 1_000_000 векторов должно хватить для расстановки контрольных точек. Эти датасеты можно скачать отсюда. (Да, я знаю, что на этом сайте не поддерживается TLS и обеспечивается скачивание только по FTP).

В этих датасетах используется формат fvecs/ivecs, по сути он обычный векторный:

| dim (4 bytes) | vector (4 * dim bytes) |
| dim (4 bytes) | vector (4 * dim bytes) |
...
| dim (4 bytes) | vector (4 * dim bytes) |

Скрипт для чтения и записи можете скачать здесь.

Инструмент-профилировщик

Для профилирования кода Rust я пользуюсь инструментом samply. Он хорошо интегрируется с профилировщиком Firefox. Кроме того, работая с ним, можно поделиться с коллегами полученными результатами профилирования, выложив их в облако. Вот пример, в котором профилирована версия на C++. Наиболее распространённые представления — FlameGraph и CallTree. Не забудьте предоставить доступ для анализа событий производительности и увеличьте предел mlock:

echo '1' | sudo tee /proc/sys/kernel/perf_event_paranoid
sudo sysctl kernel.perf_event_mlock_kb=2048

Вам также пригодится обозреватель компилятора GodBolt, в котором удобно сравнивать код функций сборки между C++ и Rust.

Профиль Cargo

Чтобы включить в релизную сборку отладочную информацию, можно добавить ещё один профиль в Cargo.toml:

[profile.perf]
inherits = "release"
debug = true
codegen-units = 16

Удобство профилирования сильно зависит от того, каковы издержки на компиляцию, и как быстро выполняется программа.

• У cargo build скорость компиляции выше, но код может выполняться медленнее, чем на чистом Python

• cargo build --release выполняется быстро, но на его компиляцию может потребоваться много времени

Для расстановки контрольных точек у нас нет иных вариантов, кроме как использовать opt-level = 3.

Несколько раз мне встречался совет выставлять следующие настройки:

codegen-units = 1
lto = "fat"
panic = "abort"

Но в моей практике они лишь замедляли компиляцию и совершенно не улучшали производительность.

Инструменты для бенчмаркинга

Criterion — хороший инструмент для бенчмаркинга (расстановки контрольных точек), работа которого основана на статистических принципах. Я создал ещё один репозиторий, в котором сложил весь бенчмаркинговый код. Но, оказывается, весь материал нужно держать в одном и том же репозитории.

Здесь, в частности, следует отметить, что результаты бенчмаркинга получаются не очень стабильными. Я наблюдал разницу до ±10%, не внося никаких изменений в код. Если вы работаете на ноутбуке, то ситуация может быть даже хуже, поскольку из-за высокой температуры ЦП может работать с пониженной тактовой частотой.

Советую расставлять контрольные точки для нескольких разных параметров функции. Я в данном случае использую различные размерности. Если по всем размерностям результаты положительные, это обычно означает, что оптимизация проходит эффективно.

Метрики

Не забудьте с самого начала добавить несколько метрик. Многие баги и проблемы с производительностью можно выявить, если проверить метрики. Я напрямую использую AtomicU64, поскольку актуальные требования просты. Позже, возможно, переключусь на метрики Prometheus.

Обратите внимание: если переборщить с метриками/логированием/трассировкой, это также может сказаться на производительности. Поэтому добавляйте их с осторожностью.

Ресурсы

Расставляя контрольные точки, я заметил, что сквозная величина QPS (запросы в секунду) крайне нестабильна. Буквально за ночь я мог получить 15% в сторону улучшения или деградации, даже не перекомпилируя код. Затем я заметил, что ЦП работают не совсем вхолостую, поскольку я использую VSCode + анализатор Rust. Создаётся впечатление, будто они не слишком сильно нагружают процессор, но их работа очень серьёзно отражается на результатах бенчмаркинга. Хотя я и работаю с Intel Core i7-13700K, у которого 8 высокопроизводительных и 8 энергоэффективных ядер, программа у нас однопоточная.

При помощи taskset я привязываю процесс к конкретному ядру ЦП. Поэтому на процессах не скажется смешанное распределение ядер.

Обратите внимание: ядра процессора Intel Core 13/14 страдают из-за нестабильности по причине очень высокого напряжения. Я исправил эту проблему на уровне BIOS.

На облачные виртуальные машины температура процессора, пожалуй, не влияет, но у облачных провайдеров могут действовать собственные политики пропуска тактов процессора и оверселлинга.

Пошаговое улучшение

Начинаем с упрощённой реализации

В моём первом релизе алгоритм RaBitQ был основан на алгебраической библиотеке под названием nalgebra. Такой выбор был обусловлен, прежде всего, тем, что для получения ортогональной матрицы мне требовалось задействовать QR-разложение, а это ключевой шаг алгоритма RaBitQ. Кроме того, в зрелой библиотеке для линейной алгебры предоставляется множество полезных функций для операций над матрицами и векторами — и мне эта библиотека действительно облегчила реализацию алгоритма. Представьте себе, каково в Python без библиотеки numpy реализовывать алгоритм, в рамках которого происходит перемножение, проецирование и разложение матриц. Это кошмар.

Я рассчитывал на хорошую производительность, поскольку nalgebra оптимизирована для сценариев именно такого рода. Но по данным бенчмаркинга алгоритм работает гораздо медленнее, чем я ожидал. Думаю, он станет работать намного быстрее, если заново реализовать его с применением numpy :(

По данным профилирования, здесь выполняется множество вызовов f32::clone(). На них уходит примерно 33% всего времени или даже 44%, если сосредоточиться на функции query_one. Здесь вспомню о том, что можно заранее выделять память для нескольких векторов и повторно использовать её при итерации — это довольно распространённый фокус. Поэтому вместо (x - y).norm_squared() мне нужно заранее объявить другой вектор, в котором хранится результат (x - y), в итоге приобретающий вид x.sub_to(y, &mut z); z.norm_squared(). См. коммит 23f9aff.

Здесь, как и в большинстве алгебраических библиотек, матрица хранится в порядке развёртывания по столбцам — и это значит, что перебор по столбцам получается быстрее, чем по строкам. Это немного раздражает, поскольку перед итерацией матрицу приходится транспонировать, а не при всех операциях с перемножением векторов/матриц удаётся обнаружить на этапе компиляции ошибку несовпадения размерности (1 x dyn или dyn x 1).

Целевой ЦП

В RaBitQ для оценки приблизительного расстояния используется бинарное скалярное произведение, которое вычисляется так:

fn binary_dot_product(x: &[u64], y: &[u64]) -> u32 {
    assert_eq!(x.len(), y.len());
    let mut res = 0;
    for i in 0..x.len() {
        res += (x[i] & y[i]).count_ones();
    }
    res
}

Я думал, что здесь u64::count_ones() будет непосредственно использовать встроенные функции. Но оказывается, что я всё равно должен сам активировать фичу popcnt на этапе компиляции. Это можно сделать при помощи RUSTFLAGS="-C target-feature=+popcnt", но я предпочитаю вариант RUSTFLAGS="-C target-cpu=native", при котором активируются сразу все фичи, поддерживаемые действующим ЦП. Правда, при этом также исчезает возможность переносить двоичный формат, но на данном этапе это нормально. В следующих разделах также требуется включить env, чтобы задействовать возможности AVX2.

Возможности вашего ЦП вы можете проверить следующей командой:

rustc --print=cfg -C target-cpu=native | rg target_feature

ОКМД (одиночный поток команд, множественный поток данных) или SIMD

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

{
    y.sub_to(x, &mut residual);
    residual.norm_squared()
}

После профилирования я обнаружил, что f32::clone() по-прежнему здесь. Проверив исходный код nalgebra, я выяснил, что по какой-то причине в этой библиотеке много clone, почему — не знаю. Тогда я решил написать ОКМД вручную. К счастью, в hnswlib (популярная реализация алгоритма HNSW — поиск ближайших соседей с помощью маленького мира) такой функционал уже реализован.

В результате удаётся исключить f32::clone() из вычисления расстояния, тогда показатель QPS для SIFT (масштабно-инвариантного преобразования признаков) улучшается на 28%. Убедитесь сами, посмотрите коммит 5f82fcc.

Мой ЦП не поддерживает AVX512, поэтому я пользуюсь версией AVX2. Можете посмотреть Steam Hardware Stats, здесь в разделе «Other Settings» (Другие настройки) указана поддержка ОКМД. У 100% пользователей есть SSE3, у 94,61% есть AVX2 и только у 13,06% пользователей есть AVX512F. Разумеется, эта статистика не вполне объективна, так как в большинстве ЦП Intel поддерживается AVX512, но геймеры — это никак не все пользователи.

Наиболее полезное руководство по работе с ОКМД —Intel Intrinsics Guide. Лучше скачать весь сайт, так как онлайн с ним работать не очень удобно. Во встроенных функциях не забывайте проверять показатели "latency" (задержка) и "throughput" (пропускная способность), иначе код может работать медленнее, чем в обычной версии.

Другой хороший ресурс — Шпаргалка по встроенным функциям x86. Хорошо подойдёт таким новичкам, как я.

@ashvardanian написал пост о «масках», в котором рассказано, как решить проблему с хвостовыми элементами (требуется AVX512).

Добавьте это, чтобы код работал и на других платформах:

#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
{
    if is_x86_feature_detected!("avx2") {
        // Версия для AVX2 
    } else {
        // обычная версия
    }
}

В Rust есть несколько полезных пакетов, которые помогают ещё лучше писать cfg для ОКМД, но пока давайте не усложнять.

Ещё об ОКМД

ОКМД напоминает молоток, так что теперь мне нужно поискать в коде гвозди, которые можно было бы им позабивать.

• Можно переписать функцию binarize_vector под AVX2 как сделано в коммите f114fc1. Тогда показатель QPS для GIST улучшается на 32%.

Данная реализация по сравнению с исходной версией для C++, к тому же, не содержит ветвлений. При включённом opt-level=3 это может быть оптимизировано на уровне компилятора. См. ассемблерный код.

Как указал @andrewaylett, opt-level=3 позволяет оптимизировать вот это:

- let shift = if (i / 32) % 2 == 0 { 32 } else { 0 };
+ let shift = ((i >> 5) & 1) << 5;  // противоположная версия, так как я храню бинарник в u64, где действует противоположный порядок старшинства, нежели в версии на C++ 

@NovaX впервые указал, что этот код эквивалентен i & 32, а эта версия более удобочитаемая.

Скалярное квантование

Чтобы дополнительно вычистить f32::clone() из кода, я решил заменить ещё больше функций из nalgebra на реализации, которые сам выполнил бы вручную. Наиболее распространённые из таких функций — min и max. Вот как выглядит версия из nalgebra:

let lower_bound = residual.min();
let upper_bound = residual.max();
То же самое делается так:
fn min_max(vec: &[f32]) -> (f32, f32) {
    let mut min = f32::MAX;
    let mut max = f32::MIN;
    for v in vec.iter() {
        if *v < min {
            min = *v;
        }
        if *v > max {
            max = *v;
        }
    }
    (min, max)
}

Я привык работать с f32::min() и f32::max(), поскольку они удобные. Но при работе с невозрастающими (неснижающимися)  векторами if оказывается более производительным.

Чтобы не перебирать вектор по несколько раз в пределах цепочки функций и, следовательно, не вычислять скалярное квантование в рамках разных итераций:

let y_scaled = residual.add_scalar(-lower_bound) * one_over_delta + &self.rand_bias;
let y_quantized = y_scaled.map(|v| v.to_u8().expect("convert to u8 error"));
let scalar_sum = y_quantized.iter().fold(0u32, |acc, &v| acc + v as u32);

можно написать следующий цикл:

{
    let mut sum = 0u32;
    for i in 0..vec.len() {
        let q = ((vec[i] - lower_bound) * multiplier + bias[i]) as u8;
        quantized[i] = q;
        sum += q as u32;
    }
    sum
}

При скалярном квантовании можно быть уверенным, что f32 преобразуется в u8, поэтому воспользуемся as u8 вместо to_u8().unwrap().

Коммит af39c1c и коммит d2d51b0 позволили улучшить QPS для GIST на 31%.

Следующую часть также можно переписать с применением ОКМД, что позволяет улучшить QPS для GIST на 12%:

• min/max: коммит c97be68 и коммит e5a4af0

• скалярное квантование: коммит 28efe09

Я также попытался заменить tr_mul на ОКМД, то есть применить векторную проекцию. Оказывается, что nalgebra в данном случае использует BLAS, поэтому производительность остаётся без изменений.

Faer: ещё один алгебраический пакет

Изучая, как решить проблему с  f32::clone() , я обнаружил faer — ещё один алгебраический пакет для Rust. Он оптимизирован благодаря активному применению ОКМД и обеспечивает более высокую производительность при переборе по строкам и по столбцам. QR-разложение с ним также проходит гораздо быстрее, чем с nalgebra. В коммите 0411821 мы смогли ускорить этап обучения.

Кроме того, теперь (после коммита 0d969bd) я могу использовать эти векторы как обычный сегмент, без ColRef или RowRef в качестве обёрток.

Должен признать, что, если бы использовал faer с самого начала, то избавился бы от множества проблем. В любом случае, этот опыт был для меня очень плодотворным.

Двоичное скалярное произведение

Я думал, что popcnt уже решает проблему двоичного скалярного произведения, но на FlameGraph видно, что count_ones() берёт всего 7% от binary_dot_product. Хотя в AVX512 есть команда vpopcntq, я предпочёл бы воспользоваться симуляцией AVX2, так как она более распространена.

Вот хорошая справка по реализации popcnt для AVX2. В коммите edabd4a эта функция повторно реализована на Rust, благодаря чему удаётся улучшить QPS для GIST на 11%. Этот фокус работает лишь в случаях, когда у вектора более 256 измерений, то есть 256 разрядов для двоичного представления.

Атрибут Inline

Атрибутом #[inline] пользуйтесь с осторожностью. Если добавить его ко всем функциям с ОКМД, то  QPS для GIST улучшается на 5%.

Ввод/вывод

Здесь для начала дать небольшой контекст.

Современная реализация основана на алгоритме IVF, в котором используются k-средние для кластеризации векторов, и в памяти хранятся центроиды. Вектор запроса сравнивается только с кластерами, обладающими меньшим значением l2_squared_distance(query, centroid).

Есть параметр под названием n_probe, от которого зависит, сколько ближайших кластеров будет прозондировано. При большом значении n_probe полнота увеличится, но QPS снизится.

В RaBitQ двоичное скалярное произведение используется для приблизительной оценки расстояния. Если значение ниже порогового, то алгоритм заново выполнит ранжирование, используя исходный квадрат расстояния L2 и соответствующим образом обновит порог.

Ранее я пользовался вариантом slice::select_nth_unstable, который только выбирает n ближайших соседей, но не сортирует их по порядку. Если затрагивать кластеры, расположенные далеко от запроса, то возрастает вероятность повторного ранжирования, а из-за этого возрастает вычислительная нагрузка для расчёта квадрата расстояния L2. При пересортировке выбранных n кластеров QPS для GIST улучшался на 4%.

Другой приём — отсортировать векторы в каждом кластере по расстоянию, отделяющему их от центроид. В коммите ea13ebc также удалось улучшить QPS для GIST на 4%.

Есть ряд метаданных, помогающих оценить приблизительное расстояние для каждого вектора:

• factor_ip: f32

• factor_ppc: f32

• error: f32

• x_c_distance_square: f32

Раньше для их хранения я использовал 4 Vec<f32>, которые не очень удобны для работы с вводом/выводом, поскольку при расчётах для каждого из них требуется vector[i]. Когда я скомбинировал их в одну структуру struct в коммите bb440e3, показатель QPS для GIST улучшился на 2,5%. Срабатывает хорошо, поскольку это 4xf32, так что я могу использовать представление C напрямую:

#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
#[repr(C)]
struct Factor {
    factor_ip: f32,
    factor_ppc: f32,
    error_bound: f32,
    center_distance_square: f32,
}

К сожалению, faer не поддерживает векторы u64. Поэтому приходится сохранить двоичное представление вектора в Vec<Vec<u64>>. Когда я поменял его на Vec<u64> в коммите 48236b2, показатель QPS для GIST улучшился на 2%.

Константные дженерики

В версии на C++ мы используем шаблон, чтобы сгенерировать код для разных измерений. Такая возможность также доступна и в Rust. Я не буду её пробовать, поскольку перекомпиляция кода для другого набора измерений может быть применима только в конкретных случаях, например, в компании, где используется небольшое фиксированное количество измерений. В общедоступной библиотеке лучше представлять пользователям универсальное решение, чтобы им не приходилось перекомпилировать код самостоятельно.

Другие инструменты

Существует сборник рецептов по проверке границ, в котором на нескольких примерах проиллюстрировано, как избавиться от проверки границ в безопасном Rust.

Я пробовал PGO и BOLT, но не добился никаких улучшений.

Переход на jemalloc или mimalloc также никак не поспособствовал производительности.

Заключение

• ОКМД — отличный метод, если пользоваться им правильно

• Ввод/вывод также важен, особенно для больших датасетов

У актуальной версии на множестве данных GIST производительность такая же, как и у версии на C++. Там, где я наращиваю использование ОКМД, в версии на C++ задействуются константные дженерики.  

Ссылка

• Algorithmica / HPC