最近在研究名为RaBitQ的新近似最近邻搜索算法，作者提供了C++实现运行速度较快。作者尝试用 Rust 重写（https://github.com/kemingy/ra...），但发现实现比原版本慢很多，以下是逐步提升性能的过程：

• 环境准备：

  • 数据集：重要的是有合理的数据集，使用论文中 sift_dim128_1m_l2和gist_dim960_1m_l2数据集，可从这里下载，格式为fvecs/ivecs，可从gist获取读写脚本。
  • 分析工具：使用samply分析 Rust 代码，与Firefox Profiler集成，可上传到云端分享分析结果，如GIST 上 C++版本的分析示例，常用 FlameGraph 和 CallTree 视图，需授予性能事件权限并增加mlock限制，GodBolt编译器浏览器可比较 C++和 Rust 的汇编函数代码。
  • Cargo 配置文件：在Cargo.toml中添加[profile.perf]配置以在发布版本中包含调试信息，cargo build编译速度快但代码可能比纯 Python 慢，cargo build --release运行快但编译时间长，基准测试需使用opt-level = 3，尝试使用codegen-units = 1、lto = "fat"、panic = "abort"等设置但效果不佳。
  • 基准测试工具：Criterion是良好的统计驱动基准测试工具，创建repo存储相关基准测试代码，注意基准测试结果不稳定，建议用不同参数进行测试。
  • 指标：从开始就添加指标，便于发现问题，使用AtomicU64，也可后期切换为Prometheus metrics，注意过多指标可能影响性能。
  • 资源：基准测试中发现端到端 QPS 极不稳定，因有 VSCode + Rust Analyzer 导致 CPU 不完全空闲影响基准测试结果，使用taskset将进程绑定到特定 CPU 避免混合核心调度影响，Intel Core 13 代/14 代 CPU 因电压高受不稳定问题影响，可在 BIOS 中修复，云虚拟机可能不受 CPU 温度影响但有自己的 CPU 节流和超额预订策略。

• 逐步提升性能：

  • 初始实现：基于nalgebra库实现 RaBitQ 算法，以为性能良好但实际比预期慢，分析发现有很多f32::clone()调用，可预分配内存重用，该库按列主序存储矩阵，迭代时需转置，部分向量/矩阵乘法在编译时不能检测维度不匹配错误。
  • CPU 目标：RaBitQ 用二进制点积距离估计近似距离，u64::count_ones()需启用popcnt特性，可通过RUSTFLAGS="-C target-feature=+popcnt"或RUSTFLAGS="-C target-cpu=native"实现，后者使二进制不可移植，可通过rustc --print=cfg -C target-cpu=native | rg target_feature检查 CPU 特性。
  • SIMD：最近邻搜索的关键函数是距离函数，用 L2 平方距离避免开方计算，分析发现仍有f32::clone()，手动编写 SIMD 消除该调用，提升 SIFT 的 QPS 约 28%，需根据 CPU 支持的 SIMD 版本编写代码，如不支持 AVX512 则使用 AVX2，可参考Intel Intrinsics Guide和x86 Intrinsics Cheat Sheet，用#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]进行平台适配。
  • 更多 SIMD：用 AVX2 重写binarize_vector函数提升 GIST 的 QPS 约 32%，启用opt-level=3可由编译器优化，还可重写min、max函数和标量量化部分代码以提升 GIST 的 QPS 约 31%，部分部分也可改写为 SIMD 提升 QPS 约 12%，尝试用 SIMD 替换tr_mul但性能不变。
  • 标量量化：用手动实现替换更多nalgebra函数以消除f32::clone()，如重写min、max函数和标量量化部分代码，提升 GIST 的 QPS 约 31%，部分部分也可改写为 SIMD 提升 QPS 约 12%。
  • 另一个代数库：faer：发现faer代数库，优化了很多 SIMD 并提供更好的行列迭代性能，QR 分解比nalgebra快，使用该库后训练部分更快，可直接使用向量切片无需ColRef或RowRef包装。
  • 二进制点积：以为popcnt已解决二进制点积问题，但分析发现count_ones()仅占binary_dot_product的 7%，用 AVX2 模拟实现提升 GIST 的 QPS 约 11%，该技巧仅在向量维度大于 256 时有效。
  • 内联：谨慎使用#[inline]属性，添加到所有 SIMD 函数可提升 GIST 的 QPS 约 5%。
  • IO：当前实现基于 IVF 算法，用slice::select_nth_unstable选择最近邻但未排序，重新排序提升 GIST 的 QPS 约 4%，对每个簇的向量按到质心的距离排序也提升 GIST 的 QPS 约 4%，将存储向量相关元数据的 4 个Vec<f32>合并为一个struct提升 GIST 的 QPS 约 2.5%，将存储向量二进制表示的Vec<Vec<u64>>改为Vec<u64>提升 GIST 的 QPS 约 2%。
  • 常量泛型：C++版本用模板为不同维度生成代码，Rust 也有此特性但未尝试，对于公共库提供通用解决方案更好。
  • 其他工具：bounds-check-cookbook提供消除安全 Rust 中边界检查的示例，尝试PGO和BOLT未提升性能，切换到jemalloc或mimalloc也未提升性能。
• 结论：SIMD 恰当使用效果好，IO 也很重要，当前性能在 GIST 数据集上与 C++版本相同，Rust 版本使用更多 SIMD 而 C++版本使用常量泛型。
• 参考文献：Algorithmica / HPC