faiss 0.6.2 → 0.6.3

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

@@ -7,6 +7,7 @@
 
 #include <faiss/impl/HNSW.h>
 
+#include <array>
 #include <cinttypes>
 #include <cstddef>
 #include <cstdlib>
@@ -172,17 +173,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;
@@ -1092,8 +1097,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 +1137,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);
-
-        if (do_dis_check) {
-            // tricky stopping condition: there are more than ef
-            // distances that are processed already that are smaller
-            // than 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];
 
-            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];
-
-            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));
+            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 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;
+            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_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;
-                } else {
-                    candidates.push(idx_3, new_exact_3);
-                }
-            }
-
-            // 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);
+                    candidates.push(index_array[i], ne);
                 }
             }
 
-            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 +1274,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;
         }

data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp

@@ -597,8 +597,8 @@ void LocalSearchQuantizer::icm_encode_step(
         const float* binaries,
         size_t n,
         size_t n_iters) const {
-    FAISS_THROW_IF_NOT(M != 0 && K != 0);
-    FAISS_THROW_IF_NOT(binaries != nullptr);
+    FAISS_THROW_IF_MSG(M == 0 || K == 0, "M and K must be nonzero");
+    FAISS_THROW_IF_NOT(binaries);
 
     // Resolve SIMD level once, not per iteration of the n × n_iters × M loop.
     with_simd_level_256bit([&]<SIMDLevel SL>() {

data/vendor/faiss/faiss/impl/NSG.cpp

@@ -142,7 +142,9 @@ void NSG::build(
         idx_t n,
         const nsg::Graph<idx_t>& knn_graph,
         bool verbose) {
-    FAISS_THROW_IF_NOT(!is_built && ntotal == 0);
+    FAISS_THROW_IF_MSG(
+            is_built || ntotal != 0,
+            "NSG graph must be empty and not yet built");
 
     if (verbose) {
         printf("NSG::build R=%d, L=%d, C=%d\n", R, L, C);

data/vendor/faiss/faiss/impl/Panorama.h

@@ -49,18 +49,20 @@ static inline void compute_level_dot_kernel(
         const uint32_t* active_indices,
         const size_t num_active,
         const size_t level_width_dims,
-        float* FAISS_RESTRICT dot_products) {
+        float* FAISS_RESTRICT dot_products,
+        size_t stride = 0) {
     const size_t width = LevelWidth > 0 ? LevelWidth : level_width_dims;
+    const size_t row_stride = stride == 0 ? width : stride;
     size_t i = 0;
     for (; i + 4 <= num_active; i += 4) {
         const float* y0 = level_storage +
-                (AllActive ? (i + 0) : active_indices[i + 0]) * width;
+                (AllActive ? (i + 0) : active_indices[i + 0]) * row_stride;
         const float* y1 = level_storage +
-                (AllActive ? (i + 1) : active_indices[i + 1]) * width;
+                (AllActive ? (i + 1) : active_indices[i + 1]) * row_stride;
         const float* y2 = level_storage +
-                (AllActive ? (i + 2) : active_indices[i + 2]) * width;
+                (AllActive ? (i + 2) : active_indices[i + 2]) * row_stride;
         const float* y3 = level_storage +
-                (AllActive ? (i + 3) : active_indices[i + 3]) * width;
+                (AllActive ? (i + 3) : active_indices[i + 3]) * row_stride;
 
         float dp0 = 0, dp1 = 0, dp2 = 0, dp3 = 0;
         FAISS_PRAGMA_IMPRECISE_LOOP
@@ -78,8 +80,8 @@ static inline void compute_level_dot_kernel(
         dot_products[i + 3] = dp3;
     }
     for (; i < num_active; i++) {
-        const float* yj =
-                level_storage + (AllActive ? i : active_indices[i]) * width;
+        const float* yj = level_storage +
+                (AllActive ? i : active_indices[i]) * row_stride;
         float dp = 0;
         FAISS_PRAGMA_IMPRECISE_LOOP
         for (size_t j = 0; j < width; j++) {