Partial Optimization for IVFPQ#4893
Partial Optimization for IVFPQ#4893littleniuer wants to merge 2 commits intofacebookresearch:mainfrom
Conversation
|
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. |
Summary
This PR introduces ARM NEON optimizations for IVFPQ index, focusing on 8-bit lookup table operations and distance computations.
Changes
faiss/sra_krl/for clear organizationDesign Decisions
faiss/sra_krl/) to keep the file structure clean and facilitate initial code review__aarch64__macro, no impact on x86 buildsTesting
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.