Add Sapphire Rapids optimizations for ScalarQuantizer (L2, IP)

Adds an AVX512_SPR specialization path for ScalarQuantizer that uses
Sapphire Rapids-specific instructions for byte-code distance computation
on QT_8bit_direct and QT_8bit_direct_signed.

Inner product (8-bit codes):

  Replaces the AVX512 path that processes 16 bytes per iteration via
  cvtepu8_epi32 + mullo_epi32 with a VNNI loop that processes 64 bytes
  per iteration using _mm512_dpbusd_epi32. VNNI computes unsigned*signed
  dot products, so the standard bias trick is used to bridge
  unsigned*unsigned: subtract 128 from code2, run dpbusd, then add the
  128 * sum(code1) correction. A scalar tail handles d % 64.

  For QT_8bit_direct_signed (storage = value + 128), the same VNNI loop
  runs and an additional closed-form correction is applied:
      (a-128) * (b-128) = a*b - 128*(a+b) + 16384
  sum(a) and sum(b) are accumulated cheaply via _mm512_sad_epu8 (one
  PSADBW per 64-byte iteration).

L2 (8-bit codes):

  Replaces the 16-bytes-per-iter cvtepu8_epi32 + sub + mullo_epi32 path
  with a 16-bit pipeline: load 64 bytes, zero-extend to 16-bit lanes via
  _mm512_cvtepu8_epi16, subtract in 16-bit, square-and-accumulate to
  32-bit with _mm512_madd_epi16. Squared differences of two uint8_t
  values fit in 16 bits (max 255^2 = 65025), so the widened
  representation is safe. Falls through to a 32-byte step and a scalar
  tail for arbitrary d. The same kernel is bit-exact for the signed
  variant: (a - 128) - (b - 128) == a - b, so no correction is needed.

Signed-off-by: Mulugeta Mammo <[email protected]>

Author: mulugetam

faiss/CMakeLists.txt

@@ -34,6 +34,9 @@ set(FAISS_SIMD_AVX512_SRC
   utils/simd_impl/rabitq_avx512.cpp
   utils/simd_impl/super_kmeans_kernels_avx512.cpp
 )
+set(FAISS_SIMD_AVX512_SPR_SRC
+  impl/scalar_quantizer/sq-avx512-spr.cpp
+)
 set(FAISS_SIMD_NEON_SRC
   impl/fast_scan/impl-neon.cpp
   impl/scalar_quantizer/sq-neon.cpp
@@ -61,7 +64,7 @@ set(FAISS_SIMD_RVV_SRC
 )
 # Select SIMD sources based on target architecture
 if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64|amd64|AMD64)")
-  set(FAISS_SIMD_SRC ${FAISS_SIMD_AVX2_SRC} ${FAISS_SIMD_AVX512_SRC})
+  set(FAISS_SIMD_SRC ${FAISS_SIMD_AVX2_SRC} ${FAISS_SIMD_AVX512_SRC} ${FAISS_SIMD_AVX512_SPR_SRC})
 elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
   set(FAISS_SIMD_SRC ${FAISS_SIMD_NEON_SRC} ${FAISS_SIMD_SVE_SRC})
 elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(riscv64|riscv)")
@@ -461,13 +464,13 @@ endif()
 if(NOT WIN32)
   # Architecture mode to support AVX512 extensions available since Intel(R) Sapphire Rapids.
   # Ref: https://networkbuilders.intel.com/solutionslibrary/intel-avx-512-fp16-instruction-set-for-intel-xeon-processor-based-products-technology-guide
-  target_compile_options(faiss_avx512_spr PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vpopcntdq -mpopcnt -mavx512fp16 -mavx512bf16 ${FAISS_BMI2_FLAGS}>)
+  target_compile_options(faiss_avx512_spr PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vnni -mavx512vpopcntdq -mpopcnt -mavx512fp16 -mavx512bf16 ${FAISS_BMI2_FLAGS}>)
 else()
   target_compile_options(faiss_avx512_spr PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/arch:AVX512>)
   # we need bigobj for the swig wrapper
   add_compile_options(/bigobj)
 endif()
-target_sources(faiss_avx512_spr PRIVATE ${FAISS_SIMD_AVX2_SRC} ${FAISS_SIMD_AVX512_SRC})
+target_sources(faiss_avx512_spr PRIVATE ${FAISS_SIMD_AVX2_SRC} ${FAISS_SIMD_AVX512_SRC} ${FAISS_SIMD_AVX512_SPR_SRC})
 target_compile_definitions(faiss_avx512_spr PRIVATE COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512 COMPILE_SIMD_AVX512_SPR )
 
 add_library(faiss_sve ${FAISS_SRC})
@@ -525,6 +528,11 @@ if(FAISS_OPT_LEVEL STREQUAL "dd")
         PROPERTIES COMPILE_OPTIONS
           "-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mfma;-mf16c;-mpopcnt"
       )
+      set_source_files_properties(${FAISS_SIMD_AVX512_SPR_SRC}
+        TARGET_DIRECTORY faiss
+        PROPERTIES COMPILE_OPTIONS
+          "-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vnni;-mavx512vpopcntdq;-mfma;-mf16c;-mpopcnt;-mavx512fp16;-mavx512bf16"
+      )
     else()
       # Per-file SIMD flags (MSVC)
       add_compile_options(/bigobj)

faiss/impl/ScalarQuantizer.cpp

@@ -154,7 +154,7 @@ void ScalarQuantizer::train(size_t n, const float* x) {
 }
 
 ScalarQuantizer::SQuantizer* ScalarQuantizer::select_quantizer() const {
-    return with_simd_level([&]<SIMDLevel SL>() -> SQuantizer* {
+    return with_simd_level_spr([&]<SIMDLevel SL>() -> SQuantizer* {
         if constexpr (SL != SIMDLevel::NONE) {
             auto* q = scalar_quantizer::sq_select_quantizer<SL>(
                     qtype, d, trained);
@@ -197,7 +197,7 @@ void ScalarQuantizer::decode(const uint8_t* codes, float* x, size_t n) const {
 ScalarQuantizer::SQDistanceComputer* ScalarQuantizer::get_distance_computer(
         MetricType metric) const {
     FAISS_THROW_IF_NOT(metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT);
-    return with_simd_level([&]<SIMDLevel SL>() -> SQDistanceComputer* {
+    return with_simd_level_spr([&]<SIMDLevel SL>() -> SQDistanceComputer* {
         if constexpr (SL != SIMDLevel::NONE) {
             auto* dc = scalar_quantizer::sq_select_distance_computer<SL>(
                     metric, qtype, d, trained);
@@ -216,7 +216,7 @@ InvertedListScanner* ScalarQuantizer::select_InvertedListScanner(
         bool store_pairs,
         const IDSelector* sel,
         bool by_residual) const {
-    return with_simd_level([&]<SIMDLevel SL>() -> InvertedListScanner* {
+    return with_simd_level_spr([&]<SIMDLevel SL>() -> InvertedListScanner* {
         if constexpr (SL != SIMDLevel::NONE) {
             auto* s = scalar_quantizer::sq_select_InvertedListScanner<SL>(
                     qtype,

faiss/impl/scalar_quantizer/distance_computers.h

@@ -77,6 +77,12 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::NONE> : SQDistanceComputer {
 template <class Similarity, SIMDLevel SL>
 struct DistanceComputerByte : SQDistanceComputer {};
 
+// Byte-domain distance computer for QT_8bit_direct_signed (storage is
+// value+128). Only specialized for AVX512_SPR; other levels fall back to
+// the float-domain DCTemplate path via the dispatch logic.
+template <class Similarity, SIMDLevel SL>
+struct DistanceComputerByteSigned : SQDistanceComputer {};
+
 template <class Similarity>
 struct DistanceComputerByte<Similarity, SIMDLevel::NONE> : SQDistanceComputer {
     using Sim = Similarity;