faiss 0.6.2 → 0.6.4

data/vendor/faiss/faiss/impl/HNSW.cpp

@@ -7,6 +7,7 @@
 
 #include <faiss/impl/HNSW.h>
 
+#include <array>
 #include <cinttypes>
 #include <cstddef>
 #include <cstdlib>
@@ -16,8 +17,10 @@
 
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/IDSelector.h>
+#include <faiss/impl/RaBitQuantizer.h>
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/impl/VisitedTable.h>
+#include <faiss/impl/hnsw/LockVector.h>
 #include <faiss/impl/hnsw/MinimaxHeap.h>
 
 namespace faiss {
@@ -172,17 +175,21 @@ void HNSW::print_neighbor_stats(int level) const {
 }
 
 void HNSW::fill_with_random_links(size_t n) {
-    int max_level_2 = prepare_level_tab(n);
+    if (n == 0) {
+        return;
+    }
+    max_level = prepare_level_tab(n);
+    entry_point = 0;
+
     RandomGenerator rng2(456);
 
-    for (int level = max_level_2 - 1; level >= 0; --level) {
+    for (int level = max_level - 1; level >= 0; --level) {
         std::vector<int> elts;
         for (size_t i = 0; i < n; i++) {
             if (levels[i] > level) {
                 elts.push_back(i);
             }
         }
-        printf("linking %zd elements in level %d\n", elts.size(), level);
 
         if (elts.size() == 1) {
             continue;
@@ -217,17 +224,17 @@ int HNSW::prepare_level_tab(size_t n, bool preset_levels) {
         }
     }
 
-    int max_level_2 = 0;
+    int local_max_level = 0;
     for (size_t i = 0; i < n; i++) {
         int pt_level = levels[i + n0] - 1;
-        if (pt_level > max_level_2) {
-            max_level_2 = pt_level;
+        if (pt_level > local_max_level) {
+            local_max_level = pt_level;
         }
         offsets.push_back(offsets.back() + cum_nb_neighbors(pt_level + 1));
     }
     neighbors.resize(offsets.back(), -1);
 
-    return max_level_2;
+    return local_max_level;
 }
 
 /** Enumerate vertices from nearest to farthest from query, keep a
@@ -300,6 +307,11 @@ namespace {
 
 using storage_idx_t = HNSW::storage_idx_t;
 
+inline size_t pruned_neighbor_size(size_t max_size, float prune_headroom) {
+    return static_cast<size_t>(
+            max_size - max_size * std::clamp(prune_headroom, 0.0f, 0.5f));
+}
+
 // Map a (high-level) HNSW comparator C — which uses int64_t IDs — to the
 // (low-level) MinimaxHeap comparator HC, which uses int32_t IDs.
 template <class C>
@@ -779,83 +791,226 @@ HNSWStats hnsw_detail::greedy_update_nearest(
             hnsw, qdis, level, nearest, d_nearest);
 }
 
+/**************************************************************
+ * Deterministic addition subroutines
+ **************************************************************/
+
 namespace {
 
 template <class C>
-void add_with_locks_impl(
+void compute_forward_links_impl(
         HNSW& hnsw,
         DistanceComputer& ptdis,
         int pt_level,
-        int pt_id,
-        LockVector& locks,
+        storage_idx_t pt_id,
         VisitedTable& vt,
+        std::vector<std::pair<storage_idx_t, int>>& pt_reverse_edges,
         bool keep_max_size_level0) {
     storage_idx_t nearest = hnsw.entry_point;
-    if (nearest == -1) { // avoid locking after the first point.
-#pragma omp critical
-        if (hnsw.entry_point == -1) { // double-check under lock.
-            hnsw.max_level = pt_level;
-            hnsw.entry_point = pt_id;
-            // leave nearest = -1 to trigger early exit after critical block.
-        } else {
-            // else: Another thread set the entry point.
-            nearest = hnsw.entry_point;
-        }
-    }
-
-    if (nearest < 0) {
-        return;
-    }
-
-    locks.lock(pt_id);
-
-    int level = hnsw.max_level; // level at which we start adding neighbors
+    FAISS_ASSERT(nearest >= 0);
     float d_nearest = ptdis(nearest);
 
-    //  greedy search on upper levels
+    int level = hnsw.max_level;
+    // greedy descent on the upper levels the point does not live on
     for (; level > pt_level; level--) {
         greedy_update_nearest_impl<C>(hnsw, ptdis, level, nearest, d_nearest);
     }
 
+    // levels the point lives on: search from the same entry point
     for (; level >= 0; level--) {
-        add_links_starting_from_impl<C>(
+        std::priority_queue<HNSW::NodeDistCloserT<C>> link_targets;
+        search_neighbors_to_add_dispatch<C>(
                 hnsw,
                 ptdis,
-                pt_id,
+                link_targets,
                 nearest,
                 d_nearest,
                 level,
-                locks,
                 vt,
-                keep_max_size_level0);
+                false);
+
+        int M = hnsw.nb_neighbors(level);
+        shrink_neighbor_list_inner<C>(
+                ptdis, link_targets, M, keep_max_size_level0 && (level == 0));
+
+        // pt_id is not reachable yet, so the snapshot stays immutable for the
+        // other points in this batch
+        size_t begin, end;
+        hnsw.neighbor_range(pt_id, level, &begin, &end);
+        size_t i = begin;
+        while (!link_targets.empty()) {
+            storage_idx_t other_id = link_targets.top().id;
+            link_targets.pop();
+            // pt_id is in its own candidate list when it is the entry point
+            if (other_id == pt_id) {
+                continue;
+            }
+            FAISS_ASSERT(i < end);
+            hnsw.neighbors[i++] = other_id;
+            pt_reverse_edges.emplace_back(other_id, level);
+        }
+        while (i < end) {
+            hnsw.neighbors[i++] = -1;
+        }
     }
+}
 
-    locks.unlock(pt_id);
+template <class C>
+void merge_reverse_links_impl(
+        HNSW& hnsw,
+        DistanceComputer& dis,
+        storage_idx_t node,
+        int level,
+        std::vector<storage_idx_t>& incoming,
+        bool keep_max_size_level0) {
+    size_t begin, end;
+    hnsw.neighbor_range(node, level, &begin, &end);
+    size_t max_size = end - begin;
 
-#pragma omp critical
-    {
-        if (pt_level > hnsw.max_level) {
-            hnsw.max_level = pt_level;
-            hnsw.entry_point = pt_id;
+    std::vector<storage_idx_t> cands;
+    cands.reserve(max_size + incoming.size());
+    for (size_t i = begin; i < end; i++) {
+        if (hnsw.neighbors[i] < 0) {
+            break;
         }
+        cands.push_back(hnsw.neighbors[i]);
+    }
+    size_t n_existing = cands.size();
+
+    std::sort(incoming.begin(), incoming.end());
+    incoming.erase(
+            std::unique(incoming.begin(), incoming.end()), incoming.end());
+    for (storage_idx_t v : incoming) {
+        if (v == node) {
+            continue;
+        }
+        bool present = false;
+        for (size_t i = 0; i < n_existing; i++) {
+            if (cands[i] == v) {
+                present = true;
+                break;
+            }
+        }
+        if (!present) {
+            cands.push_back(v);
+        }
+    }
+
+    if (cands.size() <= max_size) {
+        size_t i = begin;
+        for (storage_idx_t v : cands) {
+            hnsw.neighbors[i++] = v;
+        }
+        while (i < end) {
+            hnsw.neighbors[i++] = -1;
+        }
+        return;
+    }
+
+    // Ties broken by id, so the result is independent of collection order.
+    std::vector<std::pair<float, storage_idx_t>> arr;
+    arr.reserve(cands.size());
+    for (storage_idx_t v : cands) {
+        arr.emplace_back(dis.symmetric_dis(node, v), v);
+    }
+    std::sort(
+            arr.begin(),
+            arr.end(),
+            [](const std::pair<float, storage_idx_t>& a,
+               const std::pair<float, storage_idx_t>& b) {
+                if (a.first != b.first) {
+                    // C::cmp(x, y) is true when y is "better" than x, so
+                    // a before b iff a is better than b.
+                    return C::cmp(b.first, a.first);
+                }
+                return a.second < b.second;
+            });
+
+    // Same headroom as add_link's reciprocal pruning, so degrees stay
+    // comparable between the two builds.
+    size_t pruned_size = pruned_neighbor_size(max_size, hnsw.prune_headroom);
+    if (pruned_size < 1) {
+        pruned_size = 1;
+    }
+
+    std::vector<std::pair<float, storage_idx_t>> kept;
+    std::vector<std::pair<float, storage_idx_t>> outsiders;
+    for (const auto& cand : arr) {
+        float dist_v1_q = cand.first;
+        bool good = true;
+        for (const auto& k : kept) {
+            float dist_v1_v2 = dis.symmetric_dis(k.second, cand.second);
+            if (C::cmp(dist_v1_q, dist_v1_v2)) {
+                good = false;
+                break;
+            }
+        }
+        if (good) {
+            kept.push_back(cand);
+            if (kept.size() >= pruned_size) {
+                break;
+            }
+        } else if (keep_max_size_level0) {
+            outsiders.push_back(cand);
+        }
+    }
+    for (size_t idx = 0; keep_max_size_level0 && kept.size() < max_size &&
+         idx < outsiders.size();
+         idx++) {
+        kept.push_back(outsiders[idx]);
+    }
+
+    size_t i = begin;
+    for (const auto& k : kept) {
+        hnsw.neighbors[i++] = k.second;
+    }
+    while (i < end) {
+        hnsw.neighbors[i++] = -1;
     }
 }
 
 } // namespace
 
-void HNSW::add_with_locks(
+void HNSW::compute_forward_links_deterministic(
         DistanceComputer& ptdis,
         int pt_level,
-        int pt_id,
-        LockVector& locks,
+        storage_idx_t pt_id,
         VisitedTable& vt,
+        std::vector<std::pair<storage_idx_t, int>>& pt_reverse_edges,
         bool keep_max_size_level0) {
     if (is_similarity) {
-        add_with_locks_impl<C_similarity>(
-                *this, ptdis, pt_level, pt_id, locks, vt, keep_max_size_level0);
+        compute_forward_links_impl<C_similarity>(
+                *this,
+                ptdis,
+                pt_level,
+                pt_id,
+                vt,
+                pt_reverse_edges,
+                keep_max_size_level0);
     } else {
-        add_with_locks_impl<C_distance>(
-                *this, ptdis, pt_level, pt_id, locks, vt, keep_max_size_level0);
+        compute_forward_links_impl<C_distance>(
+                *this,
+                ptdis,
+                pt_level,
+                pt_id,
+                vt,
+                pt_reverse_edges,
+                keep_max_size_level0);
+    }
+}
+
+void HNSW::merge_reverse_links_deterministic(
+        DistanceComputer& dis,
+        storage_idx_t node,
+        int level,
+        std::vector<storage_idx_t>& incoming,
+        bool keep_max_size_level0) {
+    if (is_similarity) {
+        merge_reverse_links_impl<C_similarity>(
+                *this, dis, node, level, incoming, keep_max_size_level0);
+    } else {
+        merge_reverse_links_impl<C_distance>(
+                *this, dis, node, level, incoming, keep_max_size_level0);
     }
 }
 
@@ -886,13 +1041,82 @@ inline void extract_search_params(
     }
 }
 
-/** Templated body of `search_from_candidates` — instantiated once per
- * VisitedTable subclass × comparator.
+struct DefaultCandidateDistanceEvaluator {
+    DistanceComputer& qdis;
+
+    template <typename GetThreshold, typename AddResult>
+    size_t evaluate(
+            const storage_idx_t* ids,
+            int count,
+            GetThreshold&& /* get_threshold */,
+            AddResult&& add_result) {
+        if (count == 4) {
+            float distances[4];
+            qdis.distances_batch_4(
+                    ids[0],
+                    ids[1],
+                    ids[2],
+                    ids[3],
+                    distances[0],
+                    distances[1],
+                    distances[2],
+                    distances[3]);
+            for (int i = 0; i < 4; ++i) {
+                add_result(ids[i], distances[i]);
+            }
+        } else {
+            for (int i = 0; i < count; ++i) {
+                add_result(ids[i], qdis(ids[i]));
+            }
+        }
+        return count;
+    }
+};
+
+/** RaBitQ's staged evaluator plugs into the same graph traversal as the
+ * default distance computer. Each candidate first receives a 1-bit estimate;
+ * the full code is evaluated only when its error bound can still beat the
+ * current result threshold.
  */
-template <typename VTType, class C>
+struct RaBitQCandidateDistanceEvaluator {
+    RaBitQDistanceComputer& rq;
+    const bool is_similarity;
+
+    template <typename GetThreshold, typename AddResult>
+    size_t evaluate(
+            const storage_idx_t* ids,
+            int count,
+            GetThreshold&& get_threshold,
+            AddResult&& add_result) {
+        size_t ndis = 0;
+        for (int i = 0; i < count; ++i) {
+            const uint8_t* code =
+                    rq.codes + static_cast<size_t>(ids[i]) * rq.code_size;
+            const float estimate = rq.distance_to_code_1bit(code);
+            rq.stats.n_1bit++;
+            ndis++;
+
+            float distance = estimate;
+            if (rq.should_refine(
+                        code, estimate, get_threshold(), is_similarity)) {
+                distance = rq.distance_to_code_full(code);
+                rq.stats.n_refine++;
+                ndis++;
+            }
+            add_result(ids[i], distance);
+        }
+        return ndis;
+    }
+};
+
+/** Templated body shared by the default and staged RaBitQ searches. The
+ * evaluator changes only how a group of unvisited neighbors receives its
+ * distance; traversal, result handling and stopping conditions stay identical.
+ */
+template <typename VTType, class C, typename DistanceEvaluator>
 int search_from_candidates_fixVT(
         const HNSW& hnsw,
-        DistanceComputer& qdis,
+        DistanceEvaluator& evaluator,
         ResultHandler& res,
         MinimaxHeapT<HC_for<C>>& candidates,
         VTType& vt,
@@ -901,7 +1125,7 @@ int search_from_candidates_fixVT(
         int nres_in,
         const SearchParameters* params) {
     int nres = nres_in;
-    int ndis = 0;
+    size_t ndis = 0;
 
     bool do_dis_check;
     int efSearch;
@@ -959,7 +1183,7 @@ int search_from_candidates_fixVT(
         }
 
         int counter = 0;
-        size_t saved_j[4];
+        storage_idx_t saved_j[4];
 
         threshold = res.threshold;
 
@@ -982,32 +1206,18 @@ int search_from_candidates_fixVT(
             counter += vt.set(v1) ? 1 : 0;
 
             if (counter == 4) {
-                float dis[4];
-                qdis.distances_batch_4(
-                        saved_j[0],
-                        saved_j[1],
-                        saved_j[2],
-                        saved_j[3],
-                        dis[0],
-                        dis[1],
-                        dis[2],
-                        dis[3]);
-
-                for (size_t id4 = 0; id4 < 4; id4++) {
-                    add_to_heap(saved_j[id4], dis[id4]);
-                }
-
-                ndis += 4;
-
+                ndis += evaluator.evaluate(
+                        saved_j,
+                        counter,
+                        [&] { return threshold; },
+                        add_to_heap);
                 counter = 0;
             }
         }
 
-        for (int icnt = 0; icnt < counter; icnt++) {
-            float dis = qdis(saved_j[icnt]);
-            add_to_heap(saved_j[icnt], dis);
-
-            ndis += 1;
+        if (counter > 0) {
+            ndis += evaluator.evaluate(
+                    saved_j, counter, [&] { return threshold; }, add_to_heap);
         }
 
         nstep++;
@@ -1028,12 +1238,12 @@ int search_from_candidates_fixVT(
     return nres;
 }
 
-/// Dispatches the VisitedTable concrete type for a given C, then calls
-/// the templated `search_from_candidates_fixVT<VTType, C>`.
-template <class C>
-int search_from_candidates_dispatch(
+/// Dispatches the VisitedTable concrete type while keeping the distance
+/// evaluator statically selected outside the graph traversal loop.
+template <class C, typename DistanceEvaluator>
+int search_from_candidates_evaluator_dispatch(
         const HNSW& hnsw,
-        DistanceComputer& qdis,
+        DistanceEvaluator& evaluator,
         ResultHandler& res,
         MinimaxHeapT<HC_for<C>>& candidates,
         VisitedTable& vt,
@@ -1044,7 +1254,7 @@ int search_from_candidates_dispatch(
     auto call = [&]<typename VTType>(VTType& vt_concrete) -> int {
         return search_from_candidates_fixVT<VTType, C>(
                 hnsw,
-                qdis,
+                evaluator,
                 res,
                 candidates,
                 vt_concrete,
@@ -1060,6 +1270,59 @@ int search_from_candidates_dispatch(
     return call(vts);
 }
 
+template <class C>
+int search_from_candidates_dispatch(
+        const HNSW& hnsw,
+        DistanceComputer& qdis,
+        ResultHandler& res,
+        MinimaxHeapT<HC_for<C>>& candidates,
+        VisitedTable& vt,
+        HNSWStats& stats,
+        int level,
+        int nres_in,
+        const SearchParameters* params) {
+    DefaultCandidateDistanceEvaluator evaluator{qdis};
+    return search_from_candidates_evaluator_dispatch<C>(
+            hnsw,
+            evaluator,
+            res,
+            candidates,
+            vt,
+            stats,
+            level,
+            nres_in,
+            params);
+}
+
+template <class C>
+int search_from_candidates_rabitq_dispatch(
+        const HNSW& hnsw,
+        DistanceComputer& qdis,
+        ResultHandler& res,
+        MinimaxHeapT<HC_for<C>>& candidates,
+        VisitedTable& vt,
+        HNSWStats& stats,
+        int level,
+        int nres_in,
+        const SearchParameters* params) {
+    auto* rq = dynamic_cast<RaBitQDistanceComputer*>(&qdis);
+    FAISS_THROW_IF_NOT_MSG(
+            rq, "staged RaBitQ search requires an IndexRaBitQ storage");
+    FAISS_THROW_IF_NOT_MSG(
+            rq->nb_bits >= 2, "staged RaBitQ search requires nb_bits >= 2");
+    RaBitQCandidateDistanceEvaluator evaluator{*rq, hnsw.is_similarity};
+    return search_from_candidates_evaluator_dispatch<C>(
+            hnsw,
+            evaluator,
+            res,
+            candidates,
+            vt,
+            stats,
+            level,
+            nres_in,
+            params);
+}
+
 } // namespace
 
 /** Do a BFS on the candidates list. Public dispatcher: only handles the
@@ -1092,8 +1355,8 @@ int hnsw_detail::search_from_candidates_panorama(
         const SearchParameters* params) {
     // Panorama's progressive-bound math is L2-specific: refuse to run in
     // similarity mode.
-    FAISS_THROW_IF_NOT_MSG(
-            !hnsw.is_similarity,
+    FAISS_THROW_IF_MSG(
+            hnsw.is_similarity,
             "search_from_candidates_panorama does not support is_similarity=true");
 
     using C = HNSW::C_distance;
@@ -1132,177 +1395,129 @@ int hnsw_detail::search_from_candidates_panorama(
             flat_codes_qdis,
             "DistanceComputer must be a FlatCodesDistanceComputer");
 
-    // Allocate space for the index array and exact distances.
-    size_t M = hnsw.nb_neighbors(0);
-    std::vector<idx_t> index_array(M);
-    std::vector<float> exact_distances(M);
+    const auto& pano = panorama_index->pano;
+    const size_t nb_per_parent = static_cast<size_t>(hnsw.nb_neighbors(level));
+    const size_t num_panorama_levels = pano.n_levels;
+    const size_t level_width_floats = pano.level_width_floats;
 
-    const float* query = flat_codes_qdis->q;
-    std::vector<float> query_cum_sums(panorama_index->pano.n_levels + 1);
-    panorama_index->pano.compute_query_cum_sums(query, query_cum_sums.data());
-    float query_norm_sq = query_cum_sums[0] * query_cum_sums[0];
+    constexpr size_t kTargetBatch = 64;
+    const size_t buf_cap = kTargetBatch + nb_per_parent;
+    std::vector<uint32_t> index_array(buf_cap);
+    std::vector<float> exact_distances(buf_cap);
+    std::vector<float> dot_buffer(buf_cap);
+    std::vector<float> query_cum_sums_buf(num_panorama_levels + 1);
 
-    int nstep = 0;
+    const float* query = flat_codes_qdis->q;
     const size_t d = static_cast<size_t>(panorama_index->d);
 
-    PanoramaStats local_pano_stats;
-    local_pano_stats.reset();
+    const float* cum_base = panorama_index->get_cum_sum(0);
+    const size_t cum_stride = num_panorama_levels + 1;
+    const auto* flat_storage =
+            static_cast<const IndexFlat*>(panorama_index->storage);
+    const float* xb_base =
+            reinterpret_cast<const float*>(flat_storage->codes.data());
+    const size_t feat_stride = d;
 
-    while (candidates.size() > 0) {
-        float d0 = 0;
-        int v0 = candidates.pop_min(&d0);
+    pano.compute_query_cum_sums(query, query_cum_sums_buf.data());
+    const float* query_cum_sums = query_cum_sums_buf.data();
+    const float query_norm_sq = query_cum_sums[0] * query_cum_sums[0];
 
-        if (do_dis_check) {
-            // tricky stopping condition: there are more than ef
-            // distances that are processed already that are smaller
-            // than d0
-
-            int n_dis_below = candidates.count_below(d0);
-            if (n_dis_below >= efSearch) {
-                break;
-            }
-        }
+    int nstep = 0;
 
-        size_t begin, end;
-        hnsw.neighbor_range(v0, level, &begin, &end);
+    PanoramaStats local_pano_stats;
+    local_pano_stats.reset();
 
-        // Unlike the vanilla HNSW, we already remove (and compact) the visited
-        // nodes from the candidates list at this stage. We also remove nodes
-        // that are not selected.
+    bool stop_flag = false;
+    while (candidates.size() > 0 && !stop_flag) {
         size_t initial_size = 0;
-        for (size_t j = begin; j < end; j++) {
-            int v1 = hnsw.neighbors[j];
-            if (v1 < 0) {
-                break;
+        size_t k_popped = 0;
+        while (initial_size < kTargetBatch && candidates.size() > 0) {
+            float d0 = 0;
+            int v0 = candidates.pop_min(&d0);
+            if (do_dis_check) {
+                int n_dis_below = candidates.count_below(d0);
+                if (n_dis_below >= efSearch) {
+                    if (k_popped == 0) {
+                        // Standard early-stop: nothing queued this
+                        // iteration, terminate the outer loop entirely
+                        // (matches the single-pop path's behavior).
+                        stop_flag = true;
+                    } else {
+                        // We already have parents queued; un-pop this
+                        // one so the next outer iteration sees it and
+                        // re-applies the stop check from a clean state.
+                        candidates.push(v0, d0);
+                    }
+                    break;
+                }
             }
+            k_popped++;
 
-            const float* cum_sums_v1 = panorama_index->get_cum_sum(v1);
-            index_array[initial_size] = v1;
-            exact_distances[initial_size] =
-                    query_norm_sq + cum_sums_v1[0] * cum_sums_v1[0];
+            size_t begin, end;
+            hnsw.neighbor_range(v0, level, &begin, &end);
+            for (size_t j = begin; j < end; j++) {
+                int v1 = hnsw.neighbors[j];
+                if (v1 < 0) {
+                    break;
+                }
 
-            bool is_selected = !sel || sel->is_member(v1);
-            initial_size += is_selected && vt.set(v1) ? 1 : 0;
+                bool is_new = vt.set(v1);
+                bool is_selected = !sel || sel->is_member(v1);
+                if (is_new && is_selected) {
+                    const float vsum =
+                            cum_base[static_cast<size_t>(v1) * cum_stride];
+                    index_array[initial_size] = v1;
+                    exact_distances[initial_size] = query_norm_sq + vsum * vsum;
+                    initial_size++;
+                }
+            }
         }
 
         local_pano_stats.total_dims += initial_size * d;
+
         size_t batch_size = initial_size;
         size_t curr_panorama_level = 0;
-        const size_t num_panorama_levels = panorama_index->pano.n_levels;
         while (curr_panorama_level < num_panorama_levels && batch_size > 0) {
-            float query_cum_norm = query_cum_sums[curr_panorama_level + 1];
+            const size_t cs_level_idx = curr_panorama_level + 1;
+            const float query_cum_norm = query_cum_sums[cs_level_idx];
+            const float two_qc = 2.0f * query_cum_norm;
+
+            const size_t start_dim = curr_panorama_level * level_width_floats;
+            size_t end_dim = (curr_panorama_level + 1) * level_width_floats;
+            end_dim = std::min(end_dim, d);
+            const size_t dim_span = end_dim - start_dim;
+
+            const float* level_base = xb_base + start_dim;
+            with_level_width(dim_span, [&]<size_t W>() {
+                compute_level_dot_kernel<false, W>(
+                        query + start_dim,
+                        level_base,
+                        index_array.data(),
+                        batch_size,
+                        dim_span,
+                        dot_buffer.data(),
+                        feat_stride);
+            });
+            ndis += batch_size;
 
-            size_t start_dim = curr_panorama_level *
-                    panorama_index->pano.level_width_floats;
-            size_t end_dim = (curr_panorama_level + 1) *
-                    panorama_index->pano.level_width_floats;
-            end_dim = std::min(end_dim, static_cast<size_t>(panorama_index->d));
-
-            size_t i = 0;
             size_t next_batch_size = 0;
-            for (; i + 3 < batch_size; i += 4) {
-                idx_t idx_0 = index_array[i];
-                idx_t idx_1 = index_array[i + 1];
-                idx_t idx_2 = index_array[i + 2];
-                idx_t idx_3 = index_array[i + 3];
-
-                float dp[4];
-                flat_codes_qdis->partial_dot_product_batch_4(
-                        idx_0,
-                        idx_1,
-                        idx_2,
-                        idx_3,
-                        dp[0],
-                        dp[1],
-                        dp[2],
-                        dp[3],
-                        start_dim,
-                        end_dim - start_dim);
-                ndis += 4;
 
-                float new_exact_0 = exact_distances[i + 0] - 2 * dp[0];
-                float new_exact_1 = exact_distances[i + 1] - 2 * dp[1];
-                float new_exact_2 = exact_distances[i + 2] - 2 * dp[2];
-                float new_exact_3 = exact_distances[i + 3] - 2 * dp[3];
-
-                float cum_sum_0 = panorama_index->get_cum_sum(
-                        idx_0)[curr_panorama_level + 1];
-                float cum_sum_1 = panorama_index->get_cum_sum(
-                        idx_1)[curr_panorama_level + 1];
-                float cum_sum_2 = panorama_index->get_cum_sum(
-                        idx_2)[curr_panorama_level + 1];
-                float cum_sum_3 = panorama_index->get_cum_sum(
-                        idx_3)[curr_panorama_level + 1];
-
-                float cs_bound_0 = 2.0f * cum_sum_0 * query_cum_norm;
-                float cs_bound_1 = 2.0f * cum_sum_1 * query_cum_norm;
-                float cs_bound_2 = 2.0f * cum_sum_2 * query_cum_norm;
-                float cs_bound_3 = 2.0f * cum_sum_3 * query_cum_norm;
-
-                float lower_bound_0 = new_exact_0 - cs_bound_0;
-                float lower_bound_1 = new_exact_1 - cs_bound_1;
-                float lower_bound_2 = new_exact_2 - cs_bound_2;
-                float lower_bound_3 = new_exact_3 - cs_bound_3;
-
-                // The following code is not the most branch friendly (due to
-                // the maintenance of the candidate heap), but micro-benchmarks
-                // have shown that it is not worth it to write horrible code to
-                // squeeze out those cycles.
-                if (!C::cmp(lower_bound_0, threshold)) {
-                    exact_distances[next_batch_size] = new_exact_0;
-                    index_array[next_batch_size] = idx_0;
-                    next_batch_size += 1;
-                } else {
-                    candidates.push(idx_0, new_exact_0);
-                }
-                if (!C::cmp(lower_bound_1, threshold)) {
-                    exact_distances[next_batch_size] = new_exact_1;
-                    index_array[next_batch_size] = idx_1;
-                    next_batch_size += 1;
-                } else {
-                    candidates.push(idx_1, new_exact_1);
-                }
-                if (!C::cmp(lower_bound_2, threshold)) {
-                    exact_distances[next_batch_size] = new_exact_2;
-                    index_array[next_batch_size] = idx_2;
-                    next_batch_size += 1;
-                } else {
-                    candidates.push(idx_2, new_exact_2);
-                }
-                if (!C::cmp(lower_bound_3, threshold)) {
-                    exact_distances[next_batch_size] = new_exact_3;
-                    index_array[next_batch_size] = idx_3;
-                    next_batch_size += 1;
+            for (size_t i = 0; i < batch_size; i++) {
+                float ne = exact_distances[i] - 2.0f * dot_buffer[i];
+                float cum = cum_base
+                        [static_cast<size_t>(index_array[i]) * cum_stride +
+                         cs_level_idx];
+                float lb = ne - two_qc * cum;
+                if (lb <= threshold) {
+                    exact_distances[next_batch_size] = ne;
+                    index_array[next_batch_size] = index_array[i];
+                    next_batch_size++;
                 } else {
-                    candidates.push(idx_3, new_exact_3);
+                    candidates.push(index_array[i], ne);
                 }
             }
 
-            // Process the remaining candidates.
-            for (; i < batch_size; i++) {
-                idx_t idx = index_array[i];
-
-                float dp = flat_codes_qdis->partial_dot_product(
-                        idx, start_dim, end_dim - start_dim);
-                ndis += 1;
-                float new_exact = exact_distances[i] - 2.0f * dp;
-
-                float cum_sum = panorama_index->get_cum_sum(
-                        idx)[curr_panorama_level + 1];
-                float cs_bound = 2.0f * cum_sum * query_cum_norm;
-                float lower_bound = new_exact - cs_bound;
-
-                if (!C::cmp(lower_bound, threshold)) {
-                    exact_distances[next_batch_size] = new_exact;
-                    index_array[next_batch_size] = idx;
-                    next_batch_size += 1;
-                } else {
-                    candidates.push(idx, new_exact);
-                }
-            }
-
-            local_pano_stats.total_dims_scanned +=
-                    batch_size * (end_dim - start_dim);
+            local_pano_stats.total_dims_scanned += batch_size * dim_span;
             batch_size = next_batch_size;
             curr_panorama_level++;
         }
@@ -1317,7 +1532,7 @@ int hnsw_detail::search_from_candidates_panorama(
             candidates.push(idx, exact_distances[i]);
         }
 
-        nstep++;
+        nstep += static_cast<int>(k_popped);
         if (!do_dis_check && nstep > efSearch) {
             break;
         }
@@ -1534,37 +1749,49 @@ HNSWStats search_impl(
     }
 
     int ef = std::max(cur_efSearch, k);
-    if (bounded_queue) { // this is the most common branch, for now we only
-                         // support Panorama search in this branch
+    if (bounded_queue) {
         MinimaxHeapT<HC_for<C>> candidates(ef);
 
         candidates.push(nearest, d_nearest);
 
-        if (!hnsw.is_panorama) {
-            search_from_candidates_dispatch<C>(
-                    hnsw, qdis, res, candidates, vt, stats, 0, 0, params);
-        } else {
-            // Panorama is L2-specific and is only valid for C_distance.
-            // The public dispatch ensures we never reach this code path
-            // with C != C_distance, but assert in debug builds.
-            if constexpr (std::is_same_v<C, HNSW::C_distance>) {
-                hnsw_detail::search_from_candidates_panorama(
-                        hnsw,
-                        index,
-                        qdis,
-                        res,
-                        candidates,
-                        vt,
-                        stats,
-                        0,
-                        0,
-                        params);
-            } else {
-                FAISS_THROW_MSG(
-                        "Panorama search is not supported with is_similarity=true");
-            }
+        switch (hnsw.search_method) {
+            case HNSW::SM_DEFAULT:
+                search_from_candidates_dispatch<C>(
+                        hnsw, qdis, res, candidates, vt, stats, 0, 0, params);
+                break;
+            case HNSW::SM_RABITQ:
+                search_from_candidates_rabitq_dispatch<C>(
+                        hnsw, qdis, res, candidates, vt, stats, 0, 0, params);
+                break;
+            case HNSW::SM_PANORAMA:
+                if constexpr (std::is_same_v<C, HNSW::C_distance>) {
+                    hnsw_detail::search_from_candidates_panorama(
+                            hnsw,
+                            index,
+                            qdis,
+                            res,
+                            candidates,
+                            vt,
+                            stats,
+                            0,
+                            0,
+                            params);
+                } else {
+                    FAISS_THROW_MSG(
+                            "Panorama search does not support similarity metrics");
+                }
+                break;
+            default:
+                FAISS_THROW_MSG("invalid HNSW search method");
         }
     } else {
+        FAISS_THROW_IF_NOT_MSG(
+                hnsw.search_method != HNSW::SM_RABITQ,
+                "staged RaBitQ search requires bounded_queue=true");
+        FAISS_THROW_IF_NOT_MSG(
+                hnsw.search_method == HNSW::SM_DEFAULT ||
+                        hnsw.search_method == HNSW::SM_PANORAMA,
+                "invalid HNSW search method");
         auto call = [&]<typename VTType>(VTType& vt_concrete) {
             return search_from_candidate_unbounded_fixVT<VTType, C>(
                     hnsw,
@@ -1622,6 +1849,50 @@ void search_level_0_impl(
 
     int k = extract_k_from_ResultHandler<C>(res);
 
+    auto search_candidates = [&](MinimaxHeapT<HC_for<C>>& candidates,
+                                 int nres) {
+        switch (hnsw.search_method) {
+            case HNSW::SM_DEFAULT:
+                return search_from_candidates_dispatch<C>(
+                        hnsw,
+                        qdis,
+                        res,
+                        candidates,
+                        vt,
+                        search_stats,
+                        0,
+                        nres,
+                        params);
+            case HNSW::SM_RABITQ:
+                return search_from_candidates_rabitq_dispatch<C>(
+                        hnsw,
+                        qdis,
+                        res,
+                        candidates,
+                        vt,
+                        search_stats,
+                        0,
+                        nres,
+                        params);
+            case HNSW::SM_PANORAMA:
+                // Preserve the pre-enum behavior: search_level_0 did not use
+                // Panorama's progressive traversal and fell back to the
+                // ordinary HNSW candidate search.
+                return search_from_candidates_dispatch<C>(
+                        hnsw,
+                        qdis,
+                        res,
+                        candidates,
+                        vt,
+                        search_stats,
+                        0,
+                        nres,
+                        params);
+            default:
+                FAISS_THROW_MSG("invalid HNSW search method");
+        }
+    };
+
     if (search_type == 1) {
         int nres = 0;
 
@@ -1641,16 +1912,7 @@ void search_level_0_impl(
 
             candidates.push(cj, nearest_d[j]);
 
-            nres = search_from_candidates_dispatch<C>(
-                    hnsw,
-                    qdis,
-                    res,
-                    candidates,
-                    vt,
-                    search_stats,
-                    0,
-                    nres,
-                    params);
+            nres = search_candidates(candidates, nres);
             nres = std::min(nres, candidates_size);
         }
     } else if (search_type == 2) {
@@ -1667,8 +1929,7 @@ void search_level_0_impl(
             candidates.push(cj, nearest_d[j]);
         }
 
-        search_from_candidates_dispatch<C>(
-                hnsw, qdis, res, candidates, vt, search_stats, 0, 0, params);
+        search_candidates(candidates, 0);
     }
 }