faiss 0.5.0 → 0.5.1

data/vendor/faiss/faiss/IndexIVFFlatPanorama.cpp

@@ -70,29 +70,8 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
     float q_norm = 0.0f;
     void set_query(const float* query) override {
         this->xi = query;
-
-        const size_t d = vd.d;
-        const size_t level_width_floats = storage->level_width / sizeof(float);
-
-        std::vector<float> suffix_sums(d + 1);
-        suffix_sums[d] = 0.0f;
-
-        for (int j = d - 1; j >= 0; j--) {
-            float squared_val = query[j] * query[j];
-            suffix_sums[j] = suffix_sums[j + 1] + squared_val;
-        }
-
-        for (size_t level = 0; level < storage->n_levels; level++) {
-            size_t start_idx = level * level_width_floats;
-            if (start_idx < d) {
-                cum_sums[level] = sqrt(suffix_sums[start_idx]);
-            } else {
-                cum_sums[level] = 0.0f;
-            }
-        }
-
-        cum_sums[storage->n_levels] = 0.0f;
-        q_norm = suffix_sums[0];
+        this->storage->pano.compute_query_cum_sums(query, cum_sums.data());
+        q_norm = cum_sums[0] * cum_sums[0];
     }
 
     void set_list(idx_t list_no, float /* coarse_dis */) override {
@@ -107,95 +86,6 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
                 "IndexIVFFlatPanorama does not support distance_to_code");
     }
 
-    /// Helper function for progressive filtering that both scan_codes and
-    /// scan_codes_range use. Processes a batch of vectors through all levels,
-    /// computing exact distances and pruning based on a threshold.
-    /// Returns the number of active survivors after all levels.
-    size_t progressive_filter_batch(
-            size_t batch_no,
-            size_t list_size,
-            const uint8_t* codes_base,
-            const float* cum_sums_data,
-            float threshold,
-            std::vector<float>& exact_distances,
-            std::vector<uint32_t>& active_indices,
-            const idx_t* ids,
-            PanoramaStats& local_stats) const {
-        const size_t d = vd.d;
-        const size_t level_width_floats = storage->level_width / sizeof(float);
-
-        size_t batch_start = batch_no * storage->kBatchSize;
-        size_t curr_batch_size =
-                std::min(list_size - batch_start, storage->kBatchSize);
-
-        size_t cumsum_batch_offset =
-                batch_no * storage->kBatchSize * (storage->n_levels + 1);
-        const float* batch_cum_sums = cum_sums_data + cumsum_batch_offset;
-
-        size_t batch_offset = batch_no * storage->kBatchSize * code_size;
-        const uint8_t* storage_base = codes_base + batch_offset;
-
-        // Initialize active set with ID-filtered vectors.
-        size_t num_active = 0;
-        for (size_t i = 0; i < curr_batch_size; i++) {
-            size_t global_idx = batch_start + i;
-            bool include = !use_sel || sel->is_member(ids[global_idx]);
-
-            active_indices[num_active] = i;
-            float cum_sum = batch_cum_sums[i];
-            exact_distances[i] = cum_sum * cum_sum + q_norm;
-
-            num_active += include;
-        }
-
-        if (num_active == 0) {
-            return 0;
-        }
-
-        size_t total_active = num_active;
-
-        const float* level_cum_sums = batch_cum_sums + storage->kBatchSize;
-
-        // Progressive filtering through levels.
-        for (size_t level = 0; level < storage->n_levels; level++) {
-            local_stats.total_dims_scanned += num_active;
-            local_stats.total_dims += total_active;
-
-            float query_cum_norm = cum_sums[level + 1];
-
-            size_t level_offset =
-                    level * storage->level_width * storage->kBatchSize;
-            const float* level_storage =
-                    (const float*)(storage_base + level_offset);
-
-            size_t next_active = 0;
-            for (size_t i = 0; i < num_active; i++) {
-                uint32_t idx = active_indices[i];
-                const float* yj = level_storage + idx * level_width_floats;
-                const float* query_level = xi + level * level_width_floats;
-
-                size_t actual_level_width = std::min(
-                        level_width_floats, d - level * level_width_floats);
-                float dot_product =
-                        fvec_inner_product(query_level, yj, actual_level_width);
-
-                exact_distances[idx] -= 2.0f * dot_product;
-
-                float cum_sum = level_cum_sums[idx];
-                float cauchy_schwarz_bound = 2.0f * cum_sum * query_cum_norm;
-                float lower_bound = exact_distances[idx] - cauchy_schwarz_bound;
-
-                active_indices[next_active] = idx;
-                next_active += C::cmp(threshold, lower_bound) ? 1 : 0;
-            }
-
-            num_active = next_active;
-            level_cum_sums += storage->kBatchSize;
-        }
-
-        return num_active;
-    }
-
     size_t scan_codes(
             size_t list_size,
             const uint8_t* codes,
@@ -208,7 +98,6 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
         const size_t n_batches =
                 (list_size + storage->kBatchSize - 1) / storage->kBatchSize;
 
-        const uint8_t* codes_base = codes;
         const float* cum_sums_data = storage->get_cum_sums(list_no);
 
         std::vector<float> exact_distances(storage->kBatchSize);
@@ -217,34 +106,25 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
         PanoramaStats local_stats;
         local_stats.reset();
 
-        // Panorama's IVFFlat core progressive filtering algorithm:
-        // Process vectors in batches for cache efficiency. For each batch:
-        // 1. Apply ID selection filter and initialize distances
-        // (||y||^2 + ||x||^2).
-        // 2. Maintain an "active set" of candidate indices that haven't been
-        //    pruned yet.
-        // 3. For each level, refine distances incrementally and compact the
-        //    active set:
-        //    - Compute dot product for current level: exact_dist -= 2*<x,y>.
-        //    - Use Cauchy-Schwarz bound on remaining levels to get lower bound
-        //    - Prune candidates whose lower bound exceeds k-th best distance.
-        //    - Compact active_indices to remove pruned candidates (branchless)
-        // 4. After all levels, survivors are exact distances; update heap.
-        // This achieves early termination while maintaining SIMD-friendly
-        // sequential access patterns in the level-oriented storage layout.
         for (size_t batch_no = 0; batch_no < n_batches; batch_no++) {
             size_t batch_start = batch_no * storage->kBatchSize;
 
-            size_t num_active = progressive_filter_batch(
-                    batch_no,
-                    list_size,
-                    codes_base,
-                    cum_sums_data,
-                    simi[0],
-                    exact_distances,
-                    active_indices,
-                    ids,
-                    local_stats);
+            size_t num_active =
+                    storage->pano
+                            .progressive_filter_batch<CMax<float, int64_t>>(
+                                    codes,
+                                    cum_sums_data,
+                                    xi,
+                                    cum_sums.data(),
+                                    batch_no,
+                                    list_size,
+                                    sel,
+                                    ids,
+                                    use_sel,
+                                    active_indices,
+                                    exact_distances,
+                                    simi[0],
+                                    local_stats);
 
             // Add batch survivors to heap.
             for (size_t i = 0; i < num_active; i++) {
@@ -274,7 +154,6 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
         const size_t n_batches =
                 (list_size + storage->kBatchSize - 1) / storage->kBatchSize;
 
-        const uint8_t* codes_base = codes;
         const float* cum_sums_data = storage->get_cum_sums(list_no);
 
         std::vector<float> exact_distances(storage->kBatchSize);
@@ -288,16 +167,22 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
         for (size_t batch_no = 0; batch_no < n_batches; batch_no++) {
             size_t batch_start = batch_no * storage->kBatchSize;
 
-            size_t num_active = progressive_filter_batch(
-                    batch_no,
-                    list_size,
-                    codes_base,
-                    cum_sums_data,
-                    radius,
-                    exact_distances,
-                    active_indices,
-                    ids,
-                    local_stats);
+            size_t num_active =
+                    storage->pano
+                            .progressive_filter_batch<CMax<float, int64_t>>(
+                                    codes,
+                                    cum_sums_data,
+                                    xi,
+                                    cum_sums.data(),
+                                    batch_no,
+                                    list_size,
+                                    sel,
+                                    ids,
+                                    use_sel,
+                                    active_indices,
+                                    exact_distances,
+                                    radius,
+                                    local_stats);
 
             // Add batch survivors to range result.
             for (size_t i = 0; i < num_active; i++) {

data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp

@@ -24,9 +24,10 @@ IndexIVFRaBitQ::IndexIVFRaBitQ(
         const size_t d,
         const size_t nlist,
         MetricType metric,
-        bool own_invlists)
+        bool own_invlists,
+        uint8_t nb_bits_in)
         : IndexIVF(quantizer, d, nlist, 0, metric, own_invlists),
-          rabitq(d, metric) {
+          rabitq(d, metric, nb_bits_in) {
     code_size = rabitq.code_size;
     if (own_invlists) {
         invlists->code_size = code_size;
@@ -153,6 +154,8 @@ struct RaBitInvertedListScanner : InvertedListScanner {
     std::vector<float> query_vector;
 
     std::unique_ptr<FlatCodesDistanceComputer> dc;
+    RaBitQDistanceComputer* rabitq_dc =
+            nullptr; // For multi-bit adaptive filtering
 
     uint8_t qb = 0;
     bool centered = false;
@@ -194,6 +197,84 @@ struct RaBitInvertedListScanner : InvertedListScanner {
         return dc->distance_to_code(code);
     }
 
+    /// Override scan_codes to implement adaptive filtering for multi-bit codes
+    size_t scan_codes(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float* simi,
+            idx_t* idxi,
+            size_t k) const override {
+        size_t ex_bits = ivf_rabitq.rabitq.nb_bits - 1;
+
+        // For 1-bit codes, use default implementation
+        if (ex_bits == 0 || rabitq_dc == nullptr) {
+            return InvertedListScanner::scan_codes(
+                    list_size, codes, ids, simi, idxi, k);
+        }
+
+        // Multi-bit: Two-stage search with adaptive filtering
+        size_t nup = 0;
+
+        // Stats tracking for multi-bit two-stage search
+        // n_1bit_evaluations: candidates evaluated using 1-bit lower bound
+        // n_multibit_evaluations: candidates requiring full multi-bit distance
+        size_t local_1bit_evaluations = 0;
+        size_t local_multibit_evaluations = 0;
+
+        for (size_t j = 0; j < list_size; j++) {
+            if (sel != nullptr) {
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                if (!sel->is_member(id)) {
+                    codes += code_size;
+                    continue;
+                }
+            }
+
+            local_1bit_evaluations++;
+
+            // Stage 1: Compute lower bound using 1-bit codes
+            float lower_bound = rabitq_dc->lower_bound_distance(codes);
+
+            // Stage 2: Adaptive filtering
+            // L2 (min-heap): filter if lower_bound < simi[0]
+            // IP (max-heap): filter if lower_bound > simi[0]
+            // Note: Using simi[0] directly (not cached) enables more aggressive
+            // filtering as the heap is updated with better candidates
+            bool should_refine = keep_max ? (lower_bound > simi[0])
+                                          : (lower_bound < simi[0]);
+
+            if (should_refine) {
+                local_multibit_evaluations++;
+                // Lower bound is promising, compute full distance
+                float dis = distance_to_code(codes);
+
+                // Check if distance improves heap
+                bool improves_heap =
+                        keep_max ? (dis > simi[0]) : (dis < simi[0]);
+
+                if (improves_heap) {
+                    int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                    if (keep_max) {
+                        minheap_replace_top(k, simi, idxi, dis, id);
+                    } else {
+                        maxheap_replace_top(k, simi, idxi, dis, id);
+                    }
+                    nup++;
+                }
+            }
+            codes += code_size;
+        }
+
+        // Update global stats atomically
+#pragma omp atomic
+        rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
+#pragma omp atomic
+        rabitq_stats.n_multibit_evaluations += local_multibit_evaluations;
+
+        return nup;
+    }
+
     void internal_try_setup_dc() {
         if (!query_vector.empty() && !reconstructed_centroid.empty()) {
             // both query_vector and centroid are available!
@@ -202,6 +283,9 @@ struct RaBitInvertedListScanner : InvertedListScanner {
                     qb, reconstructed_centroid.data(), centered));
 
             dc->set_query(query_vector.data());
+
+            // Try to cast to RaBitQDistanceComputer for multi-bit support
+            rabitq_dc = dynamic_cast<RaBitQDistanceComputer*>(dc.get());
         }
     }
 };

data/vendor/faiss/faiss/IndexIVFRaBitQ.h

@@ -13,6 +13,7 @@
 #include <faiss/Index.h>
 #include <faiss/IndexIVF.h>
 
+#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQuantizer.h>
 
 namespace faiss {
@@ -35,7 +36,8 @@ struct IndexIVFRaBitQ : IndexIVF {
             const size_t d,
             const size_t nlist,
             MetricType metric = METRIC_L2,
-            bool own_invlists = true);
+            bool own_invlists = true,
+            uint8_t nb_bits = 1);
 
     IndexIVFRaBitQ();