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)