Add AVX2 and AVX512 byte-domain kernels for QT_8bit_direct_signed (#5409)

Summary:
faiss ships a byte-domain integer distance computer for `QT_8bit_direct_signed` only on AVX512-SPR (`DistanceComputerByteSigned<AVX512_SPR>`, added in https://github.com/facebookresearch/faiss/issues/5173). On AVX2 and plain AVX512 the direct_signed path falls back to the float-decode + FMA `DCTemplate`. This adds the missing AVX2 and AVX512 specializations, mirroring the existing unsigned `DistanceComputerByte<AVX2>`/`<AVX512>` kernels.

The signed kernels widen the stored bytes to int16, unbias for inner product (`value = stored - 128`) and use the natively-signed `madd_epi16`, so unlike the SPR `dpbusd` path they need none of the +128/-128 sum corrections. L2 is byte-for-byte the unsigned path, since the +128 bias cancels in the difference. Dispatch adds the AVX2 (`d % 16`) and AVX512 (`d % 32`) arms next to the existing unsigned `QT_8bit_direct` case.

`TestSQByte.test_8bit_direct` now runs `d` in {32, 64, 128} on top of {13, 16, 24}. These exercise the new AVX2 and AVX512 kernels and, since the dims also hit `d % 64`, give the shipped SPR kernel its first test coverage. The integer kernels are exact (int32 accumulation), so the test asserts bit-exact equality against an `IndexFlat` baseline.

Measured on AVX2, end-to-end `IndexIVFScalarQuantizer` search, single thread, the int8 kernel against the float path it replaces:

| metric | d | float | int8 | speedup |
| --- | --- | --- | --- | --- |
| L2 | 128 | 42.8 ms | 20.2 ms | 2.12x |
| L2 | 256 | 107.5 ms | 48.7 ms | 2.21x |
| L2 | 256, out of cache | 276.7 ms | 172.2 ms | 1.61x |
| IP | 128 | 31.3 ms | 21.4 ms | 1.46x |
| IP | 256 | 80.9 ms | 50.8 ms | 1.59x |

The L2 speedup holds around 1.6x even once the index no longer fits in cache. This is in line with the ~2.3x that https://github.com/facebookresearch/faiss/issues/5173 reported for the SPR kernel.

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

Reviewed By: alibeklfc

Differential Revision: D113115708

Pulled By: mnorris11

fbshipit-source-id: ccc4fadcc7f15e0203c3351c77834d496a63ecb5

faiss/impl/scalar_quantizer/sq-avx2.cpp

@@ -787,6 +787,68 @@ struct DistanceComputerByte<Similarity, SIMDLevel::AVX2> : SQDistanceComputer {
     }
 };
 
+template <class Similarity>
+struct DistanceComputerByteSigned<Similarity, SIMDLevel::AVX2>
+        : SQDistanceComputer {
+    using Sim = Similarity;
+
+    int d;
+    std::vector<uint8_t> tmp;
+
+    DistanceComputerByteSigned(int d, const std::vector<float>&)
+            : d(d), tmp(d) {}
+
+    int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
+            const {
+        // codes store value + 128. madd_epi16 is signed, so IP unbiases the
+        // bytes before multiplying; for L2 the +128 cancels in the difference.
+        // Only dispatched for d % 16 == 0, so the loop needs no tail.
+        __m256i accu = _mm256_setzero_si256();
+        for (int i = 0; i < d; i += 16) {
+            __m256i c1 = _mm256_cvtepu8_epi16(
+                    _mm_loadu_si128((const __m128i*)(code1 + i)));
+            __m256i c2 = _mm256_cvtepu8_epi16(
+                    _mm_loadu_si128((const __m128i*)(code2 + i)));
+            __m256i prod32;
+            if (Sim::metric_type == METRIC_INNER_PRODUCT) {
+                const __m256i bias = _mm256_set1_epi16(128);
+                c1 = _mm256_sub_epi16(c1, bias);
+                c2 = _mm256_sub_epi16(c2, bias);
+                prod32 = _mm256_madd_epi16(c1, c2);
+            } else {
+                __m256i diff = _mm256_sub_epi16(c1, c2);
+                prod32 = _mm256_madd_epi16(diff, diff);
+            }
+            accu = _mm256_add_epi32(accu, prod32);
+        }
+        __m128i sum = _mm256_extractf128_si256(accu, 0);
+        sum = _mm_add_epi32(sum, _mm256_extractf128_si256(accu, 1));
+        sum = _mm_hadd_epi32(sum, sum);
+        sum = _mm_hadd_epi32(sum, sum);
+        return _mm_cvtsi128_si32(sum);
+    }
+
+    void set_query(const float* x) final {
+        for (int i = 0; i < d; i++) {
+            tmp[i] = uint8_t(int(x[i]) + 128);
+        }
+    }
+
+    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);
+    }
+};
+
 /**********************************************************
  * TurboQuant masked_sum AVX2 specialization
  **********************************************************/

faiss/impl/scalar_quantizer/sq-avx512.cpp

@@ -732,6 +732,61 @@ struct DistanceComputerByte<Similarity, SIMDLevel::AVX512>
     }
 };
 
+template <class Similarity>
+struct DistanceComputerByteSigned<Similarity, SIMDLevel::AVX512>
+        : SQDistanceComputer {
+    using Sim = Similarity;
+
+    int d;
+    std::vector<uint8_t> tmp;
+
+    DistanceComputerByteSigned(int d, const std::vector<float>&)
+            : d(d), tmp(d) {}
+
+    int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
+            const {
+        // codes store value + 128. madd_epi16 is signed, so IP unbiases the
+        // bytes before multiplying; for L2 the +128 cancels in the difference.
+        // Only dispatched for d % 32 == 0, so the loop needs no tail.
+        __m512i accu = _mm512_setzero_si512();
+        constexpr int kLanes = 32;
+        for (int i = 0; i + kLanes <= d; i += kLanes) {
+            __m512i c1 = _mm512_cvtepu8_epi16(_mm256_loadu_epi8(code1 + i));
+            __m512i c2 = _mm512_cvtepu8_epi16(_mm256_loadu_epi8(code2 + i));
+            if (Sim::metric_type == METRIC_INNER_PRODUCT) {
+                const __m512i bias = _mm512_set1_epi16(128);
+                c1 = _mm512_sub_epi16(c1, bias);
+                c2 = _mm512_sub_epi16(c2, bias);
+                accu = _mm512_add_epi32(accu, _mm512_madd_epi16(c1, c2));
+            } else {
+                __m512i diff = _mm512_sub_epi16(c1, c2);
+                accu = _mm512_add_epi32(accu, _mm512_madd_epi16(diff, diff));
+            }
+        }
+        return _mm512_reduce_add_epi32(accu);
+    }
+
+    void set_query(const float* x) final {
+        for (int i = 0; i < d; i++) {
+            tmp[i] = uint8_t(int(x[i]) + 128);
+        }
+    }
+
+    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);
+    }
+};
+
 /**********************************************************
  * TurboQuant masked_sum AVX512 specialization
  **********************************************************/

faiss/impl/scalar_quantizer/sq-dispatch.h

@@ -546,6 +546,16 @@ SQDistanceComputer* select_distance_computer_body(
                     return new DistanceComputerByteSigned<Sim, SL2>(
                             static_cast<int>(d), trained);
                 }
+            } else if constexpr (SL2 == SIMDLevel::AVX512) {
+                if (d % 32 == 0) {
+                    return new DistanceComputerByteSigned<Sim, SL2>(
+                            static_cast<int>(d), trained);
+                }
+            } else if constexpr (SL2 == SIMDLevel::AVX2) {
+                if (d % 16 == 0) {
+                    return new DistanceComputerByteSigned<Sim, SL2>(
+                            static_cast<int>(d), trained);
+                }
             }
             return new DCTemplate<Quantizer8bitDirectSigned<SL2>, Sim, SL2>(
                     d, trained);
@@ -750,6 +760,16 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
                         return scan.template operator()<
                                 DistanceComputerByteSigned<Similarity, SL2>>();
                     }
+                } else if constexpr (SL2 == SIMDLevel::AVX512) {
+                    if (d % 32 == 0) {
+                        return scan.template operator()<
+                                DistanceComputerByteSigned<Similarity, SL2>>();
+                    }
+                } else if constexpr (SL2 == SIMDLevel::AVX2) {
+                    if (d % 16 == 0) {
+                        return scan.template operator()<
+                                DistanceComputerByteSigned<Similarity, SL2>>();
+                    }
                 }
                 return scan.template operator()<DCTemplate<
                         Quantizer8bitDirectSigned<SL2>,

tests/test_index_accuracy.py

@@ -386,7 +386,8 @@ class TestSQByte(unittest.TestCase):
             faiss.ScalarQuantizer.QT_8bit_direct,
             faiss.ScalarQuantizer.QT_8bit_direct_signed,
         ):
-            for d in 13, 16, 24:
+            # d % 16 / 32 / 64 exercise the AVX2 / AVX512 / SPR byte kernels
+            for d in 13, 16, 24, 32, 64, 128:
                 for metric_type in faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT:
                     self.subtest_8bit_direct(metric_type, d, quantizer)