From 3b52da37c40973d4816ff7d2a4e3fb9f74379988 Mon Sep 17 00:00:00 2001 From: Aaryan Mehta <73230976+blazingphoenix7@users.noreply.github.com> Date: Wed, 22 Jul 2026 15:22:53 -0700 Subject: [PATCH] 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`, 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`/`` 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 | 62 +++++++++++++++++++++++ faiss/impl/scalar_quantizer/sq-avx512.cpp | 55 ++++++++++++++++++++ faiss/impl/scalar_quantizer/sq-dispatch.h | 20 ++++++++ tests/test_index_accuracy.py | 3 +- 4 files changed, 139 insertions(+), 1 deletion(-) diff --git a/faiss/impl/scalar_quantizer/sq-avx2.cpp b/faiss/impl/scalar_quantizer/sq-avx2.cpp index 4e2231b9b..3cf97e022 100644 --- a/faiss/impl/scalar_quantizer/sq-avx2.cpp +++ b/faiss/impl/scalar_quantizer/sq-avx2.cpp @@ -787,6 +787,68 @@ struct DistanceComputerByte : SQDistanceComputer { } }; +template +struct DistanceComputerByteSigned + : SQDistanceComputer { + using Sim = Similarity; + + int d; + std::vector tmp; + + DistanceComputerByteSigned(int d, const std::vector&) + : 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 **********************************************************/ diff --git a/faiss/impl/scalar_quantizer/sq-avx512.cpp b/faiss/impl/scalar_quantizer/sq-avx512.cpp index 9e17fa96f..ae06627f8 100644 --- a/faiss/impl/scalar_quantizer/sq-avx512.cpp +++ b/faiss/impl/scalar_quantizer/sq-avx512.cpp @@ -732,6 +732,61 @@ struct DistanceComputerByte } }; +template +struct DistanceComputerByteSigned + : SQDistanceComputer { + using Sim = Similarity; + + int d; + std::vector tmp; + + DistanceComputerByteSigned(int d, const std::vector&) + : 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 **********************************************************/ diff --git a/faiss/impl/scalar_quantizer/sq-dispatch.h b/faiss/impl/scalar_quantizer/sq-dispatch.h index 1b65c8c58..e42285c06 100644 --- a/faiss/impl/scalar_quantizer/sq-dispatch.h +++ b/faiss/impl/scalar_quantizer/sq-dispatch.h @@ -546,6 +546,16 @@ SQDistanceComputer* select_distance_computer_body( return new DistanceComputerByteSigned( static_cast(d), trained); } + } else if constexpr (SL2 == SIMDLevel::AVX512) { + if (d % 32 == 0) { + return new DistanceComputerByteSigned( + static_cast(d), trained); + } + } else if constexpr (SL2 == SIMDLevel::AVX2) { + if (d % 16 == 0) { + return new DistanceComputerByteSigned( + static_cast(d), trained); + } } return new DCTemplate, Sim, SL2>( d, trained); @@ -750,6 +760,16 @@ InvertedListScanner* sq_select_InvertedListScanner( return scan.template operator()< DistanceComputerByteSigned>(); } + } else if constexpr (SL2 == SIMDLevel::AVX512) { + if (d % 32 == 0) { + return scan.template operator()< + DistanceComputerByteSigned>(); + } + } else if constexpr (SL2 == SIMDLevel::AVX2) { + if (d % 16 == 0) { + return scan.template operator()< + DistanceComputerByteSigned>(); + } } return scan.template operator(), diff --git a/tests/test_index_accuracy.py b/tests/test_index_accuracy.py index 349588b47..8049310b0 100644 --- a/tests/test_index_accuracy.py +++ b/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)