faiss 0.6.1 → 0.6.3

data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp

@@ -9,6 +9,7 @@
 
 #include <faiss/utils/rabitq_simd.h>
 #include <immintrin.h>
+#include <limits>
 
 namespace faiss::rabitq {
 
@@ -167,8 +168,164 @@ inline uint64_t reduce_add_128(__m128i v) {
     return lanes[0] + lanes[1];
 }
 
+inline __m512i round_nonnegative_ps_to_i32(__m512 x) {
+    return _mm512_cvttps_epi32(_mm512_add_ps(x, _mm512_set1_ps(0.5f)));
+}
+
+inline void store_i32_as_u8_16(__m512i values, uint8_t* out) {
+    const __m128i packed = _mm512_cvtusepi32_epi8(values);
+    _mm_storeu_si128(reinterpret_cast<__m128i*>(out), packed);
+}
+
 } // namespace
 
+template <>
+void lut_minmax_16<SIMDLevel::AVX512>(const float* tab, float& mn, float& mx) {
+    const __m512 v = _mm512_loadu_ps(tab);
+    mn = _mm512_reduce_min_ps(v);
+    mx = _mm512_reduce_max_ps(v);
+}
+
+template <>
+void minmax_values<SIMDLevel::AVX512>(
+        const float* values,
+        size_t n,
+        float& mn,
+        float& mx) {
+    if (n == 0) {
+        return;
+    }
+
+    size_t i = 0;
+    __m512 mn_vec = _mm512_set1_ps(std::numeric_limits<float>::max());
+    __m512 mx_vec = _mm512_set1_ps(std::numeric_limits<float>::lowest());
+    for (; i + 16 <= n; i += 16) {
+        const __m512 v = _mm512_loadu_ps(values + i);
+        mn_vec = _mm512_min_ps(mn_vec, v);
+        mx_vec = _mm512_max_ps(mx_vec, v);
+    }
+
+    if (i < n) {
+        const __mmask16 mask =
+                static_cast<__mmask16>((uint32_t(1) << (n - i)) - 1);
+        const __m512 v = _mm512_maskz_loadu_ps(mask, values + i);
+        mn_vec = _mm512_mask_min_ps(mn_vec, mask, mn_vec, v);
+        mx_vec = _mm512_mask_max_ps(mx_vec, mask, mx_vec, v);
+    }
+
+    mn = _mm512_reduce_min_ps(mn_vec);
+    mx = _mm512_reduce_max_ps(mx_vec);
+}
+
+template <>
+void lut_quantize_16_to_uint8<SIMDLevel::AVX512>(
+        const float* tab,
+        float mn,
+        float a,
+        uint8_t* out) {
+    const __m512 values = _mm512_loadu_ps(tab);
+    const __m512 a_vec = _mm512_set1_ps(a);
+    const __m512 scaled =
+            _mm512_fmsub_ps(values, a_vec, _mm512_set1_ps(mn * a));
+    __m512i rounded = round_nonnegative_ps_to_i32(scaled);
+    rounded = _mm512_max_epi32(rounded, _mm512_setzero_si512());
+    store_i32_as_u8_16(rounded, out);
+}
+
+template <>
+void quantize_query_values<SIMDLevel::AVX512>(
+        const float* rq,
+        size_t d,
+        float v_min,
+        float inv_delta,
+        uint8_t max_code,
+        bool centered,
+        uint8_t* rqq,
+        size_t& sum_qq,
+        int64_t& sum2_signed_odd_int) {
+    const __m512 inv_delta_vec = _mm512_set1_ps(inv_delta);
+    const __m512 v_min_times_inv_delta_vec = _mm512_set1_ps(v_min * inv_delta);
+    const __m512 zero = _mm512_setzero_ps();
+    const __m512 max_code_ps = _mm512_set1_ps(static_cast<float>(max_code));
+    const __m512i max_code_i32 = _mm512_set1_epi32(max_code);
+    const __m512i two = _mm512_set1_epi32(2);
+
+    size_t i = 0;
+    if (centered) {
+        __m512i sum_acc_lo = _mm512_setzero_si512();
+        __m512i sum_acc_hi = _mm512_setzero_si512();
+        __m512i sq_acc_lo = _mm512_setzero_si512();
+        __m512i sq_acc_hi = _mm512_setzero_si512();
+        for (; i + 16 <= d; i += 16) {
+            const __m512 values = _mm512_loadu_ps(rq + i);
+            __m512 scaled = _mm512_fmsub_ps(
+                    values, inv_delta_vec, v_min_times_inv_delta_vec);
+            scaled = _mm512_min_ps(_mm512_max_ps(scaled, zero), max_code_ps);
+            __m512i rounded = round_nonnegative_ps_to_i32(scaled);
+
+            sum_acc_lo = _mm512_add_epi64(
+                    sum_acc_lo,
+                    _mm512_cvtepi32_epi64(_mm512_castsi512_si256(rounded)));
+            sum_acc_hi = _mm512_add_epi64(
+                    sum_acc_hi,
+                    _mm512_cvtepi32_epi64(
+                            _mm512_extracti64x4_epi64(rounded, 1)));
+            const __m512i signed_odd = _mm512_sub_epi32(
+                    _mm512_mullo_epi32(rounded, two), max_code_i32);
+            const __m512i signed_odd_sqr =
+                    _mm512_mullo_epi32(signed_odd, signed_odd);
+            sq_acc_lo = _mm512_add_epi64(
+                    sq_acc_lo,
+                    _mm512_cvtepi32_epi64(
+                            _mm512_castsi512_si256(signed_odd_sqr)));
+            sq_acc_hi = _mm512_add_epi64(
+                    sq_acc_hi,
+                    _mm512_cvtepi32_epi64(
+                            _mm512_extracti64x4_epi64(signed_odd_sqr, 1)));
+            store_i32_as_u8_16(rounded, rqq + i);
+        }
+        sum_qq += static_cast<uint64_t>(
+                _mm512_reduce_add_epi64(sum_acc_lo) +
+                _mm512_reduce_add_epi64(sum_acc_hi));
+        sum2_signed_odd_int += _mm512_reduce_add_epi64(sq_acc_lo);
+        sum2_signed_odd_int += _mm512_reduce_add_epi64(sq_acc_hi);
+    } else {
+        __m512i sum_acc_lo = _mm512_setzero_si512();
+        __m512i sum_acc_hi = _mm512_setzero_si512();
+        for (; i + 16 <= d; i += 16) {
+            const __m512 values = _mm512_loadu_ps(rq + i);
+            __m512 scaled = _mm512_fmsub_ps(
+                    values, inv_delta_vec, v_min_times_inv_delta_vec);
+            scaled = _mm512_min_ps(_mm512_max_ps(scaled, zero), max_code_ps);
+            __m512i rounded = round_nonnegative_ps_to_i32(scaled);
+
+            sum_acc_lo = _mm512_add_epi64(
+                    sum_acc_lo,
+                    _mm512_cvtepi32_epi64(_mm512_castsi512_si256(rounded)));
+            sum_acc_hi = _mm512_add_epi64(
+                    sum_acc_hi,
+                    _mm512_cvtepi32_epi64(
+                            _mm512_extracti64x4_epi64(rounded, 1)));
+            store_i32_as_u8_16(rounded, rqq + i);
+        }
+        sum_qq += static_cast<uint64_t>(
+                _mm512_reduce_add_epi64(sum_acc_lo) +
+                _mm512_reduce_add_epi64(sum_acc_hi));
+    }
+
+    for (; i < d; i++) {
+        const uint8_t v_qq = round_clamped_byte_scalar(
+                (rq[i] - v_min) * inv_delta, max_code);
+        rqq[i] = v_qq;
+        sum_qq += v_qq;
+
+        if (centered) {
+            const int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
+            sum2_signed_odd_int += signed_odd_int * signed_odd_int;
+        }
+    }
+}
+
 template <>
 uint64_t bitwise_and_dot_product<SIMDLevel::AVX512>(
         const uint8_t* query,
@@ -237,6 +394,87 @@ uint64_t bitwise_and_dot_product<SIMDLevel::AVX512>(
     return sum;
 }
 
+template <>
+BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
+        SIMDLevel::AVX512>(
+        const uint8_t* query,
+        const uint8_t* data,
+        size_t size,
+        size_t qb) {
+    uint64_t dot_product = 0;
+    uint64_t popcount_sum = 0;
+    size_t offset = 0;
+    if (size_t step = 512 / 8; offset + step <= size) {
+        __m512i dot_512 = _mm512_setzero_si512();
+        __m512i pop_512 = _mm512_setzero_si512();
+        for (; offset + step <= size; offset += step) {
+            __m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset));
+            pop_512 = _mm512_add_epi64(pop_512, popcount_512(v_x));
+            for (int j = 0; j < qb; j++) {
+                __m512i v_q = _mm512_loadu_si512(
+                        (const __m512i*)(query + j * size + offset));
+                __m512i v_and = _mm512_and_si512(v_q, v_x);
+                __m512i v_popcnt = popcount_512(v_and);
+                __m512i v_shifted = _mm512_slli_epi64(v_popcnt, j);
+                dot_512 = _mm512_add_epi64(dot_512, v_shifted);
+            }
+        }
+        dot_product += _mm512_reduce_add_epi64(dot_512);
+        popcount_sum += _mm512_reduce_add_epi64(pop_512);
+    }
+    if (size_t step = 256 / 8; offset + step <= size) {
+        __m256i dot_256 = _mm256_setzero_si256();
+        __m256i pop_256 = _mm256_setzero_si256();
+        for (; offset + step <= size; offset += step) {
+            __m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset));
+            pop_256 = _mm256_add_epi64(pop_256, popcount_256(v_x));
+            for (int j = 0; j < qb; j++) {
+                __m256i v_q = _mm256_loadu_si256(
+                        (const __m256i*)(query + j * size + offset));
+                __m256i v_and = _mm256_and_si256(v_q, v_x);
+                __m256i v_popcnt = popcount_256(v_and);
+                __m256i v_shifted = _mm256_slli_epi64(v_popcnt, j);
+                dot_256 = _mm256_add_epi64(dot_256, v_shifted);
+            }
+        }
+        dot_product += reduce_add_256(dot_256);
+        popcount_sum += reduce_add_256(pop_256);
+    }
+    __m128i dot_128 = _mm_setzero_si128();
+    __m128i pop_128 = _mm_setzero_si128();
+    for (size_t step = 128 / 8; offset + step <= size; offset += step) {
+        __m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset));
+        pop_128 = _mm_add_epi64(pop_128, popcount_128(v_x));
+        for (int j = 0; j < qb; j++) {
+            __m128i v_q = _mm_loadu_si128(
+                    (const __m128i*)(query + j * size + offset));
+            __m128i v_and = _mm_and_si128(v_q, v_x);
+            __m128i v_popcnt = popcount_128(v_and);
+            __m128i v_shifted = _mm_slli_epi64(v_popcnt, j);
+            dot_128 = _mm_add_epi64(dot_128, v_shifted);
+        }
+    }
+    dot_product += reduce_add_128(dot_128);
+    popcount_sum += reduce_add_128(pop_128);
+    for (size_t step = 64 / 8; offset + step <= size; offset += step) {
+        const auto yv = *(const uint64_t*)(data + offset);
+        popcount_sum += popcount64(yv);
+        for (int j = 0; j < qb; j++) {
+            const auto qv = *(const uint64_t*)(query + j * size + offset);
+            dot_product += popcount64(qv & yv) << j;
+        }
+    }
+    for (; offset < size; ++offset) {
+        const auto yv = *(data + offset);
+        popcount_sum += popcount32(yv);
+        for (int j = 0; j < qb; j++) {
+            const auto qv = *(query + j * size + offset);
+            dot_product += popcount32(qv & yv) << j;
+        }
+    }
+    return {dot_product, popcount_sum};
+}
+
 template <>
 uint64_t bitwise_xor_dot_product<SIMDLevel::AVX512>(
         const uint8_t* query,
@@ -344,6 +582,41 @@ uint64_t popcount<SIMDLevel::AVX512>(const uint8_t* data, size_t size) {
     return sum;
 }
 
+template <>
+void rearrange_bit_planes<SIMDLevel::AVX512>(
+        const uint8_t* rotated_qq,
+        size_t d,
+        size_t qb,
+        uint8_t* out) {
+    const size_t offset = (d + 7) / 8;
+    memset(out, 0, offset * qb);
+    size_t idim = 0;
+    for (; idim + 64 <= d; idim += 64) {
+        __m512i vals = _mm512_loadu_si512((const __m512i*)(rotated_qq + idim));
+        for (size_t iv = 0; iv < qb; iv++) {
+            __m512i mask = _mm512_set1_epi8(static_cast<char>(1 << iv));
+            __mmask64 bits = _mm512_test_epi8_mask(vals, mask);
+            memcpy(&out[iv * offset + idim / 8], &bits, 8);
+        }
+    }
+    for (; idim + 32 <= d; idim += 32) {
+        __m256i vals = _mm256_loadu_si256((const __m256i*)(rotated_qq + idim));
+        for (size_t iv = 0; iv < qb; iv++) {
+            __m256i mask = _mm256_set1_epi8(static_cast<char>(1 << iv));
+            __m256i bits =
+                    _mm256_cmpeq_epi8(_mm256_and_si256(vals, mask), mask);
+            uint32_t packed = static_cast<uint32_t>(_mm256_movemask_epi8(bits));
+            memcpy(&out[iv * offset + idim / 8], &packed, 4);
+        }
+    }
+    for (; idim < d; idim++) {
+        for (size_t iv = 0; iv < qb; iv++) {
+            const bool bit = ((rotated_qq[idim] & (1 << iv)) != 0);
+            out[iv * offset + idim / 8] |= bit ? (1 << (idim % 8)) : 0;
+        }
+    }
+}
+
 } // namespace faiss::rabitq
 
 namespace faiss::rabitq::multibit {