faiss 0.5.0 → 0.5.1

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -11,12 +11,15 @@
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/ResultHandler.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/extra_distances.h>
 #include <faiss/utils/prefetch.h>
 #include <faiss/utils/sorting.h>
+#include <omp.h>
 #include <cstring>
+#include <numeric>
 
 namespace faiss {
 
@@ -100,15 +103,24 @@ namespace {
 struct FlatL2Dis : FlatCodesDistanceComputer {
     size_t d;
     idx_t nb;
-    const float* q;
     const float* b;
     size_t ndis;
+    size_t npartial_dot_products;
 
     float distance_to_code(const uint8_t* code) final {
         ndis++;
         return fvec_L2sqr(q, (float*)code, d);
     }
 
+    float partial_dot_product(
+            const idx_t i,
+            const uint32_t offset,
+            const uint32_t num_components) final override {
+        npartial_dot_products++;
+        return fvec_inner_product(
+                q + offset, b + i * d + offset, num_components);
+    }
+
     float symmetric_dis(idx_t i, idx_t j) override {
         return fvec_L2sqr(b + j * d, b + i * d, d);
     }
@@ -116,12 +128,13 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
     explicit FlatL2Dis(const IndexFlat& storage, const float* q = nullptr)
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
-                      storage.code_size),
+                      storage.code_size,
+                      q),
               d(storage.d),
               nb(storage.ntotal),
-              q(q),
               b(storage.get_xb()),
-              ndis(0) {}
+              ndis(0),
+              npartial_dot_products(0) {}
 
     void set_query(const float* x) override {
         q = x;
@@ -159,6 +172,50 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
         dis2 = dp2;
         dis3 = dp3;
     }
+
+    void partial_dot_product_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dp0,
+            float& dp1,
+            float& dp2,
+            float& dp3,
+            const uint32_t offset,
+            const uint32_t num_components) final override {
+        npartial_dot_products += 4;
+
+        // compute first, assign next
+        const float* __restrict y0 =
+                reinterpret_cast<const float*>(codes + idx0 * code_size);
+        const float* __restrict y1 =
+                reinterpret_cast<const float*>(codes + idx1 * code_size);
+        const float* __restrict y2 =
+                reinterpret_cast<const float*>(codes + idx2 * code_size);
+        const float* __restrict y3 =
+                reinterpret_cast<const float*>(codes + idx3 * code_size);
+
+        float dp0_ = 0;
+        float dp1_ = 0;
+        float dp2_ = 0;
+        float dp3_ = 0;
+        fvec_inner_product_batch_4(
+                q + offset,
+                y0 + offset,
+                y1 + offset,
+                y2 + offset,
+                y3 + offset,
+                num_components,
+                dp0_,
+                dp1_,
+                dp2_,
+                dp3_);
+        dp0 = dp0_;
+        dp1 = dp1_;
+        dp2 = dp2_;
+        dp3 = dp3_;
+    }
 };
 
 struct FlatIPDis : FlatCodesDistanceComputer {
@@ -519,4 +576,317 @@ void IndexFlat1D::search(
     done:;
     }
 }
+
+/**************************************************************
+ * shared flat Panorama search code
+ **************************************************************/
+
+namespace {
+
+template <bool use_radius, typename BlockHandler>
+inline void flat_pano_search_core(
+        const IndexFlatPanorama& index,
+        BlockHandler& handler,
+        idx_t n,
+        const float* x,
+        float radius,
+        const SearchParameters* params) {
+    using SingleResultHandler = typename BlockHandler::SingleResultHandler;
+
+    IDSelector* sel = params ? params->sel : nullptr;
+    bool use_sel = sel != nullptr;
+
+    [[maybe_unused]] int nt = std::min(int(n), omp_get_max_threads());
+    size_t n_batches = (index.ntotal + index.batch_size - 1) / index.batch_size;
+
+#pragma omp parallel num_threads(nt)
+    {
+        SingleResultHandler res(handler);
+
+        std::vector<float> query_cum_norms(index.n_levels + 1);
+        std::vector<float> exact_distances(index.batch_size);
+        std::vector<uint32_t> active_indices(index.batch_size);
+
+#pragma omp for
+        for (int64_t i = 0; i < n; i++) {
+            const float* xi = x + i * index.d;
+            index.pano.compute_query_cum_sums(xi, query_cum_norms.data());
+
+            PanoramaStats local_stats;
+            local_stats.reset();
+
+            res.begin(i);
+
+            for (size_t batch_no = 0; batch_no < n_batches; batch_no++) {
+                size_t batch_start = batch_no * index.batch_size;
+
+                float threshold;
+                if constexpr (use_radius) {
+                    threshold = radius;
+                } else {
+                    threshold = res.heap_dis[0];
+                }
+
+                size_t num_active =
+                        index.pano
+                                .progressive_filter_batch<CMax<float, int64_t>>(
+                                        index.codes.data(),
+                                        index.cum_sums.data(),
+                                        xi,
+                                        query_cum_norms.data(),
+                                        batch_no,
+                                        index.ntotal,
+                                        sel,
+                                        nullptr,
+                                        use_sel,
+                                        active_indices,
+                                        exact_distances,
+                                        threshold,
+                                        local_stats);
+
+                for (size_t j = 0; j < num_active; j++) {
+                    res.add_result(
+                            exact_distances[active_indices[j]],
+                            batch_start + active_indices[j]);
+                }
+            }
+
+            res.end();
+            indexPanorama_stats.add(local_stats);
+        }
+    }
+}
+
+} // anonymous namespace
+
+/***************************************************
+ * IndexFlatPanorama
+ ***************************************************/
+
+void IndexFlatPanorama::add(idx_t n, const float* x) {
+    size_t offset = ntotal;
+    ntotal += n;
+    size_t num_batches = (ntotal + batch_size - 1) / batch_size;
+
+    codes.resize(num_batches * batch_size * code_size);
+    cum_sums.resize(num_batches * batch_size * (n_levels + 1));
+
+    const uint8_t* code = reinterpret_cast<const uint8_t*>(x);
+    pano.copy_codes_to_level_layout(codes.data(), offset, n, code);
+    pano.compute_cumulative_sums(cum_sums.data(), offset, n, x);
+}
+
+void IndexFlatPanorama::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT(k > 0);
+    FAISS_THROW_IF_NOT(batch_size >= k);
+
+    HeapBlockResultHandler<CMax<float, int64_t>, false> handler(
+            size_t(n), distances, labels, size_t(k), nullptr);
+
+    flat_pano_search_core<false>(*this, handler, n, x, 0.0f, params);
+}
+
+void IndexFlatPanorama::range_search(
+        idx_t n,
+        const float* x,
+        float radius,
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    RangeSearchBlockResultHandler<CMax<float, int64_t>, false> handler(
+            result, radius, nullptr);
+
+    flat_pano_search_core<true>(*this, handler, n, x, radius, params);
+}
+
+void IndexFlatPanorama::reset() {
+    IndexFlat::reset();
+    cum_sums.clear();
+}
+
+void IndexFlatPanorama::reconstruct(idx_t key, float* recons) const {
+    pano.reconstruct(key, recons, codes.data());
+}
+
+void IndexFlatPanorama::reconstruct_n(idx_t i, idx_t n, float* recons) const {
+    Index::reconstruct_n(i, n, recons);
+}
+
+size_t IndexFlatPanorama::remove_ids(const IDSelector& sel) {
+    idx_t j = 0;
+    for (idx_t i = 0; i < ntotal; i++) {
+        if (sel.is_member(i)) {
+            // should be removed
+        } else {
+            if (i > j) {
+                pano.copy_entry(
+                        codes.data(),
+                        codes.data(),
+                        cum_sums.data(),
+                        cum_sums.data(),
+                        j,
+                        i);
+            }
+            j++;
+        }
+    }
+    size_t nremove = ntotal - j;
+    if (nremove > 0) {
+        ntotal = j;
+        size_t num_batches = (ntotal + batch_size - 1) / batch_size;
+        codes.resize(num_batches * batch_size * code_size);
+        cum_sums.resize(num_batches * batch_size * (n_levels + 1));
+    }
+    return nremove;
+}
+
+void IndexFlatPanorama::merge_from(Index& otherIndex, idx_t add_id) {
+    FAISS_THROW_IF_NOT_MSG(add_id == 0, "cannot set ids in FlatPanorama index");
+    check_compatible_for_merge(otherIndex);
+    IndexFlatPanorama* other = static_cast<IndexFlatPanorama*>(&otherIndex);
+
+    std::vector<float> buffer(other->ntotal * code_size);
+    otherIndex.reconstruct_n(0, other->ntotal, buffer.data());
+
+    add(other->ntotal, buffer.data());
+    other->reset();
+}
+
+void IndexFlatPanorama::add_sa_codes(
+        idx_t /* n */,
+        const uint8_t* /* codes_in */,
+        const idx_t* /* xids */) {
+    FAISS_THROW_MSG("add_sa_codes not implemented for IndexFlatPanorama");
+}
+
+void IndexFlatPanorama::permute_entries(const idx_t* perm) {
+    MaybeOwnedVector<uint8_t> new_codes(codes.size());
+    std::vector<float> new_cum_sums(cum_sums.size());
+
+    for (idx_t i = 0; i < ntotal; i++) {
+        pano.copy_entry(
+                new_codes.data(),
+                codes.data(),
+                new_cum_sums.data(),
+                cum_sums.data(),
+                i,
+                perm[i]);
+    }
+
+    std::swap(codes, new_codes);
+    std::swap(cum_sums, new_cum_sums);
+}
+
+void IndexFlatPanorama::search_subset(
+        idx_t n,
+        const float* x,
+        idx_t k_base,
+        const idx_t* base_labels,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const {
+    using SingleResultHandler =
+            HeapBlockResultHandler<CMax<float, int64_t>, false>::
+                    SingleResultHandler;
+    HeapBlockResultHandler<CMax<float, int64_t>, false> handler(
+            size_t(n), distances, labels, size_t(k), nullptr);
+
+    FAISS_THROW_IF_NOT(k > 0);
+    FAISS_THROW_IF_NOT(batch_size == 1);
+
+    [[maybe_unused]] int nt = std::min(int(n), omp_get_max_threads());
+
+#pragma omp parallel num_threads(nt)
+    {
+        SingleResultHandler res(handler);
+
+        std::vector<float> query_cum_norms(n_levels + 1);
+
+        // Panorama's optimized point-wise refinement (Algorithm 2):
+        // Batch-wise Panorama, as implemented in Panorama.h, incurs overhead
+        // from maintaining active_indices and exact_distances. This optimized
+        // implementation has minimal overhead and is thus preferred for
+        // IndexRefine's use case.
+        // 1. Initialize exact distance as ||y||^2 + ||x||^2.
+        // 2. For each level, refine distance incrementally:
+        //    - Compute dot product for current level: exact_dist -= 2*<x,y>.
+        //    - Use Cauchy-Schwarz bound on remaining levels to get lower bound.
+        //    - If there are less than k points in the heap, add the point to
+        //    the heap.
+        //    - Else, prune if lower bound exceeds k-th best distance.
+        // 3. After all levels, update heap if the point survived.
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            const idx_t* __restrict idsi = base_labels + i * k_base;
+            const float* xi = x + i * d;
+
+            PanoramaStats local_stats;
+            local_stats.reset();
+
+            pano.compute_query_cum_sums(xi, query_cum_norms.data());
+            float query_cum_norm = query_cum_norms[0] * query_cum_norms[0];
+
+            res.begin(i);
+
+            for (size_t j = 0; j < k_base; j++) {
+                idx_t idx = idsi[j];
+
+                if (idx < 0) {
+                    continue;
+                }
+
+                size_t cum_sum_offset = (n_levels + 1) * idx;
+                float cum_sum = cum_sums[cum_sum_offset];
+                float exact_distance = cum_sum * cum_sum + query_cum_norm;
+                cum_sum_offset++;
+
+                const float* x_ptr = xi;
+                const float* p_ptr =
+                        reinterpret_cast<const float*>(codes.data()) + d * idx;
+
+                local_stats.total_dims += d;
+
+                bool pruned = false;
+                for (size_t level = 0; level < n_levels; level++) {
+                    local_stats.total_dims_scanned += pano.level_width_floats;
+
+                    // Refine distance
+                    size_t actual_level_width = std::min(
+                            pano.level_width_floats,
+                            d - level * pano.level_width_floats);
+                    float dot_product = fvec_inner_product(
+                            x_ptr, p_ptr, actual_level_width);
+                    exact_distance -= 2 * dot_product;
+
+                    float cum_sum = cum_sums[cum_sum_offset];
+                    float cauchy_schwarz_bound =
+                            2.0f * cum_sum * query_cum_norms[level + 1];
+                    float lower_bound = exact_distance - cauchy_schwarz_bound;
+
+                    // Prune using Cauchy-Schwarz bound
+                    if (lower_bound > res.heap_dis[0]) {
+                        pruned = true;
+                        break;
+                    }
+
+                    cum_sum_offset++;
+                    x_ptr += pano.level_width_floats;
+                    p_ptr += pano.level_width_floats;
+                }
+
+                if (!pruned) {
+                    res.add_result(exact_distance, idx);
+                }
+            }
+
+            res.end();
+            indexPanorama_stats.add(local_stats);
+        }
+    }
+}
 } // namespace faiss

data/vendor/faiss/faiss/IndexFlat.h

@@ -13,6 +13,7 @@
 #include <vector>
 
 #include <faiss/IndexFlatCodes.h>
+#include <faiss/impl/Panorama.h>
 
 namespace faiss {
 
@@ -99,6 +100,85 @@ struct IndexFlatL2 : IndexFlat {
     void clear_l2norms();
 };
 
+struct IndexFlatPanorama : IndexFlat {
+    const size_t batch_size;
+    const size_t n_levels;
+    std::vector<float> cum_sums;
+    Panorama pano;
+
+    /**
+     * @param d dimensionality of the input vectors
+     * @param metric metric type
+     * @param n_levels number of Panorama levels
+     * @param batch_size batch size for Panorama storage
+     */
+    explicit IndexFlatPanorama(
+            idx_t d,
+            MetricType metric,
+            size_t n_levels,
+            size_t batch_size)
+            : IndexFlat(d, metric),
+              batch_size(batch_size),
+              n_levels(n_levels),
+              pano(code_size, n_levels, batch_size) {
+        FAISS_THROW_IF_NOT(metric == METRIC_L2);
+    }
+
+    void add(idx_t n, const float* x) override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
+
+    void search_subset(
+            idx_t n,
+            const float* x,
+            idx_t k_base,
+            const idx_t* base_labels,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
+
+    void reset() override;
+
+    void reconstruct(idx_t key, float* recons) const override;
+
+    void reconstruct_n(idx_t i, idx_t n, float* recons) const override;
+
+    size_t remove_ids(const IDSelector& sel) override;
+
+    void merge_from(Index& otherIndex, idx_t add_id) override;
+
+    void add_sa_codes(idx_t n, const uint8_t* codes_in, const idx_t* xids)
+            override;
+
+    void permute_entries(const idx_t* perm);
+};
+
+struct IndexFlatL2Panorama : IndexFlatPanorama {
+    /**
+     * @param d dimensionality of the input vectors
+     * @param n_levels number of Panorama levels
+     * @param batch_size batch size for Panorama storage
+     */
+    explicit IndexFlatL2Panorama(
+            idx_t d,
+            size_t n_levels,
+            size_t batch_size = 512)
+            : IndexFlatPanorama(d, METRIC_L2, n_levels, batch_size) {}
+};
+
 /// optimized version for 1D "vectors".
 struct IndexFlat1D : IndexFlatL2 {
     bool continuous_update = true; ///< is the permutation updated continuously?

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -276,7 +276,7 @@ void hnsw_search(
                 res.begin(i);
                 dis->set_query(x + i * index->d);
 
-                HNSWStats stats = hnsw.search(*dis, res, vt, params);
+                HNSWStats stats = hnsw.search(*dis, index, res, vt, params);
                 n1 += stats.n1;
                 n2 += stats.n2;
                 ndis += stats.ndis;
@@ -649,6 +649,95 @@ IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
     is_trained = true;
 }
 
+/**************************************************************
+ * IndexHNSWFlatPanorama implementation
+ **************************************************************/
+
+void IndexHNSWFlatPanorama::compute_cum_sums(
+        const float* x,
+        float* dst_cum_sums,
+        int d,
+        int num_panorama_levels,
+        int panorama_level_width) {
+    // Iterate backwards through levels, accumulating sum as we go.
+    // This avoids computing the suffix sum for each vector, which takes
+    // extra memory.
+
+    float sum = 0.0f;
+    dst_cum_sums[num_panorama_levels] = 0.0f;
+    for (int level = num_panorama_levels - 1; level >= 0; level--) {
+        int start_idx = level * panorama_level_width;
+        int end_idx = std::min(start_idx + panorama_level_width, d);
+        for (int j = start_idx; j < end_idx; j++) {
+            sum += x[j] * x[j];
+        }
+        dst_cum_sums[level] = std::sqrt(sum);
+    }
+}
+
+IndexHNSWFlatPanorama::IndexHNSWFlatPanorama()
+        : IndexHNSWFlat(),
+          cum_sums(),
+          panorama_level_width(0),
+          num_panorama_levels(0) {}
+
+IndexHNSWFlatPanorama::IndexHNSWFlatPanorama(
+        int d,
+        int M,
+        int num_panorama_levels,
+        MetricType metric)
+        : IndexHNSWFlat(d, M, metric),
+          cum_sums(),
+          panorama_level_width(
+                  (d + num_panorama_levels - 1) / num_panorama_levels),
+          num_panorama_levels(num_panorama_levels) {
+    // For now, we only support L2 distance.
+    // Supporting dot product and cosine distance is a trivial addition
+    // left for future work.
+    FAISS_THROW_IF_NOT(metric == METRIC_L2);
+
+    // Enable Panorama search mode.
+    // This is not ideal, but is still more simple than making a subclass of
+    // HNSW and overriding the search logic.
+    hnsw.is_panorama = true;
+}
+
+void IndexHNSWFlatPanorama::add(idx_t n, const float* x) {
+    idx_t n0 = ntotal;
+    cum_sums.resize((ntotal + n) * (num_panorama_levels + 1));
+
+    for (size_t idx = 0; idx < n; idx++) {
+        const float* vector = x + idx * d;
+        compute_cum_sums(
+                vector,
+                &cum_sums[(n0 + idx) * (num_panorama_levels + 1)],
+                d,
+                num_panorama_levels,
+                panorama_level_width);
+    }
+
+    IndexHNSWFlat::add(n, x);
+}
+
+void IndexHNSWFlatPanorama::reset() {
+    cum_sums.clear();
+    IndexHNSWFlat::reset();
+}
+
+void IndexHNSWFlatPanorama::permute_entries(const idx_t* perm) {
+    std::vector<float> new_cum_sums(ntotal * (num_panorama_levels + 1));
+
+    for (idx_t i = 0; i < ntotal; i++) {
+        idx_t src = perm[i];
+        memcpy(&new_cum_sums[i * (num_panorama_levels + 1)],
+               &cum_sums[src * (num_panorama_levels + 1)],
+               (num_panorama_levels + 1) * sizeof(float));
+    }
+
+    std::swap(cum_sums, new_cum_sums);
+    IndexHNSWFlat::permute_entries(perm);
+}
+
 /**************************************************************
  * IndexHNSWPQ implementation
  **************************************************************/

data/vendor/faiss/faiss/IndexHNSW.h

@@ -111,7 +111,7 @@ struct IndexHNSW : Index {
 
     void link_singletons();
 
-    void permute_entries(const idx_t* perm);
+    virtual void permute_entries(const idx_t* perm);
 
     DistanceComputer* get_distance_computer() const override;
 };
@@ -125,6 +125,62 @@ struct IndexHNSWFlat : IndexHNSW {
     IndexHNSWFlat(int d, int M, MetricType metric = METRIC_L2);
 };
 
+/** Panorama implementation of IndexHNSWFlat following
+ * https://www.arxiv.org/pdf/2510.00566.
+ *
+ * Unlike cluster-based Panorama, the vectors have to be higher dimensional
+ * (i.e. typically d > 512) and/or be able to compress a lot of their energy in
+ * the early dimensions to be effective. This is because HNSW accesses vectors
+ * in a random order, which makes cache misses dominate the distance computation
+ * time.
+ *
+ * The `num_panorama_levels` parameter controls the granularity of progressive
+ * distance refinement, allowing candidates to be eliminated early using partial
+ * distance computations rather than computing full distances.
+ *
+ * NOTE: This version of HNSW handles search slightly differently than the
+ * vanilla HNSW, as it uses partial distance computations with progressive
+ * refinement bounds. Instead of computing full distances immediately for all
+ * candidates, Panorama maintains lower and upper bounds that are incrementally
+ * tightened across refinement levels. Candidates are inserted into the search
+ * beam using approximate distance estimates (LB+UB)/2 and are only fully
+ * evaluated when they survive pruning and enter the result heap. This allows
+ * the algorithm to prune unpromising candidates early using Cauchy-Schwarz
+ * bounds on partial inner products. Hence, recall is not guaranteed to be the
+ * same as vanilla HNSW due to the heterogeneous precision within the search
+ * beam (exact vs. partial distance estimates affecting traversal order).
+ */
+struct IndexHNSWFlatPanorama : IndexHNSWFlat {
+    IndexHNSWFlatPanorama();
+    IndexHNSWFlatPanorama(
+            int d,
+            int M,
+            int num_panorama_levels,
+            MetricType metric = METRIC_L2);
+
+    void add(idx_t n, const float* x) override;
+    void reset() override;
+    void permute_entries(const idx_t* perm) override;
+
+    /// Inline for performance - called frequently in search hot path.
+    const float* get_cum_sum(idx_t i) const {
+        return cum_sums.data() + i * (num_panorama_levels + 1);
+    }
+
+    /// Compute cumulative sums for a vector (used both for database points and
+    /// queries).
+    static void compute_cum_sums(
+            const float* x,
+            float* dst_cum_sums,
+            int d,
+            int num_panorama_levels,
+            int panorama_level_width);
+
+    std::vector<float> cum_sums;
+    const size_t panorama_level_width;
+    const size_t num_panorama_levels;
+};
+
 /** PQ index topped with with a HNSW structure to access elements
  *  more efficiently.
  */