Недавно мне довелось работать над новым алгоритмом приближённого поиска ближайших соседей, который называется 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
скорость компиляции выше, но код может выполняться медленнее, чем на чистом Pythoncargo 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++ задействуются константные дженерики.