Convert rabitq_simd.h to runtime SIMD dispatch (facebookresearch#4912)

Summary:
Pull Request resolved: facebookresearch#4912

Convert RaBitQ SIMD primitives from compile-time #ifdef cascades to the
DD template specialization pattern.

- Replace inline SIMD functions in rabitq_simd.h with template<SIMDLevel>
  declarations and NONE (scalar) specializations
- Create rabitq_avx2.cpp with AVX2+SSE4.1 specializations for
  bitwise_and_dot_product, bitwise_xor_dot_product, popcount, and
  compute_inner_product (including BMI2 bitplane kernel)
- Create rabitq_avx512.cpp with AVX512 specializations (falls back to
  AVX2 bitplane kernel for compute_inner_product with ex_bits >= 2)
- Add AVX512_SPR forwarding to AVX512 for all four functions
- Create rabitq_neon.cpp forwarding ARM_NEON/ARM_SVE to NONE
- Wrap calls in RaBitQuantizer.cpp and RaBitQUtils.cpp with
  with_simd_level dispatch
- Add -mbmi2 to AVX2 and AVX512 SIMD flags (required by _pext_u64 in
  bitplane kernel; BMI2 is present on all Haswell+ CPUs)
- Register new SIMD files in xplat.bzl and CMakeLists.txt

Reviewed By: mdouze

Differential Revision: D96298743

fbshipit-source-id: 1941bf099664ba2d3d1c1245d5c06d1775614955

Author: algoriddle

faiss/CMakeLists.txt

@@ -15,20 +15,23 @@ set(FAISS_SIMD_AVX2_SRC
   impl/approx_topk/avx2.cpp
   utils/simd_impl/distances_avx2.cpp
   utils/distances_fused/simdlib_based.cpp
+  utils/simd_impl/rabitq_avx2.cpp
 )
 set(FAISS_SIMD_AVX512_SRC
   impl/fast_scan/impl-avx512.cpp
   impl/pq_code_distance/pq_code_distance-avx512.cpp
   impl/scalar_quantizer/sq-avx512.cpp
   utils/simd_impl/distances_avx512.cpp
   utils/distances_fused/avx512.cpp
+  utils/simd_impl/rabitq_avx512.cpp
 )
 set(FAISS_SIMD_NEON_SRC
   impl/fast_scan/impl-neon.cpp
   impl/scalar_quantizer/sq-neon.cpp
   impl/approx_topk/neon.cpp
   utils/simd_impl/distances_aarch64.cpp
   utils/distances_fused/simdlib_based_neon.cpp
+  utils/simd_impl/rabitq_neon.cpp
 )
 set(FAISS_SIMD_SVE_SRC
   impl/pq_code_distance/pq_code_distance-sve.cpp

faiss/impl/RaBitQUtils.cpp

@@ -8,6 +8,7 @@
 #include <faiss/impl/RaBitQUtils.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/rabitq_simd.h>
 #include <algorithm>
@@ -306,6 +307,9 @@ size_t compute_per_vector_storage_size(size_t nb_bits, size_t d) {
     }
 }
 
+// Non-template wrapper with dynamic dispatch (one dispatch per call).
+// The hot path in RaBitQuantizer dispatches once at distance computer
+// construction, so per-vector dispatch only affects this utility path.
 float compute_full_multibit_distance(
         const uint8_t* sign_bits,
         const uint8_t* ex_code,
@@ -315,18 +319,18 @@ float compute_full_multibit_distance(
         size_t d,
         size_t ex_bits,
         MetricType metric_type) {
-    const float cb = -(static_cast<float>(1 << ex_bits) - 0.5f);
-
-    float ex_ip = rabitq::multibit::compute_inner_product(
-            sign_bits, ex_code, rotated_q, d, ex_bits, cb);
-
-    float dist = qr_base + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;
-
-    if (metric_type == MetricType::METRIC_L2) {
-        dist = std::max(0.0f, dist);
-    }
-
-    return dist;
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0>(
+            [&]<SIMDLevel SL>() {
+                return compute_full_multibit_distance<SL>(
+                        sign_bits,
+                        ex_code,
+                        ex_fac,
+                        rotated_q,
+                        qr_base,
+                        d,
+                        ex_bits,
+                        metric_type);
+            });
 }
 
 void populate_block_aux_from_flat_storage(

faiss/impl/RaBitQUtils.h

@@ -10,6 +10,9 @@
 #include <faiss/MetricType.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/AlignedTable.h>
+#include <faiss/utils/rabitq_simd.h>
+#include <faiss/utils/simd_levels.h>
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
@@ -337,6 +340,33 @@ float compute_full_multibit_distance(
         size_t ex_bits,
         MetricType metric_type);
 
+// SIMDLevel-templatized version — avoids per-call dynamic dispatch.
+// Inline so it can be used from templatized distance computers without
+// needing explicit instantiations in per-SIMD TUs.
+template <SIMDLevel SL>
+inline float compute_full_multibit_distance(
+        const uint8_t* sign_bits,
+        const uint8_t* ex_code,
+        const ExtraBitsFactors& ex_fac,
+        const float* rotated_q,
+        float qr_base,
+        size_t d,
+        size_t ex_bits,
+        MetricType metric_type) {
+    const float cb = -(static_cast<float>(1 << ex_bits) - 0.5f);
+
+    float ex_ip = rabitq::multibit::compute_inner_product<SL>(
+            sign_bits, ex_code, rotated_q, d, ex_bits, cb);
+
+    float dist = qr_base + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;
+
+    if (metric_type == MetricType::METRIC_L2) {
+        dist = std::max(0.0f, dist);
+    }
+
+    return dist;
+}
+
 /** Compute pointer to a vector's auxiliary data within block layout. */
 template <typename T>
 inline T* get_block_aux_ptr(