Accelerate ScalarQuantizer::QT_bf16 with AVX512-BF16. (#4889)

Summary:
This PR accelerates `ScalarQuantizer::QT_bf16` by vectorizing the BF16 encode/decode path with AVX-512 BF16 and integer instructions, targeting Intel® Sapphire Rapids and newer processors.

Distance computations are also optimized to 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.
- **L2 distance:** Leverages the norm expansion identity, with all terms evaluated using the native BF16 dot-product instruction.

Scalar quantizer benchmarks demonstrate speedups of **1.42×** for encoding, **1.32×** for decoding, and **1.16×** for distance computation.

 ---

`perf_tests/bench_scalar_quantizer_encode --d=768 --n=2000 --iterations=20 --benchmark_filter="QT_bf16"`

```
Existing impl:
--------------------------------------------------------------------------------
Benchmark                      Time             CPU   Iterations UserCounters...
--------------------------------------------------------------------------------
QT_bf16/iterations:20     689149 ns       689010 ns           20 code_size=1.536k

This PR:
--------------------------------------------------------------------------------
Benchmark                      Time             CPU   Iterations UserCounters...
--------------------------------------------------------------------------------
QT_bf16/iterations:20     485337 ns       484511 ns           20 code_size=1.536k
```

`perf_tests/bench_scalar_quantizer_decode --d=768 --n=2000 --iterations=20 --benchmark_filter="QT_bf16"`

```
Existing impl:
--------------------------------------------------------------------------------
Benchmark                      Time             CPU   Iterations UserCounters...
--------------------------------------------------------------------------------
QT_bf16/iterations:20     498607 ns       498539 ns           20 code_size=1.536k

This PR:
--------------------------------------------------------------------------------
Benchmark                      Time             CPU   Iterations UserCounters...
--------------------------------------------------------------------------------
QT_bf16/iterations:20     376614 ns       374200 ns           20 code_size=1.536k
```

`perf_tests/bench_scalar_quantizer_distance --d=768 --n=2000 --iterations=20 --benchmark_filter="QT_bf16"`

```
Existing impl:
-----------------------------------------------------------------------------------
Benchmark                         Time             CPU   Iterations UserCounters...
-----------------------------------------------------------------------------------
QT_bf16/L2/iterations:20  254155481 ns    253585062 ns           20 code_size=1.536k
QT_bf16/IP/iterations:20  221556582 ns    221156894 ns           20 code_size=1.536k

This PR:
-----------------------------------------------------------------------------------
Benchmark                         Time             CPU   Iterations UserCounters...
-----------------------------------------------------------------------------------
QT_bf16/L2/iterations:20  226742244 ns    226270273 ns           20 code_size=1.536k
QT_bf16/IP/iterations:20  190916826 ns    190591678 ns           20 code_size=1.536k
```

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4889

Reviewed By: mdouze

Differential Revision: D106149331

Pulled By: mnorris11

fbshipit-source-id: 114097d0597028c43c19727b677dcc494250be41

faiss/CMakeLists.txt

@@ -40,6 +40,7 @@ set(FAISS_SIMD_AVX512_SRC
 # flag. Not compiled into faiss_avx512 — that target stops at the
 # baseline AVX-512F/CD/VL/DQ/BW feature set.
 set(FAISS_SIMD_AVX512_SPR_SRC
+  impl/scalar_quantizer/sq-avx512-spr.cpp
   utils/simd_impl/rabitq_avx512_spr.cpp
   utils/hamming_distance/hamming_avx512_spr.cpp
 )
@@ -541,7 +542,7 @@ if(FAISS_OPT_LEVEL STREQUAL "dd")
       set_source_files_properties(${FAISS_SIMD_AVX512_SPR_SRC}
         TARGET_DIRECTORY faiss
         PROPERTIES COMPILE_OPTIONS
-          "-mavx2;-mfma;-mf16c;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vpopcntdq;-mpopcnt"
+          "-mavx2;-mfma;-mf16c;-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mavx512vpopcntdq;-mpopcnt;-mavx512vnni;-mavx512fp16;-mavx512bf16"
       )
     else()
       # Per-file SIMD flags (MSVC)

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/quantizers.h

@@ -252,16 +252,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(
@@ -276,6 +272,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
  *******************************************************************/

faiss/impl/scalar_quantizer/sq-avx512-impl.h

@@ -0,0 +1,553 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstring>
+
+namespace faiss {
+
+namespace scalar_quantizer {
+
+using simd16float32 = faiss::simd16float32_tpl<SIMDLevel::AVX512>;
+using simd512bit = faiss::simd512bit_tpl<SIMDLevel::AVX512>;
+
+/**********************************************************
+ * TurboQuant bit-unpacking helpers
+ **********************************************************/
+
+namespace {
+
+FAISS_ALWAYS_INLINE uint16_t load_u16(const uint8_t* ptr) {
+    uint16_t value;
+    std::memcpy(&value, ptr, sizeof(value));
+    return value;
+}
+
+FAISS_ALWAYS_INLINE uint32_t load_u32(const uint8_t* ptr) {
+    uint32_t value;
+    std::memcpy(&value, ptr, sizeof(value));
+    return value;
+}
+
+FAISS_ALWAYS_INLINE uint64_t load_u64(const uint8_t* ptr) {
+    uint64_t value;
+    std::memcpy(&value, ptr, sizeof(value));
+    return value;
+}
+
+FAISS_ALWAYS_INLINE uint32_t load_u24(const uint8_t* ptr) {
+    return static_cast<uint32_t>(ptr[0]) |
+            (static_cast<uint32_t>(ptr[1]) << 8) |
+            (static_cast<uint32_t>(ptr[2]) << 16);
+}
+
+FAISS_ALWAYS_INLINE __m256i unpack_8x3bit_to_u32(uint32_t packed) {
+    const __m256i shifts = _mm256_setr_epi32(0, 3, 6, 9, 12, 15, 18, 21);
+    const __m256i indices =
+            _mm256_srlv_epi32(_mm256_set1_epi32(packed), shifts);
+    return _mm256_and_si256(indices, _mm256_set1_epi32(0x7));
+}
+
+FAISS_ALWAYS_INLINE __m512i unpack_16x1bit_to_u32(const uint8_t* code, int i) {
+    const uint32_t packed = load_u16(code + (static_cast<size_t>(i) >> 3));
+    const __m512i shifts = _mm512_setr_epi32(
+            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
+    const __m512i indices =
+            _mm512_srlv_epi32(_mm512_set1_epi32(packed), shifts);
+    return _mm512_and_si512(indices, _mm512_set1_epi32(0x1));
+}
+
+FAISS_ALWAYS_INLINE __m512i unpack_16x2bit_to_u32(const uint8_t* code, int i) {
+    const uint32_t packed = load_u32(code + (static_cast<size_t>(i) >> 2));
+    const __m512i shifts = _mm512_setr_epi32(
+            0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30);
+    const __m512i indices =
+            _mm512_srlv_epi32(_mm512_set1_epi32(packed), shifts);
+    return _mm512_and_si512(indices, _mm512_set1_epi32(0x3));
+}
+
+FAISS_ALWAYS_INLINE __m512i unpack_16x3bit_to_u32(const uint8_t* code, int i) {
+    const size_t byte_offset = (static_cast<size_t>(i) >> 4) * 6;
+    const __m256i low = unpack_8x3bit_to_u32(load_u24(code + byte_offset));
+    const __m256i high = unpack_8x3bit_to_u32(load_u24(code + byte_offset + 3));
+    __m512i indices = _mm512_castsi256_si512(low);
+    return _mm512_inserti32x8(indices, high, 1);
+}
+
+FAISS_ALWAYS_INLINE __m512i unpack_16x4bit_to_u32(const uint8_t* code, int i) {
+    const uint64_t packed = load_u64(code + (static_cast<size_t>(i) >> 1));
+    const __m256i shifts = _mm256_setr_epi32(0, 4, 8, 12, 16, 20, 24, 28);
+    const __m256i low = _mm256_and_si256(
+            _mm256_srlv_epi32(_mm256_set1_epi32((uint32_t)packed), shifts),
+            _mm256_set1_epi32(0xf));
+    const __m256i high = _mm256_and_si256(
+            _mm256_srlv_epi32(
+                    _mm256_set1_epi32((uint32_t)(packed >> 32)), shifts),
+            _mm256_set1_epi32(0xf));
+    __m512i indices = _mm512_castsi256_si512(low);
+    return _mm512_inserti32x8(indices, high, 1);
+}
+
+} // namespace
+
+/**********************************************************
+ * Codecs
+ **********************************************************/
+
+template <>
+struct Codec8bit<SIMDLevel::AVX512> : Codec8bit<SIMDLevel::NONE> {
+    static FAISS_ALWAYS_INLINE simd16float32
+    decode_16_components(const uint8_t* code, size_t i) {
+        const __m128i c16 = _mm_loadu_si128((__m128i*)(code + i));
+        const __m512i i32 = _mm512_cvtepu8_epi32(c16);
+        const __m512 f16 = _mm512_cvtepi32_ps(i32);
+        const __m512 half_one_255 = _mm512_set1_ps(0.5f / 255.f);
+        const __m512 one_255 = _mm512_set1_ps(1.f / 255.f);
+        return simd16float32(_mm512_fmadd_ps(f16, one_255, half_one_255));
+    }
+};
+
+template <>
+struct Codec4bit<SIMDLevel::AVX512> : Codec4bit<SIMDLevel::NONE> {
+    static FAISS_ALWAYS_INLINE simd16float32
+    decode_16_components(const uint8_t* code, size_t i) {
+        uint64_t c8 = *(uint64_t*)(code + (i >> 1));
+        uint64_t mask = 0x0f0f0f0f0f0f0f0f;
+        uint64_t c8ev = c8 & mask;
+        uint64_t c8od = (c8 >> 4) & mask;
+
+        __m128i c16 =
+                _mm_unpacklo_epi8(_mm_set1_epi64x(c8ev), _mm_set1_epi64x(c8od));
+        __m256i c8lo = _mm256_cvtepu8_epi32(c16);
+        __m256i c8hi = _mm256_cvtepu8_epi32(_mm_srli_si128(c16, 8));
+        __m512i i16 = _mm512_castsi256_si512(c8lo);
+        i16 = _mm512_inserti32x8(i16, c8hi, 1);
+        __m512 f16 = _mm512_cvtepi32_ps(i16);
+        const __m512 half_one_255 = _mm512_set1_ps(0.5f / 15.f);
+        const __m512 one_255 = _mm512_set1_ps(1.f / 15.f);
+        return simd16float32(_mm512_fmadd_ps(f16, one_255, half_one_255));
+    }
+};
+
+template <>
+struct Codec6bit<SIMDLevel::AVX512> : Codec6bit<SIMDLevel::NONE> {
+    static FAISS_ALWAYS_INLINE simd16float32
+    decode_16_components(const uint8_t* code, size_t i) {
+        // pure AVX512 implementation (not necessarily the fastest).
+        // see:
+        // https://github.com/zilliztech/knowhere/blob/main/thirdparty/faiss/faiss/impl/ScalarQuantizerCodec_avx512.h
+
+        // clang-format off
+
+        // 16 components, 16x6 bit=12 bytes
+        const __m128i bit_6v =
+                _mm_maskz_loadu_epi8(0b0000111111111111, code + (i >> 2) * 3);
+        const __m256i bit_6v_256 = _mm256_broadcast_i32x4(bit_6v);
+
+        // 00 01 02 03 04 05 06 07 08 09 0A 0B 0C 0D 0E 0F
+        // 00          01          02          03
+        const __m256i shuffle_mask = _mm256_setr_epi16(
+                0xFF00, 0x0100, 0x0201, 0xFF02,
+                0xFF03, 0x0403, 0x0504, 0xFF05,
+                0xFF06, 0x0706, 0x0807, 0xFF08,
+                0xFF09, 0x0A09, 0x0B0A, 0xFF0B);
+        const __m256i shuffled = _mm256_shuffle_epi8(bit_6v_256, shuffle_mask);
+
+        // 0: xxxxxxxx xx543210
+        // 1: xxxx5432 10xxxxxx
+        // 2: xxxxxx54 3210xxxx
+        // 3: xxxxxxxx 543210xx
+        const __m256i shift_right_v = _mm256_setr_epi16(
+                0x0U, 0x6U, 0x4U, 0x2U,
+                0x0U, 0x6U, 0x4U, 0x2U,
+                0x0U, 0x6U, 0x4U, 0x2U,
+                0x0U, 0x6U, 0x4U, 0x2U);
+        __m256i shuffled_shifted = _mm256_srlv_epi16(shuffled, shift_right_v);
+
+        // remove unneeded bits
+        shuffled_shifted =
+                _mm256_and_si256(shuffled_shifted, _mm256_set1_epi16(0x003F));
+
+        // scale
+        const __m512 f8 =
+                _mm512_cvtepi32_ps(_mm512_cvtepi16_epi32(shuffled_shifted));
+        const __m512 half_one_255 = _mm512_set1_ps(0.5f / 63.f);
+        const __m512 one_255 = _mm512_set1_ps(1.f / 63.f);
+        return simd16float32(_mm512_fmadd_ps(f8, one_255, half_one_255));
+
+        // clang-format on
+    }
+};
+
+/**********************************************************
+ * Quantizers (uniform and non-uniform)
+ **********************************************************/
+
+template <class Codec>
+struct QuantizerTemplate<
+        Codec,
+        scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
+        SIMDLevel::AVX512>
+        : QuantizerTemplate<
+                  Codec,
+                  scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
+                  SIMDLevel::NONE> {
+    QuantizerTemplate(size_t d, const std::vector<float>& trained)
+            : QuantizerTemplate<
+                      Codec,
+                      scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
+                      SIMDLevel::NONE>(d, trained) {
+        assert(d % 16 == 0);
+    }
+
+    FAISS_ALWAYS_INLINE simd16float32
+    reconstruct_16_components(const uint8_t* code, int i) const {
+        __m512 xi = Codec::decode_16_components(code, i).f;
+        return simd16float32(_mm512_fmadd_ps(
+                xi, _mm512_set1_ps(this->vdiff), _mm512_set1_ps(this->vmin)));
+    }
+};
+
+template <class Codec>
+struct QuantizerTemplate<
+        Codec,
+        scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
+        SIMDLevel::AVX512>
+        : QuantizerTemplate<
+                  Codec,
+                  scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
+                  SIMDLevel::NONE> {
+    QuantizerTemplate(size_t d, const std::vector<float>& trained)
+            : QuantizerTemplate<
+                      Codec,
+                      scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
+                      SIMDLevel::NONE>(d, trained) {
+        assert(d % 16 == 0);
+    }
+
+    FAISS_ALWAYS_INLINE simd16float32
+    reconstruct_16_components(const uint8_t* code, int i) const {
+        __m512 xi = Codec::decode_16_components(code, i).f;
+        return simd16float32(_mm512_fmadd_ps(
+                xi,
+                _mm512_loadu_ps(this->vdiff + i),
+                _mm512_loadu_ps(this->vmin + i)));
+    }
+};
+
+/**********************************************************
+ * TurboQuant MSE quantizer
+ **********************************************************/
+
+#define DEFINE_TQMSE_AVX512_SPECIALIZATION(NBITS, INDEX_EXPR)               \
+    template <>                                                             \
+    struct QuantizerTurboQuantMSE<NBITS, SIMDLevel::AVX512>                 \
+            : QuantizerTurboQuantMSE<NBITS, SIMDLevel::NONE> {              \
+        using Base = QuantizerTurboQuantMSE<NBITS, SIMDLevel::NONE>;        \
+                                                                            \
+        QuantizerTurboQuantMSE(size_t d, const std::vector<float>& trained) \
+                : Base(d, trained) {                                        \
+            assert(d % 16 == 0);                                            \
+        }                                                                   \
+                                                                            \
+        FAISS_ALWAYS_INLINE simd16float32                                   \
+        reconstruct_16_components(const uint8_t* code, int i) const {       \
+            const __m512i indices = (INDEX_EXPR);                           \
+            return simd16float32(_mm512_i32gather_ps(                       \
+                    indices, this->centroids, sizeof(float)));              \
+        }                                                                   \
+    }
+
+DEFINE_TQMSE_AVX512_SPECIALIZATION(1, unpack_16x1bit_to_u32(code, i));
+DEFINE_TQMSE_AVX512_SPECIALIZATION(2, unpack_16x2bit_to_u32(code, i));
+DEFINE_TQMSE_AVX512_SPECIALIZATION(3, unpack_16x3bit_to_u32(code, i));
+DEFINE_TQMSE_AVX512_SPECIALIZATION(4, unpack_16x4bit_to_u32(code, i));
+
+#undef DEFINE_TQMSE_AVX512_SPECIALIZATION
+
+template <>
+struct QuantizerTurboQuantMSE<8, SIMDLevel::AVX512>
+        : QuantizerTurboQuantMSE<8, SIMDLevel::NONE> {
+    using Base = QuantizerTurboQuantMSE<8, SIMDLevel::NONE>;
+
+    QuantizerTurboQuantMSE(size_t d, const std::vector<float>& trained)
+            : Base(d, trained) {
+        assert(d % 16 == 0);
+    }
+
+    FAISS_ALWAYS_INLINE simd16float32
+    reconstruct_16_components(const uint8_t* code, int i) const {
+        const __m128i packed = _mm_loadu_si128(
+                (const __m128i*)(code + static_cast<size_t>(i)));
+        const __m512i indices = _mm512_cvtepu8_epi32(packed);
+        return simd16float32(
+                _mm512_i32gather_ps(indices, this->centroids, sizeof(float)));
+    }
+};
+
+/**********************************************************
+ * FP16 Quantizer
+ **********************************************************/
+
+template <>
+struct QuantizerFP16<SIMDLevel::AVX512> : QuantizerFP16<SIMDLevel::NONE> {
+    QuantizerFP16(size_t d, const std::vector<float>& trained)
+            : QuantizerFP16<SIMDLevel::NONE>(d, trained) {
+        assert(d % 16 == 0);
+    }
+
+    FAISS_ALWAYS_INLINE simd16float32
+    reconstruct_16_components(const uint8_t* code, int i) const {
+        __m256i codei = _mm256_loadu_si256((const __m256i*)(code + 2 * i));
+        return simd16float32(_mm512_cvtph_ps(codei));
+    }
+};
+
+/**********************************************************
+ * BF16 Quantizer
+ **********************************************************/
+
+template <>
+struct QuantizerBF16<SIMDLevel::AVX512> : QuantizerBF16<SIMDLevel::NONE> {
+    QuantizerBF16(size_t d, const std::vector<float>& trained)
+            : QuantizerBF16<SIMDLevel::NONE>(d, trained) {
+        assert(d % 16 == 0);
+    }
+
+    FAISS_ALWAYS_INLINE simd16float32
+    reconstruct_16_components(const uint8_t* code, int i) const {
+        __m256i code_256i = _mm256_loadu_si256((const __m256i*)(code + 2 * i));
+        __m512i code_512i = _mm512_cvtepu16_epi32(code_256i);
+        code_512i = _mm512_slli_epi32(code_512i, 16);
+        return simd16float32(_mm512_castsi512_ps(code_512i));
+    }
+};
+
+/**********************************************************
+ * 8bit Direct Quantizer
+ **********************************************************/
+
+template <>
+struct Quantizer8bitDirect<SIMDLevel::AVX512>
+        : Quantizer8bitDirect<SIMDLevel::NONE> {
+    Quantizer8bitDirect(size_t d, const std::vector<float>& trained)
+            : Quantizer8bitDirect<SIMDLevel::NONE>(d, trained) {
+        assert(d % 16 == 0);
+    }
+
+    FAISS_ALWAYS_INLINE simd16float32
+    reconstruct_16_components(const uint8_t* code, int i) const {
+        __m128i x16 = _mm_loadu_si128((__m128i*)(code + i)); // 16 * int8
+        __m512i y16 = _mm512_cvtepu8_epi32(x16);             // 16 * int32
+        return simd16float32(_mm512_cvtepi32_ps(y16));       // 16 * float32
+    }
+};
+
+/**********************************************************
+ * 8bit Direct Signed Quantizer
+ **********************************************************/
+
+template <>
+struct Quantizer8bitDirectSigned<SIMDLevel::AVX512>
+        : Quantizer8bitDirectSigned<SIMDLevel::NONE> {
+    Quantizer8bitDirectSigned(size_t d, const std::vector<float>& trained)
+            : Quantizer8bitDirectSigned<SIMDLevel::NONE>(d, trained) {
+        assert(d % 16 == 0);
+    }
+
+    FAISS_ALWAYS_INLINE simd16float32
+    reconstruct_16_components(const uint8_t* code, int i) const {
+        __m128i x16 = _mm_loadu_si128((__m128i*)(code + i)); // 16 * int8
+        __m512i y16 = _mm512_cvtepu8_epi32(x16);             // 16 * int32
+        __m512i c16 = _mm512_set1_epi32(128);
+        __m512i z16 = _mm512_sub_epi32(y16, c16); // subtract 128 from all lanes
+        return simd16float32(_mm512_cvtepi32_ps(z16)); // 16 * float32
+    }
+};
+
+/**********************************************************
+ * Similarities (L2 and IP)
+ **********************************************************/
+
+template <>
+struct SimilarityL2<SIMDLevel::AVX512> {
+    static constexpr int simdwidth = 16;
+    static constexpr SIMDLevel simd_level = SIMDLevel::AVX512;
+    static constexpr MetricType metric_type = METRIC_L2;
+
+    const float *y, *yi;
+
+    explicit SimilarityL2(const float* y) : y(y), yi(nullptr) {}
+
+    simd16float32 accu16;
+
+    FAISS_ALWAYS_INLINE void begin_16() {
+        accu16.clear();
+        yi = y;
+    }
+
+    FAISS_ALWAYS_INLINE void add_16_components(simd16float32 x) {
+        simd16float32 yiv(yi);
+        yi += 16;
+        simd16float32 tmp = yiv - x;
+        accu16 = accu16 + tmp * tmp;
+    }
+
+    FAISS_ALWAYS_INLINE void add_16_components_2(
+            simd16float32 x,
+            simd16float32 y_2) {
+        simd16float32 tmp = y_2 - x;
+        accu16 = accu16 + tmp * tmp;
+    }
+
+    FAISS_ALWAYS_INLINE float result_16() {
+        return horizontal_add(accu16);
+    }
+};
+
+template <>
+struct SimilarityIP<SIMDLevel::AVX512> {
+    static constexpr int simdwidth = 16;
+    static constexpr SIMDLevel simd_level = SIMDLevel::AVX512;
+    static constexpr MetricType metric_type = METRIC_INNER_PRODUCT;
+
+    const float *y, *yi;
+
+    explicit SimilarityIP(const float* y) : y(y), yi(nullptr) {}
+
+    simd16float32 accu16;
+
+    FAISS_ALWAYS_INLINE void begin_16() {
+        accu16.clear();
+        yi = y;
+    }
+
+    FAISS_ALWAYS_INLINE void add_16_components(simd16float32 x) {
+        simd16float32 yiv(yi);
+        yi += 16;
+        accu16 = accu16 + yiv * x;
+    }
+
+    FAISS_ALWAYS_INLINE void add_16_components_2(
+            simd16float32 x1,
+            simd16float32 x2) {
+        accu16 = accu16 + x1 * x2;
+    }
+
+    FAISS_ALWAYS_INLINE float result_16() {
+        return horizontal_add(accu16);
+    }
+};
+
+/**********************************************************
+ * Distance Computers
+ **********************************************************/
+
+template <class Quantizer, class Similarity>
+struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512>
+        : SQDistanceComputer {
+    using Sim = Similarity;
+
+    Quantizer quant;
+
+    DCTemplate(size_t d, const std::vector<float>& trained)
+            : quant(d, trained) {}
+
+    float compute_distance(const float* x, const uint8_t* code) const {
+        Similarity sim(x);
+        sim.begin_16();
+        for (size_t i = 0; i < quant.d; i += 16) {
+            simd16float32 xi = quant.reconstruct_16_components(code, i);
+            sim.add_16_components(xi);
+        }
+        return sim.result_16();
+    }
+
+    float compute_code_distance(const uint8_t* code1, const uint8_t* code2)
+            const {
+        Similarity sim(nullptr);
+        sim.begin_16();
+        for (size_t i = 0; i < quant.d; i += 16) {
+            simd16float32 x1 = quant.reconstruct_16_components(code1, i);
+            simd16float32 x2 = quant.reconstruct_16_components(code2, i);
+            sim.add_16_components_2(x1, x2);
+        }
+        return sim.result_16();
+    }
+
+    void set_query(const float* x) final {
+        q = x;
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return compute_code_distance(
+                codes + i * code_size, codes + j * code_size);
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        return compute_distance(q, code);
+    }
+};
+
+template <class Similarity>
+struct DistanceComputerByte<Similarity, SIMDLevel::AVX512>
+        : SQDistanceComputer {
+    using Sim = Similarity;
+
+    int d;
+    std::vector<uint8_t> tmp;
+
+    DistanceComputerByte(int d, const std::vector<float>&) : d(d), tmp(d) {}
+
+    int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
+            const {
+        // compute 16 lanes of 32-bit products (16-bytes) at once for
+        // the supported metrics
+        __m512i accu = _mm512_setzero_si512();
+        constexpr int kLanes = 16;
+        for (int i = 0; i < d; i += kLanes) {
+            __m128i c1 = _mm_loadu_si128((__m128i*)(code1 + i));
+            __m128i c2 = _mm_loadu_si128((__m128i*)(code2 + i));
+            __m512i c1i = _mm512_cvtepu8_epi32(c1);
+            __m512i c2i = _mm512_cvtepu8_epi32(c2);
+
+            __m512i v;
+            if (Sim::metric_type == METRIC_INNER_PRODUCT) {
+                v = _mm512_mullo_epi32(c1i, c2i);
+            } else {
+                __m512i diff = _mm512_sub_epi32(c1i, c2i);
+                v = _mm512_mullo_epi32(diff, diff);
+            }
+            accu = _mm512_add_epi32(accu, v);
+        }
+        return _mm512_reduce_add_epi32(accu);
+    }
+
+    void set_query(const float* x) final {
+        for (int i = 0; i < d; i++) {
+            tmp[i] = int(x[i]);
+        }
+    }
+
+    int compute_distance(const float* x, const uint8_t* code) {
+        set_query(x);
+        return compute_code_distance(tmp.data(), code);
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return compute_code_distance(
+                codes + i * code_size, codes + j * code_size);
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        return compute_code_distance(tmp.data(), code);
+    }
+};
+
+} // namespace scalar_quantizer
+} // namespace faiss

faiss/impl/scalar_quantizer/sq-avx512-spr.cpp

@@ -0,0 +1,320 @@
+/*
+ * 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.
+ */
+
+#ifdef COMPILE_SIMD_AVX512_SPR
+
+#include <immintrin.h>
+
+#include <faiss/impl/scalar_quantizer/codecs.h>
+#include <faiss/impl/scalar_quantizer/distance_computers.h>
+#include <faiss/impl/scalar_quantizer/quantizers.h>
+#include <faiss/impl/scalar_quantizer/scanners.h>
+#include <faiss/impl/scalar_quantizer/similarities.h>
+#include <faiss/impl/simdlib/simdlib_avx512.h>
+
+#include <faiss/impl/scalar_quantizer/sq-avx512-impl.h>
+
+namespace faiss {
+namespace scalar_quantizer {
+
+/**********************************************************
+ * Codecs — inherit AVX512 implementations
+ **********************************************************/
+
+template <>
+struct Codec8bit<SIMDLevel::AVX512_SPR> : Codec8bit<SIMDLevel::AVX512> {};
+
+template <>
+struct Codec4bit<SIMDLevel::AVX512_SPR> : Codec4bit<SIMDLevel::AVX512> {};
+
+template <>
+struct Codec6bit<SIMDLevel::AVX512_SPR> : Codec6bit<SIMDLevel::AVX512> {};
+
+/**********************************************************
+ * Quantizers — inherit AVX512 implementations
+ **********************************************************/
+
+template <class Codec>
+struct QuantizerTemplate<
+        Codec,
+        QuantizerTemplateScaling::UNIFORM,
+        SIMDLevel::AVX512_SPR>
+        : QuantizerTemplate<
+                  Codec,
+                  QuantizerTemplateScaling::UNIFORM,
+                  SIMDLevel::AVX512> {
+    using QuantizerTemplate<
+            Codec,
+            QuantizerTemplateScaling::UNIFORM,
+            SIMDLevel::AVX512>::QuantizerTemplate;
+};
+
+template <class Codec>
+struct QuantizerTemplate<
+        Codec,
+        QuantizerTemplateScaling::NON_UNIFORM,
+        SIMDLevel::AVX512_SPR>
+        : QuantizerTemplate<
+                  Codec,
+                  QuantizerTemplateScaling::NON_UNIFORM,
+                  SIMDLevel::AVX512> {
+    using QuantizerTemplate<
+            Codec,
+            QuantizerTemplateScaling::NON_UNIFORM,
+            SIMDLevel::AVX512>::QuantizerTemplate;
+};
+
+template <>
+struct QuantizerFP16<SIMDLevel::AVX512_SPR> : QuantizerFP16<SIMDLevel::AVX512> {
+    using QuantizerFP16<SIMDLevel::AVX512>::QuantizerFP16;
+};
+
+template <>
+struct QuantizerBF16<SIMDLevel::AVX512_SPR> : QuantizerBF16<SIMDLevel::AVX512> {
+    using QuantizerBF16<SIMDLevel::AVX512>::QuantizerBF16;
+
+    void encode_vector(const float* x, uint8_t* code) const override {
+        encode_bf16_simd(x, (uint16_t*)code, this->d);
+    }
+
+    void decode_vector(const uint8_t* code, float* x) const override {
+        decode_bf16_simd((const uint16_t*)code, x, this->d);
+    }
+};
+
+template <>
+struct Quantizer8bitDirect<SIMDLevel::AVX512_SPR>
+        : Quantizer8bitDirect<SIMDLevel::AVX512> {
+    using Quantizer8bitDirect<SIMDLevel::AVX512>::Quantizer8bitDirect;
+};
+
+template <>
+struct Quantizer8bitDirectSigned<SIMDLevel::AVX512_SPR>
+        : Quantizer8bitDirectSigned<SIMDLevel::AVX512> {
+    using Quantizer8bitDirectSigned<
+            SIMDLevel::AVX512>::Quantizer8bitDirectSigned;
+};
+
+/**********************************************************
+ * TurboQuant MSE — inherit AVX512 implementations
+ **********************************************************/
+
+template <int NBits>
+struct QuantizerTurboQuantMSE<NBits, SIMDLevel::AVX512_SPR>
+        : QuantizerTurboQuantMSE<NBits, SIMDLevel::AVX512> {
+    using QuantizerTurboQuantMSE<NBits, SIMDLevel::AVX512>::
+            QuantizerTurboQuantMSE;
+};
+
+/**********************************************************
+ * Similarities — inherit AVX512 implementations
+ **********************************************************/
+
+template <>
+struct SimilarityL2<SIMDLevel::AVX512_SPR> : SimilarityL2<SIMDLevel::AVX512> {
+    using SimilarityL2<SIMDLevel::AVX512>::SimilarityL2;
+    static constexpr SIMDLevel simd_level = SIMDLevel::AVX512_SPR;
+};
+
+template <>
+struct SimilarityIP<SIMDLevel::AVX512_SPR> : SimilarityIP<SIMDLevel::AVX512> {
+    using SimilarityIP<SIMDLevel::AVX512>::SimilarityIP;
+    static constexpr SIMDLevel simd_level = SIMDLevel::AVX512_SPR;
+};
+
+/**********************************************************
+ * Generic DCTemplate and DistanceComputerByte
+ *
+ * Delegate to AVX512 implementations.
+ **********************************************************/
+
+template <class Quantizer, class Similarity>
+struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512_SPR>
+        : DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512> {
+    using DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512>::DCTemplate;
+};
+
+template <class Similarity>
+struct DistanceComputerByte<Similarity, SIMDLevel::AVX512_SPR>
+        : DistanceComputerByte<Similarity, SIMDLevel::AVX512> {
+    using DistanceComputerByte<Similarity, SIMDLevel::AVX512>::
+            DistanceComputerByte;
+};
+
+/**********************************************************
+ * BF16 native distance helpers using VDPBF16PS
+ **********************************************************/
+
+static FAISS_ALWAYS_INLINE float bf16_vdpbf16ps(
+        const uint16_t* a,
+        const uint16_t* b,
+        size_t d) {
+    __m512 acc = _mm512_setzero_ps();
+    size_t i = 0;
+    for (; i + 32 <= d; i += 32) {
+        __m512bh va = (__m512bh)_mm512_loadu_epi16(a + i);
+        __m512bh vb = (__m512bh)_mm512_loadu_epi16(b + i);
+        acc = _mm512_dpbf16_ps(acc, va, vb);
+    }
+    // Remainder: 16 elements (d % 16 == 0 but may not be % 32)
+    if (i < d) {
+        __m256i a_lo = _mm256_loadu_epi16(a + i);
+        __m256i b_lo = _mm256_loadu_epi16(b + i);
+        __m512bh va =
+                (__m512bh)_mm512_inserti64x4(_mm512_setzero_si512(), a_lo, 0);
+        __m512bh vb =
+                (__m512bh)_mm512_inserti64x4(_mm512_setzero_si512(), b_lo, 0);
+        acc = _mm512_dpbf16_ps(acc, va, vb);
+    }
+    return _mm512_reduce_add_ps(acc);
+}
+
+static FAISS_ALWAYS_INLINE float bf16_L2_asymmetric(
+        const uint16_t* query_bf16,
+        const uint16_t* code,
+        size_t d) {
+    __m512 acc_qc = _mm512_setzero_ps();
+    __m512 acc_cc = _mm512_setzero_ps();
+    size_t i = 0;
+    for (; i + 32 <= d; i += 32) {
+        __m512bh vq = (__m512bh)_mm512_loadu_epi16(query_bf16 + i);
+        __m512bh vc = (__m512bh)_mm512_loadu_epi16(code + i);
+        acc_qc = _mm512_dpbf16_ps(acc_qc, vq, vc);
+        acc_cc = _mm512_dpbf16_ps(acc_cc, vc, vc);
+    }
+    if (i < d) {
+        __m256i q_lo = _mm256_loadu_epi16(query_bf16 + i);
+        __m256i c_lo = _mm256_loadu_epi16(code + i);
+        __m512bh vq =
+                (__m512bh)_mm512_inserti64x4(_mm512_setzero_si512(), q_lo, 0);
+        __m512bh vc =
+                (__m512bh)_mm512_inserti64x4(_mm512_setzero_si512(), c_lo, 0);
+        acc_qc = _mm512_dpbf16_ps(acc_qc, vq, vc);
+        acc_cc = _mm512_dpbf16_ps(acc_cc, vc, vc);
+    }
+    float dot_qc = _mm512_reduce_add_ps(acc_qc);
+    float norm_c = _mm512_reduce_add_ps(acc_cc);
+    return -2.0f * dot_qc + norm_c;
+}
+
+static FAISS_ALWAYS_INLINE float bf16_L2_symmetric(
+        const uint16_t* a,
+        const uint16_t* b,
+        size_t d) {
+    __m512 acc_ab = _mm512_setzero_ps();
+    __m512 acc_aa = _mm512_setzero_ps();
+    __m512 acc_bb = _mm512_setzero_ps();
+    size_t i = 0;
+    for (; i + 32 <= d; i += 32) {
+        __m512bh va = (__m512bh)_mm512_loadu_epi16(a + i);
+        __m512bh vb = (__m512bh)_mm512_loadu_epi16(b + i);
+        acc_ab = _mm512_dpbf16_ps(acc_ab, va, vb);
+        acc_aa = _mm512_dpbf16_ps(acc_aa, va, va);
+        acc_bb = _mm512_dpbf16_ps(acc_bb, vb, vb);
+    }
+    if (i < d) {
+        __m256i a_lo = _mm256_loadu_epi16(a + i);
+        __m256i b_lo = _mm256_loadu_epi16(b + i);
+        __m512bh va =
+                (__m512bh)_mm512_inserti64x4(_mm512_setzero_si512(), a_lo, 0);
+        __m512bh vb =
+                (__m512bh)_mm512_inserti64x4(_mm512_setzero_si512(), b_lo, 0);
+        acc_ab = _mm512_dpbf16_ps(acc_ab, va, vb);
+        acc_aa = _mm512_dpbf16_ps(acc_aa, va, va);
+        acc_bb = _mm512_dpbf16_ps(acc_bb, vb, vb);
+    }
+    return _mm512_reduce_add_ps(acc_aa) - 2.0f * _mm512_reduce_add_ps(acc_ab) +
+            _mm512_reduce_add_ps(acc_bb);
+}
+
+/**********************************************************
+ * BF16 + Inner Product distance computer (SPR)
+ **********************************************************/
+
+struct DCBF16_IP : SQDistanceComputer {
+    using Sim = SimilarityIP<SIMDLevel::AVX512_SPR>;
+
+    size_t d;
+    std::vector<uint16_t> query_bf16;
+
+    DCBF16_IP(size_t d, const std::vector<float>&) : d(d), query_bf16(d) {}
+
+    void set_query(const float* x) final {
+        q = x;
+        encode_bf16_simd(x, query_bf16.data(), d);
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        return bf16_vdpbf16ps(query_bf16.data(), (const uint16_t*)code, d);
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return bf16_vdpbf16ps(
+                (const uint16_t*)(codes + i * code_size),
+                (const uint16_t*)(codes + j * code_size),
+                d);
+    }
+};
+
+/**********************************************************
+ * BF16 + L2 distance computer (SPR)
+ **********************************************************/
+
+struct DCBF16_L2 : SQDistanceComputer {
+    using Sim = SimilarityL2<SIMDLevel::AVX512_SPR>;
+
+    size_t d;
+    std::vector<uint16_t> query_bf16;
+    float query_norm_sq;
+
+    DCBF16_L2(size_t d, const std::vector<float>&)
+            : d(d), query_bf16(d), query_norm_sq(0) {}
+
+    void set_query(const float* x) final {
+        q = x;
+        encode_bf16_simd(x, query_bf16.data(), d);
+        query_norm_sq = bf16_vdpbf16ps(query_bf16.data(), query_bf16.data(), d);
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        return query_norm_sq +
+                bf16_L2_asymmetric(query_bf16.data(), (const uint16_t*)code, d);
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return bf16_L2_symmetric(
+                (const uint16_t*)(codes + i * code_size),
+                (const uint16_t*)(codes + j * code_size),
+                d);
+    }
+};
+
+template <>
+struct DCTemplate<
+        QuantizerBF16<SIMDLevel::AVX512_SPR>,
+        SimilarityIP<SIMDLevel::AVX512_SPR>,
+        SIMDLevel::AVX512_SPR> : DCBF16_IP {
+    using Sim = SimilarityIP<SIMDLevel::AVX512_SPR>;
+    using DCBF16_IP::DCBF16_IP;
+};
+
+template <>
+struct DCTemplate<
+        QuantizerBF16<SIMDLevel::AVX512_SPR>,
+        SimilarityL2<SIMDLevel::AVX512_SPR>,
+        SIMDLevel::AVX512_SPR> : DCBF16_L2 {
+    using Sim = SimilarityL2<SIMDLevel::AVX512_SPR>;
+    using DCBF16_L2::DCBF16_L2;
+};
+
+} // namespace scalar_quantizer
+} // namespace faiss
+
+#define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX512_SPR
+#include <faiss/impl/scalar_quantizer/sq-dispatch.h>
+
+#endif // COMPILE_SIMD_AVX512_SPR

faiss/impl/scalar_quantizer/sq-dispatch.h

@@ -27,7 +27,7 @@ constexpr SIMDLevel SL = THE_LEVEL_TO_DISPATCH;
 // Returns true if dimension d is compatible with the given SIMD level
 template <SIMDLevel SL2>
 constexpr bool is_dimension_compatible(size_t d) {
-    if constexpr (SL2 == SIMDLevel::AVX512) {
+    if constexpr (SL2 == SIMDLevel::AVX512 || SL2 == SIMDLevel::AVX512_SPR) {
         return d % 16 == 0;
     } else if constexpr (SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
         return d % 8 == 0;
@@ -171,7 +171,8 @@ SQDistanceComputer* select_distance_computer_body(
             return new DCTemplate<QuantizerBF16<SL2>, Sim, SL2>(d, trained);
 
         case ScalarQuantizer::QT_8bit_direct:
-            if constexpr (SL2 == SIMDLevel::AVX512) {
+            if constexpr (
+                    SL2 == SIMDLevel::AVX512 || SL2 == SIMDLevel::AVX512_SPR) {
                 if (d % 32 == 0) {
                     return new DistanceComputerByte<Sim, SL2>(
                             static_cast<int>(d), trained);
@@ -320,7 +321,9 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
                 return scan.template
                 operator()<DCTemplate<QuantizerBF16<SL2>, Similarity, SL2>>();
             case ScalarQuantizer::QT_8bit_direct:
-                if constexpr (SL2 == SIMDLevel::AVX512) {
+                if constexpr (
+                        SL2 == SIMDLevel::AVX512 ||
+                        SL2 == SIMDLevel::AVX512_SPR) {
                     if (d % 32 == 0) {
                         return scan.template
                         operator()<DistanceComputerByte<Similarity, SL2>>();

faiss/impl/simd_dispatch.h

@@ -198,6 +198,15 @@ inline auto with_simd_level(LambdaType&& action) {
             std::forward<LambdaType>(action));
 }
 
+/**
+ * Use for functions with AVX512_SPR-specific implementations.
+ */
+template <typename LambdaType>
+inline auto with_simd_level_spr(LambdaType&& action) {
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0_SPR>(
+            std::forward<LambdaType>(action));
+}
+
 /**
  * Use for functions implemented with simdXintY (256-bit) operations
  * that don't have dedicated AVX512 or SVE implementations.

faiss/utils/bf16.h

@@ -7,8 +7,13 @@
 
 #pragma once
 
+#include <cstddef>
 #include <cstdint>
 
+#if defined(__AVX512F__) || defined(__AVX512BF16__)
+#include <immintrin.h>
+#endif
+
 namespace faiss {
 
 namespace {
@@ -33,4 +38,33 @@ inline float decode_bf16(const uint16_t v) {
     return fp.as_f32;
 }
 
+inline void encode_bf16_simd(const float* src, uint16_t* dst, size_t n) {
+    size_t i = 0;
+#ifdef __AVX512BF16__
+    for (; i + 16 <= n; i += 16) {
+        __m512 v = _mm512_loadu_ps(src + i);
+        __m256bh encoded = _mm512_cvtneps_pbh(v);
+        _mm256_storeu_epi16(dst + i, (__m256i)encoded);
+    }
+#endif
+    for (; i < n; i++) {
+        dst[i] = encode_bf16(src[i]);
+    }
+}
+
+inline void decode_bf16_simd(const uint16_t* src, float* dst, size_t n) {
+    size_t i = 0;
+#if defined(__AVX512F__)
+    for (; i + 16 <= n; i += 16) {
+        __m256i v = _mm256_loadu_si256((const __m256i*)(src + i));
+        __m512i w = _mm512_cvtepu16_epi32(v);
+        w = _mm512_slli_epi32(w, 16);
+        _mm512_storeu_ps(dst + i, _mm512_castsi512_ps(w));
+    }
+#endif
+    for (; i < n; i++) {
+        dst[i] = decode_bf16(src[i]);
+    }
+}
+
 } // namespace faiss

perf_tests/bench_scalar_quantizer_distance.cpp

@@ -11,25 +11,27 @@
 #include <map>
 
 #include <benchmark/benchmark.h>
+#include <faiss/MetricType.h>
 #include <faiss/impl/ScalarQuantizer.h>
 #include <faiss/utils/random.h>
 #include "utils.h"
 
 using namespace faiss;
+
 DEFINE_uint32(d, 128, "dimension");
-DEFINE_uint32(n, 2000, "dimension");
+DEFINE_uint32(n, 2000, "number of vectors");
 DEFINE_uint32(iterations, 20, "iterations");
 
 static void bench_distance(
         benchmark::State& state,
         ScalarQuantizer::QuantizerType type,
+        MetricType metric,
         int d,
         int n) {
     std::vector<float> x(d * n);
 
     float_rand(x.data(), d * n, 12345);
 
-    // make sure it's idempotent
     ScalarQuantizer sq(d, type);
 
     omp_set_num_threads(1);
@@ -44,7 +46,7 @@ static void bench_distance(
     sq.compute_codes(x.data(), codes.data(), n);
 
     std::unique_ptr<ScalarQuantizer::SQDistanceComputer> dc(
-            sq.get_distance_computer());
+            sq.get_distance_computer(metric));
     dc->codes = codes.data();
     dc->code_size = sq.code_size;
 
@@ -56,6 +58,7 @@ static void bench_distance(
                 benchmark::DoNotOptimize(sum_dis += (*dc)(j));
             }
         }
+        benchmark::ClobberMemory();
     }
 }
 
@@ -63,6 +66,7 @@ int main(int argc, char** argv) {
     benchmark::Initialize(&argc, argv);
     gflags::AllowCommandLineReparsing();
     gflags::ParseCommandLineFlags(&argc, &argv, true);
+
     int iterations = FLAGS_iterations;
     int d = FLAGS_d;
     int n = FLAGS_n;
@@ -70,9 +74,24 @@ int main(int argc, char** argv) {
 
     for (auto& [bench_name, quantizer_type] : benchs) {
         benchmark::RegisterBenchmark(
-                bench_name.c_str(), bench_distance, quantizer_type, d, n)
+                (bench_name + "/L2").c_str(),
+                bench_distance,
+                quantizer_type,
+                METRIC_L2,
+                d,
+                n)
+                ->Iterations(iterations);
+
+        benchmark::RegisterBenchmark(
+                (bench_name + "/IP").c_str(),
+                bench_distance,
+                quantizer_type,
+                METRIC_INNER_PRODUCT,
+                d,
+                n)
                 ->Iterations(iterations);
     }
+
     benchmark::RunSpecifiedBenchmarks();
     benchmark::Shutdown();
 }