Accelerate ScalarQuantizer::QT_bf16 with AVX512-BF16.

Vectorizes the BF16 encode/decode path with AVX-512 BF16 and integer
instructions, targeting Intel Sapphire Rapids and newer processors.
Distance computations now operate directly on BF16 data, avoiding the
overhead of FP32 conversion:

- Inner product: the query is encoded to BF16 once, and dot products
  are computed natively in BF16 via VDPBF16PS.
- L2 distance: uses the norm expansion identity, with all terms
  evaluated using the native BF16 dot-product instruction.

Scalar quantizer benchmarks on Sapphire Rapids (d=768, n=2000,
iterations=20) show speedups of 1.42x for encoding, 1.32x for decoding,
and 1.16x for distance computation.

Signed-off-by: Mulugeta Mammo <mulugeta.mammo@intel.com>

Author: mulugetam

faiss/CMakeLists.txt

@@ -28,6 +28,10 @@ set(FAISS_SIMD_AVX512_SRC
   utils/distances_fused/avx512.cpp
   utils/simd_impl/rabitq_avx512.cpp
 )
+set(FAISS_SIMD_AVX512_SPR_SRC
+  impl/scalar_quantizer/sq-avx512-spr.cpp
+  utils/simd_impl/distances_avx512_spr.cpp
+)
 set(FAISS_SIMD_NEON_SRC
   impl/fast_scan/impl-neon.cpp
   impl/scalar_quantizer/sq-neon.cpp
@@ -43,7 +47,7 @@ set(FAISS_SIMD_SVE_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})
 else()
@@ -389,7 +393,7 @@ endif()
 if(NOT WIN32)
   # All modern CPUs support F, CD, VL, DQ, BW extensions.
   # Ref: https://en.wikipedia.org/wiki/AVX512
-  target_compile_options(faiss_avx512 PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mpopcnt>)
+  target_compile_options(faiss_avx512 PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-mavx -mavx2 -mfma -mf16c -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mpopcnt>)
 else()
   target_compile_options(faiss_avx512 PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/arch:AVX512>)
   # we need bigobj for the swig wrapper
@@ -405,14 +409,14 @@ 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>)
+  target_compile_options(faiss_avx512_spr PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-mavx -mavx2 -mfma -mf16c -mbmi2 -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mavx512vbmi2 -mavx512vpopcntdq -mpopcnt -mavx512fp16 -mavx512bf16>)
 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_compile_definitions(faiss_avx512_spr PRIVATE COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512 COMPILE_SIMD_AVX512_SPR )
+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})
 if(NOT FAISS_OPT_LEVEL STREQUAL "sve")
@@ -460,14 +464,21 @@ if(FAISS_OPT_LEVEL STREQUAL "dd")
       target_compile_options(faiss PRIVATE
         $<$<COMPILE_LANGUAGE:CXX>:-mpopcnt -msse4 -mno-avx -mno-avx2>)
       # Per-file SIMD flags
+      # Note: -mavx must be explicit because baseline sets -mno-avx which
+      # overrides the implicit -mavx from higher ISA flags like -mavx512f.
       set_source_files_properties(${FAISS_SIMD_AVX2_SRC}
         TARGET_DIRECTORY faiss
-        PROPERTIES COMPILE_OPTIONS "-mavx2;-mfma;-mf16c;-mpopcnt"
+        PROPERTIES COMPILE_OPTIONS "-mavx;-mavx2;-mfma;-mf16c;-mpopcnt"
       )
       set_source_files_properties(${FAISS_SIMD_AVX512_SRC}
         TARGET_DIRECTORY faiss
         PROPERTIES COMPILE_OPTIONS
-          "-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mfma;-mf16c;-mpopcnt"
+          "-mavx;-mavx2;-mfma;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mpopcnt"
+      )
+      set_source_files_properties(${FAISS_SIMD_AVX512_SPR_SRC}
+        TARGET_DIRECTORY faiss
+        PROPERTIES COMPILE_OPTIONS
+          "-mavx;-mavx2;-mfma;-mbmi2;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vbmi2;-mpopcnt;-mavx512bf16"
       )
     elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
       # ARM NEON is always available on aarch64, no special compiler flags needed

faiss/impl/scalar_quantizer/quantizers.h

@@ -165,16 +165,12 @@ struct QuantizerBF16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
     QuantizerBF16(size_t d_in, const std::vector<float>& /* unused */)
             : d(d_in) {}
 
-    void encode_vector(const float* x, uint8_t* code) const final {
-        for (size_t i = 0; i < d; i++) {
-            ((uint16_t*)code)[i] = encode_bf16(x[i]);
-        }
+    void encode_vector(const float* x, uint8_t* code) const override {
+        encode_bf16_simd(x, (uint16_t*)code, d);
     }
 
-    void decode_vector(const uint8_t* code, float* x) const final {
-        for (size_t i = 0; i < d; i++) {
-            x[i] = decode_bf16(((uint16_t*)code)[i]);
-        }
+    void decode_vector(const uint8_t* code, float* x) const override {
+        decode_bf16_simd((const uint16_t*)code, x, d);
     }
 
     FAISS_ALWAYS_INLINE float reconstruct_component(
@@ -189,6 +185,11 @@ struct QuantizerBF16 : QuantizerBF16<SIMDLevel::NONE> {
     using QuantizerBF16<SIMDLevel::NONE>::QuantizerBF16;
 };
 
+template <>
+struct QuantizerBF16<SIMDLevel::AVX512>;
+template <>
+struct QuantizerBF16<SIMDLevel::AVX512_SPR>;
+
 /*******************************************************************
  * 8bit_direct quantizer
  *******************************************************************/