Summary:
This diff adds native ARM NEON kernels for the RaBitQ bitwise popcount
operations and for the multi-bit inner product. The kernels replace the scalar
fallbacks on ARM. They handle every supported bit width and every unaligned
dimension.
- `tests/test_rabitq_simd_neon.cpp` holds the correctness test, and
`tests/BUCK` holds its `cpp_unittest` target. The target carries
`target_compatible_with` on `ovr_config//cpu:arm64`, which mirrors the
x86_64 gate on `test_rabitq_simd`.
- `tests/test_rabitq_simd_util.h` holds the shared `kDims` list and the shared
`random_bytes` helper. `test_rabitq_simd.cpp` and `test_rabitq_simd_dd.cpp`
now use them too.
- The three kernels that index a fixed array of eight accumulators assert
`qb <= kMaxQueryBits`. `RaBitQuantizer::set_query` caps `qb` at 8, but the
kernels are public templates and carry no other bound. The scalar kernel and
the AVX-512 kernel need no such assertion, because they shift inside the loop
and hold one accumulator.
The multi-bit test compares with a relative tolerance of `2e-5`. The NEON kernel
accumulates in eight lanes, so it does not reproduce the scalar sum bit for bit.
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5532
Test Plan:
`buck2 test fbcode//faiss/tests:test_ivf_index fbcode//faiss/tests:test_custom_result_handler fbcode//faiss/tests:test_hnsw fbcode//faiss/tests:test_rabitq_simd_neon fbcode//faiss/tests:test_params_override fbcode//faiss/tests:test_hnsw_rabitq fbcode//faiss/tests:test_rabitq fbcode//faiss/tests:test_search_params fbcode//faiss/tests:test_io_corrupted fbcode//faiss/tests:test_rabitq_utils`, run at the top of the stack on aarch64: 173 pass, 0 fail, 3 skip.
`buck2 build fbcode//faiss:faiss` succeeds in the static aarch64 build and with `-c faiss.dynamic_dispatch=true`.
The host holds Neoverse V2 cores with ASIMDDP, so `test_rabitq_simd_neon` runs
its kernels natively rather than skipping.
Reviewed By: alibeklfc
Differential Revision: D121846717
Pulled By: mnorris11
fbshipit-source-id: de079613a52ced326045830fb799dfc9b169fe2e
lyang24
0db2f30886