faiss 0.4.3 → 0.5.1

data/vendor/faiss/faiss/IndexFastScan.h

@@ -8,12 +8,15 @@
 #pragma once
 
 #include <faiss/Index.h>
+#include <faiss/impl/FastScanDistancePostProcessing.h>
 #include <faiss/utils/AlignedTable.h>
 
 namespace faiss {
 
 struct CodePacker;
 struct NormTableScaler;
+struct IDSelector;
+struct SIMDResultHandlerToFloat;
 
 /** Fast scan version of IndexPQ and IndexAQ. Works for 4-bit PQ and AQ for now.
  *
@@ -54,6 +57,14 @@ struct IndexFastScan : Index {
     // (set when initialized by IndexPQ or IndexAQ)
     const uint8_t* orig_codes = nullptr;
 
+    /** Initialize the fast scan index
+     *
+     * @param d         dimensionality of vectors
+     * @param M         number of subquantizers
+     * @param nbits     number of bits per subquantizer
+     * @param metric    distance metric to use
+     * @param bbs       block size for SIMD processing
+     */
     void init_fastscan(
             int d,
             size_t M,
@@ -65,6 +76,15 @@ struct IndexFastScan : Index {
 
     void reset() override;
 
+    /** Search for k nearest neighbors
+     *
+     * @param n          number of query vectors
+     * @param x          query vectors (n * d)
+     * @param k          number of nearest neighbors to find
+     * @param distances  output distances (n * k)
+     * @param labels     output labels/indices (n * k)
+     * @param params     optional search parameters
+     */
     void search(
             idx_t n,
             const float* x,
@@ -73,20 +93,70 @@ struct IndexFastScan : Index {
             idx_t* labels,
             const SearchParameters* params = nullptr) const override;
 
+    /** Add vectors to the index
+     *
+     * @param n  number of vectors to add
+     * @param x  vectors to add (n * d)
+     */
     void add(idx_t n, const float* x) override;
 
+    /** Compute codes for vectors
+     *
+     * @param codes  output codes
+     * @param n      number of vectors to encode
+     * @param x      vectors to encode (n * d)
+     */
     virtual void compute_codes(uint8_t* codes, idx_t n, const float* x)
             const = 0;
 
-    virtual void compute_float_LUT(float* lut, idx_t n, const float* x)
-            const = 0;
+    /** Compute floating-point lookup table for distance computation
+     *
+     * @param lut          output lookup table
+     * @param n            number of query vectors
+     * @param x            query vectors (n * d)
+     * @param context      processing context containing all processors
+     */
+    virtual void compute_float_LUT(
+            float* lut,
+            idx_t n,
+            const float* x,
+            const FastScanDistancePostProcessing& context) const = 0;
+
+    /** Create a KNN handler for this index type
+     *
+     * This method can be overridden by derived classes to provide
+     * specialized handlers (e.g., RaBitQHeapHandler for RaBitQ indexes).
+     * Base implementation creates standard handlers based on k and impl.
+     *
+     * @param is_max       whether to use CMax comparator (true) or CMin (false)
+     * @param impl         implementation number
+     * @param n            number of queries
+     * @param k            number of neighbors to find
+     * @param ntotal       total number of vectors in database
+     * @param distances    output distances array
+     * @param labels       output labels array
+     * @param sel          optional ID selector
+     * @param context      processing context for distance post-processing
+     * @return             pointer to created handler (never returns nullptr)
+     */
+    virtual SIMDResultHandlerToFloat* make_knn_handler(
+            bool is_max,
+            int impl,
+            idx_t n,
+            idx_t k,
+            size_t ntotal,
+            float* distances,
+            idx_t* labels,
+            const IDSelector* sel,
+            const FastScanDistancePostProcessing& context) const;
 
     // called by search function
     void compute_quantized_LUT(
             idx_t n,
             const float* x,
             uint8_t* lut,
-            float* normalizers) const;
+            float* normalizers,
+            const FastScanDistancePostProcessing& context) const;
 
     template <bool is_max>
     void search_dispatch_implem(
@@ -95,7 +165,7 @@ struct IndexFastScan : Index {
             idx_t k,
             float* distances,
             idx_t* labels,
-            const NormTableScaler* scaler) const;
+            const FastScanDistancePostProcessing& context) const;
 
     template <class Cfloat>
     void search_implem_234(
@@ -104,7 +174,7 @@ struct IndexFastScan : Index {
             idx_t k,
             float* distances,
             idx_t* labels,
-            const NormTableScaler* scaler) const;
+            const FastScanDistancePostProcessing& context) const;
 
     template <class C>
     void search_implem_12(
@@ -114,7 +184,7 @@ struct IndexFastScan : Index {
             float* distances,
             idx_t* labels,
             int impl,
-            const NormTableScaler* scaler) const;
+            const FastScanDistancePostProcessing& context) const;
 
     template <class C>
     void search_implem_14(
@@ -124,14 +194,39 @@ struct IndexFastScan : Index {
             float* distances,
             idx_t* labels,
             int impl,
-            const NormTableScaler* scaler) const;
+            const FastScanDistancePostProcessing& context) const;
 
+    /** Reconstruct a vector from its code
+     *
+     * @param key     index of vector to reconstruct
+     * @param recons  output reconstructed vector
+     */
     void reconstruct(idx_t key, float* recons) const override;
+
+    /** Remove vectors by ID selector
+     *
+     * @param sel  selector defining which vectors to remove
+     * @return     number of vectors removed
+     */
     size_t remove_ids(const IDSelector& sel) override;
 
+    /** Get the code packer for this index
+     *
+     * @return  pointer to the code packer
+     */
     CodePacker* get_CodePacker() const;
 
+    /** Merge another index into this one
+     *
+     * @param otherIndex  index to merge from
+     * @param add_id      ID offset to add to merged vectors
+     */
     void merge_from(Index& otherIndex, idx_t add_id = 0) override;
+
+    /** Check if another index is compatible for merging
+     *
+     * @param otherIndex  index to check compatibility with
+     */
     void check_compatible_for_merge(const Index& otherIndex) const override;
 
     /// standalone codes interface (but the codes are flattened)

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