Skip to content

Partial Optimization for IVFPQ#4893

Open
littleniuer wants to merge 2 commits intofacebookresearch:mainfrom
littleniuer:main
Open

Partial Optimization for IVFPQ#4893
littleniuer wants to merge 2 commits intofacebookresearch:mainfrom
littleniuer:main

Conversation

@littleniuer
Copy link

Summary

This PR introduces ARM NEON optimizations for IVFPQ index, focusing on 8-bit lookup table operations and distance computations.

Changes

  • Optimized 8-bit LUT (lookup table) construction and distance calculation
  • Implemented algorithmic improvements combined with NEON intrinsics
  • All core optimization code is placed under faiss/sra_krl/ for clear organization

Design Decisions

  • Code is organized in a separate directory (faiss/sra_krl/) to keep the file structure clean and facilitate initial code review
  • Guarded by __aarch64__ macro, no impact on x86 builds
  • Functionally equivalent to the original implementation (no changes to index format or search results)

Testing

  • Passed FAISS built-in unit tests on AArch64 platform
  • Benchmark results show noticeable performance improvements on ARM servers

Notes

We understand that naming conventions and NEON code placement may need adjustments to align with the ongoing SIMD restructuring. We are happy to collaborate with maintainers to refine the code structure as needed.

If this contribution is well-received, we have additional optimizations for HNSW, Refine, and FastScan ready for follow-up PRs.

@littleniuer
Copy link
Author

I noticed two test failures (test_precomputed_tables and test_hnsw_2level_mixed_search) on aarch64 with my changes. The root cause is a minor floating-point precision difference (~3.8e-06 max absolute error) introduced by the NEON-optimized PQ distance table computation. Specifically, the optimized kernel uses a transposed data layout and processes 16 centroids in parallel with a dimension-wise accumulation order, which differs from the original per-centroid accumulation. Since float32 addition is not associative, this leads to ~1 ULP differences in the distance values.

Both tests currently use bitwise-exact comparison (assert_array_equal / np.all(==)) for distance results. I've verified that relaxing the distance comparison to np.testing.assert_allclose(atol=1e-5, rtol=1e-6) resolves the issue. Note that only the distance (D) comparisons are relaxed — all index (I) comparisons remain strictly equal, confirming that the search results themselves are not affected.

All other 1087 tests pass without any modifications.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant