faiss 0.5.2 → 0.6.0

data/vendor/faiss/faiss/impl/RaBitQUtils.h

@@ -9,8 +9,10 @@
 
 #include <faiss/MetricType.h>
 #include <faiss/impl/platform_macros.h>
+#include <faiss/utils/AlignedTable.h>
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 #include <vector>
 
 namespace faiss {
@@ -68,6 +70,7 @@ struct QueryFactorsData {
 
     float qr_to_c_L2sqr = 0;
     float qr_norm_L2sqr = 0;
+    float q_dot_c = 0; // <query, centroid> for IP metric; 0 for L2
 
     float int_dot_scale = 1;
 
@@ -239,6 +242,41 @@ inline float compute_1bit_adjusted_distance(
     return adjusted_distance;
 }
 
+/** Determine whether a candidate should be refined in two-stage search.
+ * Consolidates the filtering logic for both L2 and IP metrics.
+ *
+ * For L2 (min-heap): uses lower_bound = est_distance - error_adjustment
+ *   - Skip if lower_bound >= threshold (can't beat current worst)
+ * For IP (max-heap): uses upper_bound = est_distance + error_adjustment
+ *   - Skip if upper_bound <= threshold (can't beat current best)
+ *
+ * @param est_distance     Estimated 1-bit distance
+ * @param f_error          Database vector error factor
+ * @param g_error          Query vector error factor
+ * @param threshold        Current heap threshold (worst result in heap)
+ * @param is_similarity    True for IP metric (max-heap), false for L2
+ * (min-heap)
+ * @return                 True if candidate should be refined with full
+ * multi-bit distance
+ */
+inline bool should_refine_candidate(
+        float est_distance,
+        float f_error,
+        float g_error,
+        float threshold,
+        bool is_similarity) {
+    float error_adjustment = f_error * g_error;
+    if (is_similarity) {
+        // IP (max-heap): use upper bound for filtering
+        float upper_bound = est_distance + error_adjustment;
+        return upper_bound > threshold;
+    } else {
+        // L2 (min-heap): use lower bound for filtering
+        float lower_bound = std::max(0.0f, est_distance - error_adjustment);
+        return lower_bound < threshold;
+    }
+}
+
 /** Extract multi-bit code on-the-fly from packed ex-bit codes.
  * This inline function extracts a single code value without unpacking the
  * entire array, enabling efficient on-the-fly decoding during distance
@@ -276,55 +314,79 @@ inline int extract_code_inline(
  *
  * The multi-bit distance combines the sign bit (1-bit) with additional
  * magnitude bits (ex_bits) to compute a more accurate distance estimate.
+ * Uses SIMD-optimized bit-plane decomposition (AVX2+BMI2) for ex_bits 1-7,
+ * with scalar fallback for non-x86 or non-BMI2 platforms.
  *
  * @param sign_bits       unpacked sign bits (1-bit codes in standard format)
  * @param ex_code         packed ex-bit codes
  * @param ex_fac          ex-bit factors (f_add_ex, f_rescale_ex)
  * @param rotated_q       rotated query vector
- * @param qr_to_c_L2sqr   precomputed ||query_rotated - centroid||^2
- * @param qr_norm_L2sqr   precomputed ||query_rotated||^2 (0 for L2 metric)
+ * @param qr_base         precomputed base term: ||q-c||^2 for L2, <q,c> for IP
  * @param d               dimensionality
  * @param ex_bits         number of extra bits (nb_bits - 1)
  * @param metric_type     distance metric (L2 or Inner Product)
  * @return                computed full multi-bit distance
  */
-inline float compute_full_multibit_distance(
+float compute_full_multibit_distance(
         const uint8_t* sign_bits,
         const uint8_t* ex_code,
         const ExtraBitsFactors& ex_fac,
         const float* rotated_q,
-        float qr_to_c_L2sqr,
-        float qr_norm_L2sqr,
+        float qr_base,
         size_t d,
         size_t ex_bits,
-        MetricType metric_type) {
-    float ex_ip = 0.0f;
-    const float cb = -(static_cast<float>(1 << ex_bits) - 0.5f);
-
-    for (size_t i = 0; i < d; i++) {
-        const size_t byte_idx = i / 8;
-        const size_t bit_offset = i % 8;
-        const bool sign_bit = (sign_bits[byte_idx] >> bit_offset) & 1;
-
-        int ex_code_val = extract_code_inline(ex_code, i, ex_bits);
-
-        int total_code = (sign_bit ? 1 : 0) << ex_bits;
-        total_code += ex_code_val;
-        float reconstructed = static_cast<float>(total_code) + cb;
-
-        ex_ip += rotated_q[i] * reconstructed;
-    }
-
-    float dist = qr_to_c_L2sqr + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;
+        MetricType metric_type);
+
+/** Compute pointer to a vector's auxiliary data within block layout. */
+template <typename T>
+inline T* get_block_aux_ptr(
+        T* block_data,
+        size_t vec_pos,
+        size_t bbs,
+        size_t packed_block_size,
+        size_t full_block_size,
+        size_t storage_size) {
+    return block_data + (vec_pos / bbs) * full_block_size + packed_block_size +
+            (vec_pos % bbs) * storage_size;
+}
 
-    if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
-        dist = -0.5f * (dist - qr_norm_L2sqr);
-    } else {
-        dist = std::max(0.0f, dist);
-    }
+/** Compute per-vector auxiliary storage size.
+ *
+ * @param nb_bits  number of quantization bits (1 = sign-bit only)
+ * @param d        dimensionality
+ * @return         storage size in bytes
+ */
+size_t compute_per_vector_storage_size(size_t nb_bits, size_t d);
 
-    return dist;
-}
+/** [LEGACY FORMAT SUPPORT] Migrate block data from old I/O format to new
+ * format.
+ *
+ * This function is used only when reading indexes saved with the legacy format
+ * (fourcc "Irfs"/"Iwrf") to convert them to the new embedded auxiliary data
+ * format. Not needed for indexes saved with the new format ("Irfn"/"Iwrn").
+ *
+ * Re-layouts blocks in-place and copies aux data from flat_storage.
+ *
+ * @param flat_storage       legacy per-vector aux data indexed by global ID
+ * @param codes              block data (will be resized and re-laid out)
+ * @param num_vectors        number of vectors in this segment
+ * @param bbs                block batch size (vectors per block)
+ * @param M2                 rounded sub-quantizer count
+ * @param old_block_stride   old block size (packed codes only, or current)
+ * @param new_block_stride   new block size (packed codes + aux region)
+ * @param storage_size       per-vector aux storage size in bytes
+ * @param id_map             maps local offset to global ID; null = sequential
+ */
+void populate_block_aux_from_flat_storage(
+        const std::vector<uint8_t>& flat_storage,
+        AlignedTable<uint8_t>& codes,
+        size_t num_vectors,
+        size_t bbs,
+        size_t M2,
+        size_t old_block_stride,
+        size_t new_block_stride,
+        size_t storage_size,
+        const int64_t* id_map = nullptr);
 
 } // namespace rabitq_utils
 } // namespace faiss

data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp

@@ -12,6 +12,7 @@
 #include <faiss/impl/RaBitQuantizerMultiBit.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/rabitq_simd.h>
+
 #include <algorithm>
 #include <cmath>
 #include <cstring>
@@ -63,7 +64,7 @@ size_t RaBitQuantizer::compute_code_size(size_t d, size_t num_bits) const {
     return base_size + ex_size;
 }
 
-void RaBitQuantizer::train(size_t n, const float* x) {
+void RaBitQuantizer::train(size_t /*n*/, const float* /*x*/) {
     // does nothing
 }
 
@@ -215,29 +216,6 @@ void RaBitQuantizer::decode_core(
     }
 }
 
-// Implementation of RaBitQDistanceComputer (declared in header)
-
-float RaBitQDistanceComputer::lower_bound_distance(const uint8_t* code) {
-    FAISS_ASSERT(code != nullptr);
-
-    // Compute estimated distance using 1-bit codes
-    float est_distance = distance_to_code_1bit(code);
-
-    // Extract f_error from the code
-    size_t size = (d + 7) / 8;
-    const SignBitFactorsWithError* base_fac =
-            reinterpret_cast<const SignBitFactorsWithError*>(code + size);
-    float f_error = base_fac->f_error;
-
-    // Compute proper lower bound using RaBitQ error formula:
-    // lower_bound = est_distance - f_error * g_error
-    // This guarantees: lower_bound ≤ true_distance
-    float lower_bound = est_distance - (f_error * g_error);
-
-    // Distance cannot be negative
-    return std::max(0.0f, lower_bound);
-}
-
 namespace {
 
 struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer {
@@ -336,13 +314,15 @@ float RaBitQDistanceComputerNotQ::distance_to_code_full(const uint8_t* code) {
             ex_code + (d * ex_bits + 7) / 8);
 
     // Call shared utility directly with rotated_q pointer
+    float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT)
+            ? query_fac.q_dot_c
+            : query_fac.qr_to_c_L2sqr;
     return rabitq_utils::compute_full_multibit_distance(
             binary_data,
             ex_code,
             *ex_fac,
             rotated_q.data(),
-            query_fac.qr_to_c_L2sqr,
-            query_fac.qr_norm_L2sqr,
+            qr_base,
             d,
             ex_bits,
             metric_type);
@@ -388,6 +368,8 @@ void RaBitQDistanceComputerNotQ::set_query(const float* x) {
     if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
         // precompute if needed
         query_fac.qr_norm_L2sqr = fvec_norm_L2sqr(x, d);
+        query_fac.q_dot_c =
+                centroid ? fvec_inner_product(x, centroid, d) : 0.0f;
     }
 }
 
@@ -502,13 +484,15 @@ float RaBitQDistanceComputerQ::distance_to_code_full(const uint8_t* code) {
             ex_code + (d * ex_bits + 7) / 8);
 
     // Call shared utility directly with rotated_q pointer
+    float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT)
+            ? query_fac.q_dot_c
+            : query_fac.qr_to_c_L2sqr;
     return rabitq_utils::compute_full_multibit_distance(
             binary_data,
             ex_code,
             *ex_fac,
             rotated_q.data(),
-            query_fac.qr_to_c_L2sqr,
-            query_fac.qr_norm_L2sqr,
+            qr_base,
             d,
             ex_bits,
             metric_type);

data/vendor/faiss/faiss/impl/RaBitQuantizer.h

@@ -103,10 +103,8 @@ struct RaBitQuantizer : Quantizer {
 //
 // 1. distance_to_code_1bit() - Fast 1-bit filtering using only sign bits
 // 2. distance_to_code_full() - Accurate multi-bit refinement using all bits
-// 3. lower_bound_distance() - Error-bounded adaptive filtering
-//                              (based on 1-bit estimator)
 //
-// These three methods implement RaBitQ's two-stage search pattern and are
+// These methods implement RaBitQ's two-stage search pattern and are
 // shared between the quantized (Q) and non-quantized (NotQ) query variants.
 // The intermediate class allows two-stage search code to work with both
 // variants via a single dynamic_cast.
@@ -116,8 +114,8 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer {
     MetricType metric_type = MetricType::METRIC_L2;
     size_t nb_bits = 1;
 
-    // Query norm for lower bound computation (g_error in rabitq-library)
-    // This is the L2 norm of the rotated query: ||query - centroid||
+    // Query error factor for bound computation (g_error in rabitq-library)
+    // Used with f_error to compute error bounds for two-stage filtering
     float g_error = 0.0f;
 
     float symmetric_dis(idx_t /*i*/, idx_t /*j*/) override {
@@ -131,11 +129,6 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer {
     // Compute full multi-bit distance (accurate)
     virtual float distance_to_code_full(const uint8_t* code) = 0;
 
-    // Compute lower bound of distance using error bounds
-    // Guarantees: actual_distance >= lower_bound_distance
-    // Used for adaptive filtering in two-stage search
-    virtual float lower_bound_distance(const uint8_t* code);
-
     // Override from FlatCodesDistanceComputer
     // Delegates to distance_to_code_full() for multi-bit distance computation
     float distance_to_code(const uint8_t* code) final {

data/vendor/faiss/faiss/impl/RaBitQuantizerMultiBit.cpp

@@ -180,9 +180,7 @@ void pack_multibit_codes(
  *
  * @param residual Original residual vector (data - centroid)
  * @param centroid Centroid vector (can be nullptr for zero centroid)
- * @param tmp_code Quantized ex-bit codes (before packing, after bit flipping)
  * @param d Dimensionality
- * @param ex_bits Number of extra bits
  * @param norm L2 norm of residual
  * @param ipnorm Unnormalized inner product between quantized and normalized
  * residual
@@ -192,9 +190,7 @@ void pack_multibit_codes(
 void compute_ex_factors(
         const float* residual,
         const float* centroid,
-        const int* tmp_code,
         size_t d,
-        size_t ex_bits,
         float norm,
         double ipnorm,
         ExtraBitsFactors& ex_factors,
@@ -210,45 +206,23 @@ void compute_ex_factors(
         ipnorm_inv = 1.0f;
     }
 
-    // Reconstruct xu_cb from total_code
-    // total_code was formed from: total_code[i] = (sign << ex_bits) +
-    // ex_code[i] Reconstruction: xu_cb[i] = total_code[i] + cb
-    const float cb = -(static_cast<float>(1 << ex_bits) - 0.5f);
-    std::vector<float> xu_cb(d);
-    for (size_t i = 0; i < d; i++) {
-        xu_cb[i] = static_cast<float>(tmp_code[i]) + cb;
-    }
-
     // Compute inner products needed for factors
-    float l2_sqr = norm * norm;
-    float ip_resi_xucb = fvec_inner_product(residual, xu_cb.data(), d);
+    float l2_sqr = norm * norm; // ||residual||^2 = ||x - c||^2
 
-    // Compute factors
     if (metric_type == MetricType::METRIC_L2) {
-        // For L2, no centroid correction needed in IVF setting
-        // because residual = x - centroid, distance computed in residual space
+        // For L2: f_add_ex = ||residual||^2
+        // No centroid correction needed in IVF setting because
+        // residual = x - centroid, distance computed in residual space
         ex_factors.f_add_ex = l2_sqr;
         ex_factors.f_rescale_ex = ipnorm_inv * -2.0f * norm;
     } else {
-        // For IP, centroid correction is needed
-        float ip_resi_cent = 0;
-        if (centroid != nullptr) {
-            ip_resi_cent = fvec_inner_product(residual, centroid, d);
-        }
-
-        float ip_cent_xucb = 0;
-        if (centroid != nullptr) {
-            ip_cent_xucb = fvec_inner_product(centroid, xu_cb.data(), d);
-        }
-
-        // When ip_resi_xucb is zero, the correction term should be zero
-        float correction_term = 0.0f;
-        if (ip_resi_xucb != 0.0f) {
-            correction_term = l2_sqr * ip_cent_xucb / ip_resi_xucb;
-        }
-
-        ex_factors.f_add_ex = 1 - ip_resi_cent + correction_term;
-        ex_factors.f_rescale_ex = ipnorm_inv * -norm;
+        // For IP: direct dot-product formulation
+        // f_add_ex = <c, r> (dot product of centroid and residual)
+        // f_rescale_ex = ||r|| / ipnorm (positive scaling)
+        float c_dot_r =
+                centroid ? fvec_inner_product(residual, centroid, d) : 0.0f;
+        ex_factors.f_add_ex = c_dot_r;
+        ex_factors.f_rescale_ex = ipnorm_inv * norm;
     }
 }
 
@@ -290,12 +264,14 @@ void quantize_ex_bits(
     float norm_sqr = fvec_norm_L2sqr(residual, d);
     float norm = std::sqrt(norm_sqr);
 
-    // Handle degenerate case
+    // Handle degenerate case: residual is (near-)zero, meaning x ≈ centroid.
+    // For both L2 and IP, f_add_ex and f_rescale_ex are trivially zero:
+    //   L2: ||r||² ≈ 0, IP: <c,r> ≈ 0 and ||r||/ipnorm ≈ 0
     if (norm < 1e-10f) {
         size_t code_size = (d * ex_bits + 7) / 8;
         memset(ex_code, 0, code_size);
         ex_factors.f_add_ex = 0.0f;
-        ex_factors.f_rescale_ex = 1.0f;
+        ex_factors.f_rescale_ex = 0.0f;
         return;
     }
 
@@ -349,9 +325,7 @@ void quantize_ex_bits(
     compute_ex_factors(
             residual,
             centroid, // Pass centroid for IP metric factor computation
-            total_code.data(),
             d,
-            ex_bits,
             norm,
             ipnorm,
             ex_factors,

data/vendor/faiss/faiss/impl/RaBitQuantizerMultiBit.h

@@ -60,9 +60,7 @@ void pack_multibit_codes(
  *
  * @param residual Original residual vector (data - centroid)
  * @param centroid Centroid vector (can be nullptr for zero centroid)
- * @param tmp_code Quantized ex-bit codes (unpacked integers)
  * @param d Dimensionality
- * @param ex_bits Number of extra bits
  * @param norm L2 norm of residual
  * @param ipnorm Unnormalized inner product
  * @param ex_factors Output factors structure
@@ -71,9 +69,7 @@ void pack_multibit_codes(
 void compute_ex_factors(
         const float* residual,
         const float* centroid,
-        const int* tmp_code,
         size_t d,
-        size_t ex_bits,
         float norm,
         double ipnorm,
         rabitq_utils::ExtraBitsFactors& ex_factors,

data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp

@@ -18,6 +18,7 @@
 #include <faiss/VectorTransform.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/residual_quantizer_encode_steps.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/utils.h>
@@ -274,10 +275,12 @@ void ResidualQuantizer::train(size_t n, const float* x) {
     // find min and max norms
     std::vector<float> norms(n);
 
-    for (size_t i = 0; i < n; i++) {
-        norms[i] = fvec_L2sqr(
-                x + i * d, residuals.data() + i * cur_beam_size * d, d);
-    }
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (size_t i = 0; i < n; i++) {
+            norms[i] = fvec_L2sqr<SL>(
+                    x + i * d, residuals.data() + i * cur_beam_size * d, d);
+        }
+    });
 
     // fvec_norms_L2sqr(norms.data(), x, d, n);
     train_norm(n, norms.data());
@@ -393,11 +396,13 @@ float ResidualQuantizer::retrain_AQ_codebook(size_t n, const float* x) {
     }
 
     float output_recons_error = 0;
-    for (size_t j = 0; j < d; j++) {
-        output_recons_error += fvec_norm_L2sqr(
-                xt.data() + total_codebook_size + n * j,
-                n - total_codebook_size);
-    }
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (size_t j = 0; j < d; j++) {
+            output_recons_error += fvec_norm_L2sqr<SL>(
+                    xt.data() + total_codebook_size + n * j,
+                    n - total_codebook_size);
+        }
+    });
     if (verbose) {
         printf("  output quantization error %g\n", output_recons_error);
     }