算法性能优化策略：RaBitQ 的演进之旅
在近期的研究中，我专注于一种新颖的近似最近邻搜索算法——RaBitQ。尽管原始的 C++ 实现已经展现出了卓越的运行效率，我仍试图将其用 Rust 语言重新构建，以期进一步提升性能。然而，初步尝试表明我的 Rust 实现在速度上远逊于原始版本。以下是我如何系统性地提升算法性能的详细过程。
环境准备
数据集选择
为了确保实验的准确性，选择一个合适的数据集至关重要。鉴于相关论文已在 sift_dim128_1m_l2 和 gist_dim960_1m_l2 数据集上展示了结果，我决定采用这两个数据集进行基准测试。这些数据集的维度分别为 128 和 960，包含 1,000,000 个向量，足以满足基准测试的需求。这些数据集可以通过 FTP 协议从特定网站下载，尽管该网站尚未部署 TLS 加密。
数据集采用 fvecs/ivecs 格式，这是一种广泛使用的向量存储格式，其结构如下：
| dim (4 bytes) | vector (4 * dim bytes) |
| dim (4 bytes) | vector (4 * dim bytes) |
...
| dim (4 bytes) | vector (4 * dim bytes) |
相关的读写脚本可以从我的 gist 仓库中获取。
分析工具
为了深入分析 Rust 代码的性能，我采用了 samply 工具，该工具与 Firefox Profiler 紧密集成，便于性能分析结果的云端共享。例如，C++ 版本的性能分析结果已上传至 GIST。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.toml 文件中添加适当的配置，可以在发布构建中包含调试信息，如下所示：
[profile.perf]
inherits = "release"
debug = true
codegen-units = 16
编译成本与运行时性能之间的平衡对于分析用户体验至关重要。
使用 cargo build命令可以加快编译速度，但可能导致运行速度低于预期。而 cargo build --release命令虽然编译时间较长，但能提供更快的运行速度。
对于基准测试，我们通常选择 opt-level = 3 以获得最佳性能。
我还尝试了一些建议的设置：
codegen-units = 1
lto = "fat"
panic = "abort"
然而，这些设置在我的案例中仅降低了编译速度，并未显著提升性能。
基准测试工具
Criterion 是一个优秀的统计驱动基准测试工具。我创建了一个专门的仓库来存储所有相关的基准测试代码。值得注意的是，基准测试结果可能会有较大的波动，例如在没有任何代码修改的情况下，我观察到了 ±10% 的性能差异。在笔记本电脑上进行测试时，由于 CPU 可能因过热而降频，情况可能更加复杂。
我建议使用不同的参数对函数进行基准测试。例如，我使用了不同的向量维度进行测试。如果所有维度的测试结果都显示出性能提升，通常意味着改进是有效的。
性能指标
从实验开始就添加一些性能指标是非常重要的。许多错误和性能问题都可以通过监控这些指标来发现。我直接使用了 AtomicU64，因为当前的需求相对简单。未来可能会考虑切换到 Prometheus 进行更复杂的指标监控。
需要注意的是，过多的指标、日志或追踪可能会对性能产生负面影响，因此在添加时应谨慎。
资源管理
在进行基准测试时，我注意到端到端的 QPS 表现出极高的不稳定性。有时在不重新编译代码的情况下，第二天早上就可能出现 15% 的性能提升或下降。进一步分析发现，CPU 并非完全空闲，因为我使用的 VSCode + Rust Analyzer 虽然看似 CPU 占用不高，但实际上对基准测试结果产生了显著影响。尽管我使用的是 Intel Core i7-13700K 处理器，拥有 8 个性能核心和 8 个高效核心，且程序是单线程运行的。
通过使用 taskset 命令将进程绑定到特定的 CPU，可以避免混合核心调度带来的影响。
需要注意的是，Intel Core 13th/14th 代 CPU 可能会受到高电压不稳定问题的影响。我已经在 BIOS 中解决了这个问题。
云 VM 可能不会受到 CPU 温度的影响，但云服务提供商可能会有自己的 CPU 限制和超额预订政策。
性能优化策略
基础实现
我的初步实现基于名为 nalgebra 的代数库，该库为 RaBitQ 算法提供了关键的 QR 分解功能。此外，成熟的线性代数库提供了许多有用的函数来操作矩阵和向量，这大大简化了算法的实现。可以想象，在没有 numpy 的情况下用 Python 实现涉及矩阵乘法、投影和分解的算法是多么困难。
尽管预期性能应该不错，因为 nalgebra 对此类场景进行了优化，但基准测试结果显示性能远低于预期。我推测如果使用 numpy 重新实现可能会有显著的性能提升。
通过分析，我发现 f32::clone() 调用非常频繁，占据了总时间的大约 33%，或者在关注 query_one 函数时，这一比例高达 44%。这提示我可以为某些向量预分配内存并在迭代中重用它，这是一个非常常见的优化技巧。因此，而不是使用 (x - y).norm_squared()，我需要预先声明另一个向量来存储 (x - y) 的结果，最终实现为 x.sub_to(y, &mut z); z.norm_squared()。参见提交 23f9aff。
与大多数代数库一样，nalgebra 以列主序方式存储矩阵，这意味着迭代列可能比迭代行更快。这有点烦人，因为我需要在迭代之前转置矩阵，而且并非所有的向量/矩阵乘法都能在编译期间检测到维度不匹配错误（1 x dyn 或 dyn x 1）。
CPU 特性优化
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"，它启用了当前 CPU 支持的所有 CPU 特性，但也使二进制文件不可移植，目前这并不是问题。接下来的部分也需要这个环境变量来启用 AVX2 特性。
可以使用以下命令来检查 CPU 特性：
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，原因我不清楚。我决定手动编写 SIMD 代码。幸运的是，hnswlib（一个流行的 HNSW 实现）已经实现了这一点。
这消除了距离计算中的 f32::clone()，并提高了 SIFT 的 QPS **28%**。查看提交 5f82fcc。
我的 CPU 不支持 AVX512，所以我使用 AVX2 版本。可以通过查看 Steam 硬件统计数据来了解 SIMD 支持情况。100% 用户有 SSE3 支持，94.61% 用户有 AVX2 支持，只有 13.06% 用户有 AVX512F 支持。当然，这个统计数据是有偏见的，大多数云 Intel CPUs 都支持 AVX512，游戏玩家不能代表所有用户。
要使用 SIMD，最有用指南是 Intel Intrinsics Guide。最好下载网站，因为在线体验不好。记住检查 " latency" 和 " throughput" 的内部函数，否则你的代码可能比正常版本慢。
另一个资源是 x86 Intrinsics Cheat Sheet。这对新手如我很有用。
@ashvardanian 有一篇关于 "mask load" 的文章，解决了尾元素问题（需要 AVX512）。
要使代码在其他平台上工作：
#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
{
if is_x86_feature_detected!("avx2") {
// AVX2 版本
} else {
// 正常版本
}
}
有一些有用的 crates 用于编写更好的 cfg 用于 SIMD，现在让我们保持简单。
更多 SIMD 优化
SIMD 就像一把锤子，现在我需要在代码中找到更多的钉子。
使用 AVX2 重写 binarize_vector函数，在提交 f114fc1 中提高了 GIST 的 QPS **32%**。
与原始的 C++ 版本相比，这个实现也是无分支的。
- let shift = if (i / 32) % 2 == 0 { 32 } else { 0 };
+ let shift = ((i >> 5) & 1) << 5;
看看汇编的差异。
嗯，去分支化并没有使整体性能提高太多，因为 binarize_vector 函数只对每个查询调用一次。但这是个好的学习机会。
标量量化
为了消除代码中更多的 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 提高了 GIST 的 QPS **31%**。
以下部分也可以使用 SIMD 重写，这提高了 GIST 的 QPS **12%**：
min/max: 提交 c97be68 & 提交 e5a4af0 标量量化：提交 28efe09
我也尝试用 SIMD 替换 tr_mul，这是一个向量投影。结果发现 nalgebra 在这里使用 BLAS，所以性能保持不变。
另一个代数库：faer
在研究 f32::clone() 问题时，我发现了另一个 Rust 代数库，称为 faer。它用大量 SIMD 优化，并提供了更好的行/列迭代性能。QR 分解也比 nalgebra 快得多。这个提交 0411821 使训练部分更快。
此外，我现在可以在提交 0d969bd 之后将这些向量用作普通切片，而不需要 ColRef 或 RowRef 包装器。
我必须承认，如果我从一开始就使用 faer，我可以避免很多麻烦。无论如何，我从这次经历中学到了很多。
二进制点积优化
我以为 popcnt 已经解决了二进制点积问题，但 FlameGraph 显示 count_ones() 只占 binary_dot_product 的 7%。尽管 AVX512 有 vpopcntq 指令，我更倾向于使用更常见的 AVX2 模拟。
这是一个很好的参考，用于 AVX2 的 popcnt 实现。提交 edabd4a 在 Rust 中重新实现这一点，提高了 GIST 的 QPS **11%**。
内联优化
#[inline] 属性应该谨慎使用。为所有 SIMD 函数添加此属性，提高了 GIST 的 QPS **5%**。
IO 优化
我需要在这里添加一些背景信息。
当前实现基于 IVF 算法，它使用 k-means 对向量进行聚类，并将中心点存储在内存中。查询向量仅与较小的 l2_squared_distance(query, centroid) 的聚类进行比较。
有一个参数称为 n_probe，它控制将探测多少最近聚类。较大的 n_probe 将增加召回率，但会降低 QPS。
RaBitQ 使用二进制点积来估计近似距离。如果它小于阈值，它将使用原始的 L2 平方距离进行重新排名，并相应更新阈值。
以前，我使用 slice::select_nth_unstable，它只选择 n-nearest 但不会按顺序对它们进行排序。遍历远离查询的聚类将增加重新排名比率，这需要更多的 L2 平方距离计算。重新排序选定的 n-th 聚类提高了 GIST 的 QPS **4%**。
另一个技巧是对每个聚类中的向量按其到中心点的距离进行排序，这个提交 ea13ebc 也提高了 GIST 的 QPS **4%**。
有一些元数据用于估计每个向量的近似距离：
factor_ip：f32 factor_ppc：f32 error：f32 x_c_distance_square：f32
以前我使用 4 个 Vec<f32> 来存储它们，这对 IO 不友好，因为计算需要每个的 vector[i]。通过在提交 bb440e3 中将它们组合成一个 struct，QPS 提高了 **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>> 中。通过在提交 48236b2 中将其更改为 Vec<u64>，QPS 提高了 **2%**。
常量泛型
C++ 版本使用模板为不同维度生成代码。这个特性在 Rust 中也有。我没有尝试，因为为不同维度重新编译代码可能只适用于特定用例，比如在只有几个固定维度的公司内部。对于公共库，最好提供通用解决方案，以便用户不必自己重新编译。
其他工具
有一个 bounds-check-cookbook，提供了几个如何在安全 Rust 中消除边界检查的示例。
我尝试了 PGO 和 BOLT，但没有得到任何改进。
切换到 jemalloc 或 mimalloc 也没有改善性能。
结论
当正确使用时，SIMD 是非常棒的 IO 也很重要，特别是对于大型数据集 在选择正确的库时要做更多的研究
参考
Algorithmica / HPC
`
