rabitq neon kernels (#5532)

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

tests/CMakeLists.txt

@@ -94,6 +94,10 @@ if(FAISS_OPT_LEVEL STREQUAL "dd")
         PROPERTIES COMPILE_OPTIONS "/arch:AVX512"
       )
     endif()
+  elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
+    list(APPEND FAISS_TEST_SRC
+      test_rabitq_simd_neon.cpp
+    )
   endif()
 endif()
 

tests/test_rabitq_simd.cpp

@@ -31,33 +31,8 @@ static std::vector<uint8_t> random_codes(size_t d, size_t qb, uint32_t seed) {
     return q;
 }
 
-static std::vector<uint8_t> random_bytes(size_t n, uint32_t seed) {
-    std::mt19937 rng(seed);
-    std::vector<uint8_t> v(n);
-    for (size_t i = 0; i < n; i++) {
-        v[i] = static_cast<uint8_t>(rng());
-    }
-    return v;
-}
-
-// 32-d chunks and chunk boundaries.
-static const std::vector<size_t> kDims = {
-        1,
-        8,
-        16,
-        31,
-        32,
-        33,
-        100,
-        128,
-        255,
-        256,
-        257,
-        384,
-        512,
-        768,
-        1024,
-        2048};
+using faiss_test::kDims;
+using faiss_test::random_bytes;
 
 // Note: scalar kernel's own correctness is covered end-to-end by
 // tests/test_rabitq.py. This target is x86-only (see BUCK).

tests/test_rabitq_simd_dd.cpp

@@ -29,37 +29,8 @@
 
 using faiss::SIMDLevel;
 
-namespace {
-
-std::vector<uint8_t> random_bytes(size_t n, uint32_t seed) {
-    std::mt19937 rng(seed);
-    std::vector<uint8_t> v(n);
-    for (size_t i = 0; i < n; i++) {
-        v[i] = static_cast<uint8_t>(rng());
-    }
-    return v;
-}
-
-// Widths on either side of every register boundary the kernels step over.
-const std::vector<size_t> kDims = {
-        1,
-        8,
-        16,
-        31,
-        32,
-        33,
-        100,
-        128,
-        255,
-        256,
-        257,
-        384,
-        512,
-        768,
-        1024,
-        2048};
-
-} // namespace
+using faiss_test::kDims;
+using faiss_test::random_bytes;
 
 TEST(RaBitQVpopcnt, BitwiseKernelsMatchScalarAcrossTails) {
     if (!faiss::SIMDConfig::is_simd_level_available(

tests/test_rabitq_simd_neon.cpp

@@ -0,0 +1,118 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <random>
+#include <vector>
+
+#include <faiss/utils/rabitq_simd.h>
+#include <faiss/utils/simd_levels.h>
+
+#include "test_rabitq_simd_util.h"
+
+using faiss::SIMDLevel;
+using faiss_test::kDims;
+using faiss_test::random_bytes;
+
+// This target is aarch64-only (see BUCK): it calls the ARM_NEON
+// specializations directly, and no other architecture compiles them.
+TEST(RaBitQNeon, BitwiseKernelsMatchScalar) {
+    if (!faiss::SIMDConfig::is_simd_level_available(SIMDLevel::ARM_NEON)) {
+        GTEST_SKIP() << "ARM_NEON is not available on this CPU";
+    }
+
+    for (size_t d : kDims) {
+        const size_t size = (d + 7) / 8;
+        const auto data = random_bytes(size, 19717 + d);
+        for (size_t qb = 1; qb <= 8; qb++) {
+            const auto query = random_bytes(size * qb, 51691 + d + qb);
+
+            const uint64_t expected_and =
+                    faiss::rabitq::bitwise_and_dot_product<SIMDLevel::NONE>(
+                            query.data(), data.data(), size, qb);
+            const uint64_t expected_xor =
+                    faiss::rabitq::bitwise_xor_dot_product<SIMDLevel::NONE>(
+                            query.data(), data.data(), size, qb);
+            const uint64_t expected_popcount =
+                    faiss::rabitq::popcount<SIMDLevel::NONE>(data.data(), size);
+
+            const uint64_t actual_and =
+                    faiss::rabitq::bitwise_and_dot_product<SIMDLevel::ARM_NEON>(
+                            query.data(), data.data(), size, qb);
+            const uint64_t actual_xor =
+                    faiss::rabitq::bitwise_xor_dot_product<SIMDLevel::ARM_NEON>(
+                            query.data(), data.data(), size, qb);
+            const uint64_t actual_popcount =
+                    faiss::rabitq::popcount<SIMDLevel::ARM_NEON>(
+                            data.data(), size);
+            const auto actual_fused =
+                    faiss::rabitq::bitwise_and_dot_product_with_popcount<
+                            SIMDLevel::ARM_NEON>(
+                            query.data(), data.data(), size, qb);
+
+            EXPECT_EQ(actual_and, expected_and) << "d=" << d << " qb=" << qb;
+            EXPECT_EQ(actual_xor, expected_xor) << "d=" << d << " qb=" << qb;
+            EXPECT_EQ(actual_popcount, expected_popcount)
+                    << "d=" << d << " qb=" << qb;
+            EXPECT_EQ(actual_fused.dot_product, expected_and)
+                    << "d=" << d << " qb=" << qb;
+            EXPECT_EQ(actual_fused.popcount, expected_popcount)
+                    << "d=" << d << " qb=" << qb;
+        }
+    }
+}
+
+TEST(RaBitQNeon, MultiBitInnerProductMatchesScalar) {
+    if (!faiss::SIMDConfig::is_simd_level_available(SIMDLevel::ARM_NEON)) {
+        GTEST_SKIP() << "ARM_NEON is not available on this CPU";
+    }
+
+    std::mt19937 rng(918273);
+    std::uniform_real_distribution<float> values(-2.0f, 2.0f);
+
+    for (size_t d : kDims) {
+        const auto sign_bits = random_bytes((d + 7) / 8, 13579 + d);
+        std::vector<float> query(d);
+        for (float& value : query) {
+            value = values(rng);
+        }
+        for (size_t ex_bits = 1; ex_bits <= 7; ex_bits++) {
+            // ip_scalar reads an 8-byte extraction window, so pad the code to
+            // match the production layout.
+            const auto extra_bits = random_bytes(
+                    (d * ex_bits + 7) / 8 + 7, 24680 + d + ex_bits);
+            const float cb = -static_cast<float>((1u << ex_bits) - 0.5f);
+            const float expected =
+                    faiss::rabitq::multibit::compute_inner_product<
+                            SIMDLevel::NONE>(
+                            sign_bits.data(),
+                            extra_bits.data(),
+                            query.data(),
+                            d,
+                            ex_bits,
+                            cb);
+            const float actual = faiss::rabitq::multibit::compute_inner_product<
+                    SIMDLevel::ARM_NEON>(
+                    sign_bits.data(),
+                    extra_bits.data(),
+                    query.data(),
+                    d,
+                    ex_bits,
+                    cb);
+
+            // The NEON kernel accumulates in eight lanes, so it does not
+            // reproduce the scalar sum bit for bit.
+            const float tolerance = 2e-5f * std::max(1.0f, std::abs(expected));
+            EXPECT_NEAR(actual, expected, tolerance)
+                    << "d=" << d << " ex_bits=" << ex_bits;
+        }
+    }
+}

tests/test_rabitq_simd_util.h

@@ -5,9 +5,11 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// Shared by test_rabitq_simd.cpp and test_rabitq_simd_dd.cpp. The AVX512
-// kernels live in a dynamic dispatch build only, so the tests that call them
-// sit in the second file, but both compare against the same scalar reference.
+// Shared by test_rabitq_simd.cpp, test_rabitq_simd_dd.cpp and
+// test_rabitq_simd_neon.cpp. The AVX512 kernels live in a dynamic dispatch
+// build only, and the NEON kernels live in an aarch64 build only, so each set
+// of tests sits in its own file. They all compare against the same scalar
+// reference.
 
 #pragma once
 
@@ -15,6 +17,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <random>
 #include <vector>
 
 #include <faiss/utils/rabitq_simd.h>
@@ -24,6 +27,34 @@ namespace faiss_test {
 
 using faiss::SIMDLevel;
 
+// Widths on either side of every register boundary the kernels step over.
+inline const std::vector<size_t> kDims = {
+        1,
+        8,
+        16,
+        31,
+        32,
+        33,
+        100,
+        128,
+        255,
+        256,
+        257,
+        384,
+        512,
+        768,
+        1024,
+        2048};
+
+inline std::vector<uint8_t> random_bytes(size_t n, uint32_t seed) {
+    std::mt19937 rng(seed);
+    std::vector<uint8_t> v(n);
+    for (size_t i = 0; i < n; i++) {
+        v[i] = static_cast<uint8_t>(rng());
+    }
+    return v;
+}
+
 template <SIMDLevel SL>
 inline void check_quantization_matches_scalar() {
     constexpr size_t d = 257;