faiss 0.6.2 → 0.6.4

data/vendor/faiss/faiss/utils/hamming_distance/hamming_impl.h

@@ -36,6 +36,7 @@
 
 #include <algorithm>
 #include <cstdio>
+#include <cstring>
 #include <limits>
 #include <memory>
 #include <vector>
@@ -148,6 +149,7 @@ void hammings_knn_mc_impl(
     std::unique_ptr<int64_t[]> all_ids_per_dis(new int64_t[na * nBuckets * k]);
 
     std::vector<HCounterState<HammingComputer>> cs;
+    cs.reserve(na);
     for (size_t i = 0; i < na; ++i) {
         cs.push_back(
                 HCounterState<HammingComputer>(
@@ -298,6 +300,132 @@ void generalized_hammings_knn_hc_impl(
     }
 }
 
+/* Both number of words and remainder are constants, so the compiler flattens
+ * the per-chunk loop. The query is copied into qw/qt once, not re-read for
+ * every candidate. */
+template <int NW, int REM>
+void hammings_ragged_fixed(
+        const uint8_t* __restrict a,
+        const uint8_t* __restrict b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* __restrict dis) {
+    for (size_t i = 0; i < na; i++) {
+        const uint8_t* __restrict ai = a + i * ncodes;
+        hamdis_t* __restrict dis_ = dis + i * nb;
+        // Read the query once here, not once per candidate.
+        uint64_t qw[NW == 0 ? 1 : NW];
+        for (int k = 0; k < NW; k++) {
+            memcpy(&qw[k], ai + k * 8, 8);
+        }
+        uint64_t qt = 0;
+        if constexpr (REM > 0) {
+            memcpy(&qt, ai + NW * 8, REM);
+        }
+        for (size_t j = 0; j < nb; j++) {
+            const uint8_t* __restrict bj = b + j * ncodes;
+            hamdis_t h = 0;
+            uint64_t y;
+            for (int k = 0; k < NW; k++) {
+                memcpy(&y, bj + k * 8, 8);
+                h += popcount64(qw[k] ^ y);
+            }
+            if constexpr (REM > 0) {
+                y = 0;
+                memcpy(&y, bj + NW * 8, REM);
+                h += popcount64(qt ^ y);
+            }
+            dis_[j] = h;
+        }
+    }
+}
+
+/* Same loop with the word count read at run time, for codes long enough that
+ * a constant count gives no performance improvement. */
+template <int REM>
+void hammings_ragged_var(
+        const uint8_t* __restrict a,
+        const uint8_t* __restrict b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* __restrict dis) {
+    const size_t nwords = ncodes / 8;
+    for (size_t i = 0; i < na; i++) {
+        const uint8_t* __restrict ai = a + i * ncodes;
+        hamdis_t* __restrict dis_ = dis + i * nb;
+        for (size_t j = 0; j < nb; j++) {
+            const uint8_t* __restrict bj = b + j * ncodes;
+            hamdis_t h = 0;
+            uint64_t x, y;
+            for (size_t k = 0; k < nwords; k++) {
+                memcpy(&x, ai + k * 8, 8);
+                memcpy(&y, bj + k * 8, 8);
+                h += popcount64(x ^ y);
+            }
+            if constexpr (REM > 0) {
+                x = 0;
+                y = 0;
+                memcpy(&x, ai + nwords * 8, REM);
+                memcpy(&y, bj + nwords * 8, REM);
+                h += popcount64(x ^ y);
+            }
+            dis_[j] = h;
+        }
+    }
+}
+
+/* Makes the word count a constant for short codes when calling
+ * hammings_ragged_fixed, worth roughly 2x perf improvement. Longer codes
+ * gain nothing from it, so they share hammings_ragged_var. */
+template <int REM>
+void hammings_ragged_by_nwords(
+        const uint8_t* __restrict a,
+        const uint8_t* __restrict b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* __restrict dis) {
+    switch (ncodes / 8) {
+        case 0:
+            return hammings_ragged_fixed<0, REM>(a, b, na, nb, ncodes, dis);
+        case 1:
+            return hammings_ragged_fixed<1, REM>(a, b, na, nb, ncodes, dis);
+        case 2:
+            return hammings_ragged_fixed<2, REM>(a, b, na, nb, ncodes, dis);
+        case 3:
+            return hammings_ragged_fixed<3, REM>(a, b, na, nb, ncodes, dis);
+        default:
+            return hammings_ragged_var<REM>(a, b, na, nb, ncodes, dis);
+    }
+}
+
+void hammings_ragged_dispatch(
+        const uint8_t* __restrict a,
+        const uint8_t* __restrict b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* __restrict dis) {
+    switch (ncodes % 8) {
+        case 1:
+            return hammings_ragged_by_nwords<1>(a, b, na, nb, ncodes, dis);
+        case 2:
+            return hammings_ragged_by_nwords<2>(a, b, na, nb, ncodes, dis);
+        case 3:
+            return hammings_ragged_by_nwords<3>(a, b, na, nb, ncodes, dis);
+        case 4:
+            return hammings_ragged_by_nwords<4>(a, b, na, nb, ncodes, dis);
+        case 5:
+            return hammings_ragged_by_nwords<5>(a, b, na, nb, ncodes, dis);
+        case 6:
+            return hammings_ragged_by_nwords<6>(a, b, na, nb, ncodes, dis);
+        default:
+            return hammings_ragged_by_nwords<7>(a, b, na, nb, ncodes, dis);
+    }
+}
+
 } // anonymous namespace
 
 /******************************************************************
@@ -356,6 +484,19 @@ void hamming_range_search_fixSL<THE_SIMD_LEVEL>(
             });
 }
 
+/* Its own entry point, so the word-multiple kernels keep the code the
+ * compiler already generates for them. */
+template <>
+void hammings_ragged_fixSL<THE_SIMD_LEVEL>(
+        const uint8_t* a,
+        const uint8_t* b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* dis) {
+    hammings_ragged_dispatch(a, b, na, nb, ncodes, dis);
+}
+
 template <>
 void hammings_fixSL<THE_SIMD_LEVEL>(
         const uint8_t* a,

data/vendor/faiss/faiss/utils/quantize_lut.cpp

@@ -75,7 +75,7 @@ void round_uint8_per_column(
             max_span = span;
         }
     }
-    float a = 255 / max_span;
+    float a = max_span > 0 ? 255.0f / max_span : 0.0f;
     float b = 0;
     for (size_t i = 0; i < n; i++) {
         b += mins[i];
@@ -111,7 +111,7 @@ void round_uint8_per_column_multi(
             max_span = span;
         }
     }
-    float a = 255 / max_span;
+    float a = max_span > 0 ? 255.0f / max_span : 0.0f;
     float b = 0;
     for (size_t i = 0; i < n; i++) {
         b += mins[i];
@@ -143,7 +143,8 @@ void quantize_LUT_and_bias(
         float* b_out) {
     float a, b;
     if (!bias) {
-        FAISS_THROW_IF_NOT(!lut_is_3d);
+        FAISS_THROW_IF_MSG(
+                lut_is_3d, "3d LUT is not supported when bias is null");
         std::vector<float> mins(M);
         float max_span_LUT = -HUGE_VAL, max_span_dis = 0;
         b = 0;
@@ -154,7 +155,12 @@ void quantize_LUT_and_bias(
             max_span_dis += span;
             b += mins[i];
         }
-        a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
+        a = std::min(
+                max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF,
+                max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF);
+        if (!std::isfinite(a)) {
+            a = 0.0f;
+        }
 
         for (size_t i = 0; i < M; i++) {
             round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub);
@@ -174,7 +180,12 @@ void quantize_LUT_and_bias(
             max_span_dis += span;
             b += mins[i];
         }
-        a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
+        a = std::min(
+                max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF,
+                max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF);
+        if (!std::isfinite(a)) {
+            a = 0.0f;
+        }
         b += bias_min;
 
         for (size_t i = 0; i < M; i++) {
@@ -208,7 +219,12 @@ void quantize_LUT_and_bias(
             b = std::min(b, b2j);
         }
 
-        a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
+        a = std::min(
+                max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF,
+                max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF);
+        if (!std::isfinite(a)) {
+            a = 0.0f;
+        }
 
         ij = 0;
         size_t ij_2 = 0;
@@ -256,7 +272,7 @@ void quantize_LUT_and_bias(
             max_span = std::max(max_span, span);
             b += mins[i];
         }
-        a = 255 / max_span;
+        a = max_span > 0 ? 255.0f / max_span : 0.0f;
         ij = 0;
         size_t ij_2 = 0;
         for (size_t j = 0; j < nprobe; j++) {
@@ -305,7 +321,12 @@ void aq_quantize_LUT_and_bias(
         max_span_dis += (i >= M - M_norm ? span * norm_scale : span);
         b += mins[i];
     }
-    a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
+    a = std::min(
+            max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF,
+            max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF);
+    if (!std::isfinite(a)) {
+        a = 0.0f;
+    }
     b += bias_min;
 
     for (size_t i = 0; i < M; i++) {

data/vendor/faiss/faiss/utils/rabitq_simd.h

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
@@ -16,6 +17,15 @@
 
 namespace faiss::rabitq {
 
+/// SIMD levels with RaBitQ query/LUT quantization implementations.
+constexpr int RABITQ_QUANTIZATION_SIMD_LEVELS = (1 << int(SIMDLevel::NONE)) |
+        (1 << int(SIMDLevel::AVX2)) | (1 << int(SIMDLevel::AVX512));
+
+struct BitwiseAndDotProductResult {
+    uint64_t dot_product;
+    uint64_t popcount;
+};
+
 /**
  * Compute dot product between query and binary data using popcount on AND.
  *
@@ -32,6 +42,17 @@ uint64_t bitwise_and_dot_product(
         size_t size,
         size_t qb);
 
+/**
+ * Compute bitwise_and_dot_product(query, data, size, qb) and popcount(data,
+ * size) in one pass over data.
+ */
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
+BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount(
+        const uint8_t* query,
+        const uint8_t* data,
+        size_t size,
+        size_t qb);
+
 /**
  * Compute dot product between query and binary data using popcount on XOR.
  *
@@ -58,6 +79,56 @@ uint64_t bitwise_xor_dot_product(
 template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
 uint64_t popcount(const uint8_t* data, size_t size);
 
+/**
+ * Rearrange per-dimension quantized query codes into bit-plane layout.
+ *
+ * @p rotated_qq holds one qb-bit code per dimension (one byte each, value in
+ * [0, 2^qb)). @p out receives qb bit-planes of ((d + 7) / 8) bytes each:
+ * bit-plane j packs bit j of code i into bit (i % 8) of byte (i / 8). @p out
+ * must have room for qb * ((d + 7) / 8) bytes and is fully overwritten.
+ *
+ * @param rotated_qq  per-dimension codes (d bytes)
+ * @param d           dimensionality
+ * @param qb          quantization bits per dimension (1..8)
+ * @param out         bit-plane output buffer (qb * ((d + 7) / 8) bytes)
+ */
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
+void rearrange_bit_planes(
+        const uint8_t* rotated_qq,
+        size_t d,
+        size_t qb,
+        uint8_t* out);
+
+/// Find min/max of one 16-entry FastScan LUT row.
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
+void lut_minmax_16(const float* tab, float& mn, float& mx);
+
+/// Find min/max of an arbitrary-length float vector. For n == 0, mn and mx
+/// are left unchanged.
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
+void minmax_values(const float* values, size_t n, float& mn, float& mx);
+
+/// Quantize one 16-entry FastScan LUT row with non-negative half-up rounding.
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
+void lut_quantize_16_to_uint8(
+        const float* tab,
+        float mn,
+        float a,
+        uint8_t* out);
+
+/// Quantize rotated query values and accumulate query correction terms.
+template <SIMDLevel SL = SINGLE_SIMD_LEVEL>
+void quantize_query_values(
+        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);
+
 // NONE specializations — scalar fallbacks
 
 template <>
@@ -85,6 +156,35 @@ inline uint64_t bitwise_and_dot_product<SIMDLevel::NONE>(
     return sum;
 }
 
+template <>
+inline BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount<
+        SIMDLevel::NONE>(
+        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;
+    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 <>
 inline uint64_t bitwise_xor_dot_product<SIMDLevel::NONE>(
         const uint8_t* query,
@@ -125,6 +225,108 @@ inline uint64_t popcount<SIMDLevel::NONE>(const uint8_t* data, size_t size) {
     return sum;
 }
 
+template <>
+inline void rearrange_bit_planes<SIMDLevel::NONE>(
+        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);
+    for (size_t idim = 0; 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;
+        }
+    }
+}
+
+inline uint8_t round_clamped_byte_scalar(float x, uint8_t max_code) {
+    if (x <= 0.0f) {
+        return 0;
+    }
+    if (x >= max_code) {
+        return max_code;
+    }
+    return static_cast<uint8_t>(static_cast<int>(x + 0.5f));
+}
+
+inline uint8_t round_nonnegative_byte_scalar(float x) {
+    return round_clamped_byte_scalar(x, 255);
+}
+
+template <>
+inline void lut_minmax_16<SIMDLevel::NONE>(
+        const float* tab,
+        float& mn,
+        float& mx) {
+    mn = tab[0];
+    mx = tab[0];
+    for (size_t s = 1; s < 16; s++) {
+        mn = std::min(mn, tab[s]);
+        mx = std::max(mx, tab[s]);
+    }
+}
+
+template <>
+inline void minmax_values<SIMDLevel::NONE>(
+        const float* values,
+        size_t n,
+        float& mn,
+        float& mx) {
+    if (n == 0) {
+        return;
+    }
+    mn = values[0];
+    mx = values[0];
+    for (size_t i = 1; i < n; i++) {
+        mn = std::min(mn, values[i]);
+        mx = std::max(mx, values[i]);
+    }
+}
+
+template <>
+inline void lut_quantize_16_to_uint8<SIMDLevel::NONE>(
+        const float* tab,
+        float mn,
+        float a,
+        uint8_t* out) {
+    for (size_t s = 0; s < 16; s++) {
+        out[s] = round_nonnegative_byte_scalar(a * (tab[s] - mn));
+    }
+}
+
+template <>
+inline void quantize_query_values<SIMDLevel::NONE>(
+        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) {
+    if (centered) {
+        for (size_t i = 0; 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;
+
+            const int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
+            sum2_signed_odd_int += signed_odd_int * signed_odd_int;
+        }
+    } else {
+        for (size_t i = 0; 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;
+        }
+    }
+}
+
 } // namespace faiss::rabitq
 
 /*********************************************************