faiss 0.5.2 → 0.6.0

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -27,6 +27,7 @@
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/VisitedTable.h>
 #include <faiss/utils/random.h>
 #include <faiss/utils/sorting.h>
 
@@ -142,7 +143,7 @@ void hnsw_add_vertices(
 
 #pragma omp parallel if (i1 > i0 + 100)
             {
-                VisitedTable vt(ntotal);
+                VisitedTable vt(ntotal, hnsw.use_visited_hashset);
 
                 std::unique_ptr<DistanceComputer> dis(
                         storage_distance_computer(index_hnsw.storage));
@@ -265,7 +266,7 @@ void hnsw_search(
 
 #pragma omp parallel if (i1 - i0 > 1)
         {
-            VisitedTable vt(index->ntotal);
+            VisitedTable vt(index->ntotal, hnsw.use_visited_hashset);
             typename BlockResultHandler::SingleResultHandler res(bres);
 
             std::unique_ptr<DistanceComputer> dis(
@@ -333,6 +334,14 @@ void IndexHNSW::range_search(
     }
 }
 
+void IndexHNSW::search1(
+        const float* x,
+        ResultHandler& handler,
+        SearchParameters* params) const {
+    SingleQueryBlockResultHandler<HNSW::C, false> bres(handler);
+    hnsw_search(this, 1, x, bres, params);
+}
+
 void IndexHNSW::add(idx_t n, const float* x) {
     FAISS_THROW_IF_NOT_MSG(
             storage,
@@ -428,7 +437,7 @@ void IndexHNSW::search_level_0(
         std::unique_ptr<DistanceComputer> qdis(
                 storage_distance_computer(storage));
         HNSWStats search_stats;
-        VisitedTable vt(ntotal);
+        VisitedTable vt(ntotal, hnsw.use_visited_hashset);
         RH::SingleResultHandler res(bres);
 
 #pragma omp for
@@ -516,7 +525,7 @@ void IndexHNSW::init_level_0_from_entry_points(
 
 #pragma omp parallel
     {
-        VisitedTable vt(ntotal);
+        VisitedTable vt(ntotal, hnsw.use_visited_hashset);
 
         std::unique_ptr<DistanceComputer> dis(
                 storage_distance_computer(storage));
@@ -616,7 +625,7 @@ void IndexHNSW::link_singletons() {
 
     std::vector<float> recons(singletons.size() * d);
     for (int i = 0; i < singletons.size(); i++) {
-        FAISS_ASSERT(!"not implemented");
+        FAISS_ASSERT(false); // not implemented
     }
 }
 
@@ -653,33 +662,8 @@ IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
  * 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) {}
+        : IndexHNSWFlat(), cum_sums(), pano(0, 1, 1), num_panorama_levels(0) {}
 
 IndexHNSWFlatPanorama::IndexHNSWFlatPanorama(
         int d,
@@ -688,8 +672,7 @@ IndexHNSWFlatPanorama::IndexHNSWFlatPanorama(
         MetricType metric)
         : IndexHNSWFlat(d, M, metric),
           cum_sums(),
-          panorama_level_width(
-                  (d + num_panorama_levels - 1) / num_panorama_levels),
+          pano(d * sizeof(float), num_panorama_levels, 1),
           num_panorama_levels(num_panorama_levels) {
     // For now, we only support L2 distance.
     // Supporting dot product and cosine distance is a trivial addition
@@ -704,18 +687,8 @@ IndexHNSWFlatPanorama::IndexHNSWFlatPanorama(
 
 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);
-    }
-
+    cum_sums.resize((ntotal + n) * (pano.n_levels + 1));
+    pano.compute_cumulative_sums(cum_sums.data(), n0, n, x);
     IndexHNSWFlat::add(n, x);
 }
 
@@ -725,13 +698,13 @@ void IndexHNSWFlatPanorama::reset() {
 }
 
 void IndexHNSWFlatPanorama::permute_entries(const idx_t* perm) {
-    std::vector<float> new_cum_sums(ntotal * (num_panorama_levels + 1));
+    std::vector<float> new_cum_sums(ntotal * (pano.n_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));
+        memcpy(&new_cum_sums[i * (pano.n_levels + 1)],
+               &cum_sums[src * (pano.n_levels + 1)],
+               (pano.n_levels + 1) * sizeof(float));
     }
 
     std::swap(cum_sums, new_cum_sums);
@@ -903,7 +876,8 @@ void IndexHNSW2Level::search(
 
 #pragma omp parallel
         {
-            VisitedTable vt(ntotal);
+            // visited table (not hash set) for tri-state flags.
+            VisitedTable vt(ntotal, /*use_hashset=*/false);
             std::unique_ptr<DistanceComputer> dis(
                     storage_distance_computer(storage));
 

data/vendor/faiss/faiss/IndexHNSW.h

@@ -9,13 +9,15 @@
 
 #pragma once
 
+#include <optional>
 #include <vector>
-#include "faiss/Index.h"
 
+#include <faiss/Index.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexPQ.h>
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/impl/HNSW.h>
+#include <faiss/impl/Panorama.h>
 #include <faiss/utils/utils.h>
 
 namespace faiss {
@@ -47,6 +49,9 @@ struct IndexHNSW : Index {
     // used when GpuIndexCagra::copyFrom(IndexHNSWCagra*) is invoked.
     bool keep_max_size_level0 = false;
 
+    // See impl/VisitedTable.h.
+    std::optional<bool> use_visited_hashset;
+
     explicit IndexHNSW(int d = 0, int M = 32, MetricType metric = METRIC_L2);
     explicit IndexHNSW(Index* storage, int M = 32);
 
@@ -73,6 +78,12 @@ struct IndexHNSW : Index {
             RangeSearchResult* result,
             const SearchParameters* params = nullptr) const override;
 
+    /** search one vector with a custom result handler */
+    void search1(
+            const float* x,
+            ResultHandler& handler,
+            SearchParameters* params = nullptr) const override;
+
     void reconstruct(idx_t key, float* recons) const override;
 
     void reset() override;
@@ -164,20 +175,11 @@ struct IndexHNSWFlatPanorama : IndexHNSWFlat {
 
     /// 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);
+        return cum_sums.data() + i * (pano.n_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;
+    Panorama pano;
     const size_t num_panorama_levels;
 };
 

data/vendor/faiss/faiss/IndexIDMap.cpp

@@ -12,7 +12,7 @@
 #include <cinttypes>
 #include <cstdint>
 #include <cstdio>
-#include "faiss/Index.h"
+#include <stdexcept>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>

data/vendor/faiss/faiss/IndexIVF.cpp

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #include <faiss/IndexIVF.h>
 
 #include <omp.h>
@@ -27,6 +25,8 @@
 #include <faiss/impl/CodePacker.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/expanded_scanners.h>
 
 namespace faiss {
 
@@ -920,6 +920,52 @@ void IndexIVF::range_search_preassigned(
     stats->ndis += ndis;
 }
 
+void IndexIVF::search1(
+        const float* x,
+        ResultHandler& handler,
+        SearchParameters* params_in) const {
+    const IVFSearchParameters* params = nullptr;
+    const SearchParameters* quantizer_params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
+        quantizer_params = params->quantizer_params;
+    }
+    const size_t nprobe =
+            std::min(nlist, params ? params->nprobe : this->nprobe);
+    size_t nx = 1;
+    std::unique_ptr<idx_t[]> keys(new idx_t[nx * nprobe]);
+    std::unique_ptr<float[]> coarse_dis(new float[nx * nprobe]);
+
+    double t0 = getmillisecs();
+    quantizer->search(
+            nx, x, nprobe, coarse_dis.get(), keys.get(), quantizer_params);
+    indexIVF_stats.quantization_time += getmillisecs() - t0;
+
+    t0 = getmillisecs();
+    invlists->prefetch_lists(keys.get(), nx * nprobe);
+
+    std::unique_ptr<InvertedListScanner> scanner(
+            get_InvertedListScanner(false, nullptr, params));
+    scanner->set_query(x);
+
+    for (idx_t i = 0; i < nprobe; i++) {
+        idx_t key = keys[i];
+        if (key < 0 || invlists->is_empty(key)) {
+            continue;
+        }
+
+        scanner->set_list(key, coarse_dis[i]);
+        InvertedLists::ScopedCodes scodes(invlists, key);
+        InvertedLists::ScopedIds ids(invlists, key);
+        size_t list_size = invlists->list_size(key);
+
+        scanner->scan_codes(list_size, scodes.get(), ids.get(), handler);
+    }
+
+    indexIVF_stats.search_time += getmillisecs() - t0;
+}
+
 InvertedListScanner* IndexIVF::get_InvertedListScanner(
         bool /*store_pairs*/,
         const IDSelector* /* sel */,
@@ -1298,6 +1344,20 @@ IndexIVFStats indexIVF_stats;
  * InvertedListScanner
  *************************************************************************/
 
+// this gets expanded in expanded_scanners
+
+size_t InvertedListScanner::scan_codes(
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) const {
+    return run_scan_codes(*this, list_size, codes, ids, handler);
+}
+
+void InvertedListScanner::set_list(idx_t list_no_in, float /* coarse_dis */) {
+    this->list_no = list_no_in;
+}
+
 size_t InvertedListScanner::scan_codes(
         size_t list_size,
         const uint8_t* codes,
@@ -1305,46 +1365,15 @@ size_t InvertedListScanner::scan_codes(
         float* simi,
         idx_t* idxi,
         size_t k) const {
-    size_t nup = 0;
-
     if (!keep_max) {
-        for (size_t j = 0; j < list_size; j++) {
-            if (sel != nullptr) {
-                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                if (!sel->is_member(id)) {
-                    codes += code_size;
-                    continue;
-                }
-            }
-
-            float dis = distance_to_code(codes);
-            if (dis < simi[0]) {
-                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                maxheap_replace_top(k, simi, idxi, dis, id);
-                nup++;
-            }
-            codes += code_size;
-        }
+        using C = CMax<float, idx_t>;
+        HeapResultHandler<C, false> handler(k, simi, idxi);
+        return scan_codes(list_size, codes, ids, handler);
     } else {
-        for (size_t j = 0; j < list_size; j++) {
-            if (sel != nullptr) {
-                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                if (!sel->is_member(id)) {
-                    codes += code_size;
-                    continue;
-                }
-            }
-
-            float dis = distance_to_code(codes);
-            if (dis > simi[0]) {
-                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                minheap_replace_top(k, simi, idxi, dis, id);
-                nup++;
-            }
-            codes += code_size;
-        }
+        using C = CMin<float, idx_t>;
+        HeapResultHandler<C, false> handler(k, simi, idxi);
+        return scan_codes(list_size, codes, ids, handler);
     }
-    return nup;
 }
 
 size_t InvertedListScanner::iterate_codes(
@@ -1386,16 +1415,14 @@ void InvertedListScanner::scan_codes_range(
         const idx_t* ids,
         float radius,
         RangeQueryResult& res) const {
-    for (size_t j = 0; j < list_size; j++) {
-        float dis = distance_to_code(codes);
-        bool keep = !keep_max
-                ? dis < radius
-                : dis > radius; // TODO templatize to remove this test
-        if (keep) {
-            int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-            res.add(dis, id);
-        }
-        codes += code_size;
+    if (!keep_max) {
+        using C = CMax<float, idx_t>;
+        RangeResultHandler<C, false> handler(&res, radius);
+        scan_codes(list_size, codes, ids, handler);
+    } else {
+        using C = CMin<float, idx_t>;
+        RangeResultHandler<C, false> handler(&res, radius);
+        scan_codes(list_size, codes, ids, handler);
     }
 }
 

data/vendor/faiss/faiss/IndexIVF.h

@@ -11,9 +11,6 @@
 #define FAISS_INDEX_IVF_H
 
 #include <stdint.h>
-#include <memory>
-#include <unordered_map>
-#include <vector>
 
 #include <faiss/Clustering.h>
 #include <faiss/Index.h>
@@ -325,6 +322,12 @@ struct IndexIVF : Index, IndexIVFInterface {
             RangeSearchResult* result,
             const SearchParameters* params = nullptr) const override;
 
+    /** search one vector with a custom result handler */
+    void search1(
+            const float* x,
+            ResultHandler& handler,
+            SearchParameters* params = nullptr) const override;
+
     /** Get a scanner for this index (store_pairs means ignore labels)
      *
      * The default search implementation uses this to compute the distances.
@@ -492,7 +495,7 @@ struct InvertedListScanner {
     virtual void set_query(const float* query_vector) = 0;
 
     /// following codes come from this inverted list
-    virtual void set_list(idx_t list_no, float coarse_dis) = 0;
+    virtual void set_list(idx_t list_no, float coarse_dis);
 
     /// compute a single query-to-code distance
     virtual float distance_to_code(const uint8_t* code) const = 0;
@@ -543,6 +546,13 @@ struct InvertedListScanner {
             RangeQueryResult& result,
             size_t& list_size) const;
 
+    // accumulate results with a ResultHandler
+    virtual size_t scan_codes(
+            size_t n,
+            const uint8_t* codes,
+            const idx_t* ids,
+            ResultHandler& handler) const;
+
     virtual ~InvertedListScanner() {}
 };
 

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp

@@ -17,6 +17,7 @@
 #include <faiss/impl/FastScanDistancePostProcessing.h>
 #include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/pq4_fast_scan.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/invlists/BlockInvertedLists.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/quantize_lut.h>
@@ -405,18 +406,20 @@ void IndexIVFAdditiveQuantizerFastScan::compute_LUT(
         // bias = coef * <q, c>
         // NOTE: q^2 is not added to `biases`
         biases.resize(n * nprobe);
+        with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel
-        {
-            std::vector<float> centroid(d);
-            float* c = centroid.data();
+            {
+                std::vector<float> centroid(d);
+                float* c = centroid.data();
 
 #pragma omp for
-            for (idx_t ij = 0; ij < n * nprobe; ij++) {
-                int i = ij / nprobe;
-                quantizer->reconstruct(cq.ids[ij], c);
-                biases[ij] = coef * fvec_inner_product(c, x + i * d, d);
+                for (idx_t ij = 0; ij < n * nprobe; ij++) {
+                    int i = ij / nprobe;
+                    quantizer->reconstruct(cq.ids[ij], c);
+                    biases[ij] = coef * fvec_inner_product<SL>(c, x + i * d, d);
+                }
             }
-        }
+        });
     }
 
     if (metric_type == METRIC_L2) {

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h

@@ -37,13 +37,13 @@ namespace faiss {
 struct IndexIVFAdditiveQuantizerFastScan : IndexIVFFastScan {
     using Search_type_t = AdditiveQuantizer::Search_type_t;
 
-    AdditiveQuantizer* aq;
+    AdditiveQuantizer* aq{};
 
     bool rescale_norm = false;
     int norm_scale = 1;
 
     // max number of training vectors
-    size_t max_train_points;
+    size_t max_train_points{};
 
     IndexIVFAdditiveQuantizerFastScan(
             Index* quantizer,

data/vendor/faiss/faiss/IndexIVFFastScan.cpp

@@ -95,18 +95,19 @@ IndexIVFFastScan::~IndexIVFFastScan() = default;
  * Code management functions
  *********************************************************/
 
-void IndexIVFFastScan::preprocess_code_metadata(
-        idx_t /* n */,
-        const uint8_t* /* flat_codes */,
-        idx_t /* start_global_idx */) {
-    // Default: no-op
-}
-
 size_t IndexIVFFastScan::code_packing_stride() const {
     // Default: use standard M-byte stride
     return 0;
 }
 
+size_t IndexIVFFastScan::get_block_stride() const {
+    std::unique_ptr<CodePacker> packer(get_CodePacker());
+    FAISS_THROW_IF_NOT_MSG(
+            packer->nvec == static_cast<size_t>(bbs),
+            "CodePacker must pack bbs vectors per block for fast-scan");
+    return packer->block_size;
+}
+
 void IndexIVFFastScan::add_with_ids(
         idx_t n,
         const float* x,
@@ -148,9 +149,6 @@ void IndexIVFFastScan::add_with_ids(
     AlignedTable<uint8_t> flat_codes(n * code_size);
     encode_vectors(n, x, idx.get(), flat_codes.get());
 
-    // Allow subclasses to preprocess metadata before packing
-    preprocess_code_metadata(n, flat_codes.get(), ntotal);
-
     DirectMapAdd dm_adder(direct_map, n, xids);
     BlockInvertedLists* bil = dynamic_cast<BlockInvertedLists*>(invlists);
     FAISS_THROW_IF_NOT_MSG(bil, "only block inverted lists supported");
@@ -206,7 +204,11 @@ void IndexIVFFastScan::add_with_ids(
                 bbs,
                 M2,
                 bil->codes[list_no].data(),
-                pack_stride);
+                pack_stride,
+                get_block_stride());
+
+        postprocess_packed_codes(
+                list_no, list_size, i1 - i0, list_codes.data());
 
         i0 = i1;
     }
@@ -1029,7 +1031,8 @@ void IndexIVFFastScan::search_implem_10(
                     codes.get(),
                     LUT,
                     handler,
-                    context.norm_scaler);
+                    context.norm_scaler,
+                    get_block_stride());
 
             ndis += ls;
             nlist_visited++;
@@ -1180,7 +1183,8 @@ void IndexIVFFastScan::search_implem_12(
                 codes.get(),
                 LUT.get(),
                 handler,
-                context.norm_scaler);
+                context.norm_scaler,
+                get_block_stride());
         // prepare for next loop
         i0 = i1;
     }
@@ -1403,7 +1407,8 @@ void IndexIVFFastScan::search_implem_14(
                     codes.get(),
                     LUT.get(),
                     *handler.get(),
-                    context.norm_scaler);
+                    context.norm_scaler,
+                    get_block_stride());
         }
 
         // labels is in-place for HeapHC
@@ -1519,6 +1524,12 @@ void IndexIVFFastScan::sa_decode(idx_t n, const uint8_t* codes, float* x)
     }
 }
 
+void IndexIVFFastScan::postprocess_packed_codes(
+        idx_t /*list_no*/,
+        size_t /*list_offset*/,
+        size_t /*n_added*/,
+        const uint8_t* /*flat_codes*/) {}
+
 IVFFastScanStats IVFFastScan_stats;
 
 } // namespace faiss

data/vendor/faiss/faiss/IndexIVFFastScan.h

@@ -360,28 +360,6 @@ struct IndexIVFFastScan : IndexIVF {
     void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
 
    protected:
-    /** Preprocess metadata from encoded vectors before packing.
-     *
-     * Called during add_with_ids after encode_vectors but before codes
-     * are packed into SIMD-friendly blocks. Subclasses can override to
-     * extract and store metadata embedded in codes or perform other
-     * pre-packing operations.
-     *
-     * Default implementation: no-op
-     *
-     * Example use case:
-     * - IndexIVFRaBitQFastScan extracts factor data from codes for use
-     *   during search-time distance corrections
-     *
-     * @param n                  number of vectors encoded
-     * @param flat_codes         encoded vectors (n * code_size bytes)
-     * @param start_global_idx   starting global index (ntotal before add)
-     */
-    virtual void preprocess_code_metadata(
-            idx_t n,
-            const uint8_t* flat_codes,
-            idx_t start_global_idx);
-
     /** Get stride for interpreting codes during SIMD packing.
      *
      * The stride determines how to read codes when packing them into
@@ -399,6 +377,32 @@ struct IndexIVFFastScan : IndexIVF {
      *         - >0: use custom stride (e.g., code_size for embedded metadata)
      */
     virtual size_t code_packing_stride() const;
+
+   public:
+    /** Get stride in bytes between consecutive SIMD blocks.
+     *
+     * Derived from get_CodePacker()->block_size so that there is a
+     * single source of truth for the block layout.
+     *
+     * @return stride in bytes
+     */
+    size_t get_block_stride() const;
+
+    /** Post-process packed codes after pq4_pack_codes_range.
+     *
+     * Called during add_with_ids after codes have been packed into
+     * SIMD-friendly blocks.
+     *
+     * @param list_no       inverted list number
+     * @param list_offset   starting offset within the list (pre-existing size)
+     * @param n_added       number of vectors added in this batch
+     * @param flat_codes    encoded vectors for this batch (n_added * code_size)
+     */
+    virtual void postprocess_packed_codes(
+            idx_t list_no,
+            size_t list_offset,
+            size_t n_added,
+            const uint8_t* flat_codes);
 };
 
 struct IVFFastScanStats {