faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexIVFFastScan.h

@@ -8,12 +8,12 @@
 #pragma once
 
 #include <faiss/IndexIVF.h>
-#include <faiss/impl/FastScanDistancePostProcessing.h>
+#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
 #include <faiss/utils/AlignedTable.h>
 
 namespace faiss {
 
-struct NormTableScaler;
 struct SIMDResultHandlerToFloat;
 struct Quantizer;
 
@@ -41,14 +41,14 @@ struct Quantizer;
 
 struct IndexIVFFastScan : IndexIVF {
     // size of the kernel
-    int bbs; // set at build time
+    int bbs = 0; // set at build time
 
-    size_t M;
-    size_t nbits;
-    size_t ksub;
+    size_t M = 0;
+    size_t nbits = 0;
+    size_t ksub = 0;
 
     // M rounded up to a multiple of 2
-    size_t M2;
+    size_t M2 = 0;
 
     // search-time implementation
     int implem = 0;
@@ -156,7 +156,7 @@ struct IndexIVFFastScan : IndexIVF {
      * @param context       processing context containing query factors
      * processor
      */
-    void compute_LUT_uint8(
+    virtual void compute_LUT_uint8(
             size_t n,
             const float* x,
             const CoarseQuantized& cq,
@@ -222,45 +222,31 @@ struct IndexIVFFastScan : IndexIVF {
             RangeSearchResult* result,
             const SearchParameters* params = nullptr) const override;
 
-    /** Create a KNN handler for this index type
+    /** Create a SIMD-dispatched scanner for knn search (IVF variant).
      *
-     * This method can be overridden by derived classes to provide
-     * specialized handlers (e.g., IVFRaBitQHeapHandler for RaBitQ indexes).
-     * Base implementation creates standard handlers based on k and impl.
+     * Returns a FastScanCodeScanner that bundles handler + accumulation
+     * kernel behind the SIMD dispatch boundary. ntotal is not passed
+     * because IVF sets it per-list via handler->ntotal.
+     * Derived classes override this to provide custom handlers
+     * (e.g. RaBitQ).
      *
-     * @param is_max        true for max-heap (inner product), false for
-     *                      min-heap (L2 distance)
-     * @param impl          implementation number:
-     *                      - even (10, 12, 14): use heap for top-k
-     *                      - odd (11, 13, 15): use reservoir sampling
-     * @param n             number of queries
-     * @param k             number of neighbors to find per query
-     * @param distances     output array for distances (n * k), will be
-     *                      populated by handler
-     * @param labels        output array for result IDs (n * k), will be
-     *                      populated by handler
-     * @param sel           optional ID selector to filter results (nullptr =
-     *                      no filtering)
-     * @param context       processing context containing additional data
-     * @param normalizers   optional array of size 2*n for converting quantized
-     *                      uint16 distances to float.
-     *
-     * @return Allocated result handler (caller owns and must delete).
-     *         Handler processes SIMD batches and populates distances/labels.
-     *
-     * @note The returned handler must be deleted by caller after use.
-     *       Typical usage: handler->begin() → process batches → handler->end()
+     * @param is_max       whether to use CMax comparator (true) or CMin
+     * @param n            number of queries
+     * @param k            number of neighbors to find
+     * @param distances    output distances array
+     * @param labels       output labels array
+     * @param sel          optional ID selector
+     * @return             scanner
      */
-    virtual SIMDResultHandlerToFloat* make_knn_handler(
+    virtual std::unique_ptr<FastScanCodeScanner> make_knn_scanner(
             bool is_max,
-            int impl,
             idx_t n,
             idx_t k,
             float* distances,
             idx_t* labels,
             const IDSelector* sel,
-            const FastScanDistancePostProcessing& context,
-            const float* normalizers = nullptr) const;
+            int impl = 0,
+            const FastScanDistancePostProcessing& context = {}) const;
 
     // dispatch to implementations and parallelize
     void search_dispatch_implem(
@@ -306,16 +292,18 @@ struct IndexIVFFastScan : IndexIVF {
             const IVFSearchParameters* params = nullptr) const;
 
     // implem 10 and 12 are not multithreaded internally, so
-    // export search stats
+    // export search stats. k is required for ensure_topk_full support.
     void search_implem_10(
             idx_t n,
             const float* x,
+            idx_t k,
             SIMDResultHandlerToFloat& handler,
             const CoarseQuantized& cq,
             size_t* ndis_out,
             size_t* nlist_out,
             const FastScanDistancePostProcessing& context,
-            const IVFSearchParameters* params = nullptr) const;
+            const IVFSearchParameters* params,
+            FastScanCodeScanner& scanner) const;
 
     void search_implem_12(
             idx_t n,
@@ -325,7 +313,8 @@ struct IndexIVFFastScan : IndexIVF {
             size_t* ndis_out,
             size_t* nlist_out,
             const FastScanDistancePostProcessing& context,
-            const IVFSearchParameters* params = nullptr) const;
+            const IVFSearchParameters* params,
+            FastScanCodeScanner& scanner) const;
 
     // implem 14 is multithreaded internally across nprobes and queries
     void search_implem_14(
@@ -359,29 +348,19 @@ 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
+    /** Get the size of the code portion packed by pq4_pack_codes.
      *
-     * Example use case:
-     * - IndexIVFRaBitQFastScan extracts factor data from codes for use
-     *   during search-time distance corrections
+     * Returns the number of bytes per vector that are interleaved into
+     * SIMD blocks by pq4_pack_codes, excluding any embedded metadata
+     * (e.g., RaBitQ factors). The meaning of these bytes depends on the
+     * quantizer: for PQ/AQ they are 4-bit sub-quantizer nibbles, for
+     * RaBitQ they are 1-bit-per-dimension sign bits packed into nibbles.
      *
-     * @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)
+     * Must be implemented by all derived classes.
      */
-    virtual void preprocess_code_metadata(
-            idx_t n,
-            const uint8_t* flat_codes,
-            idx_t start_global_idx);
+    virtual size_t fast_scan_code_size() const = 0;
 
+   protected:
     /** Get stride for interpreting codes during SIMD packing.
      *
      * The stride determines how to read codes when packing them into
@@ -399,6 +378,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 {

data/vendor/faiss/faiss/IndexIVFFlat.cpp

@@ -13,7 +13,6 @@
 
 #include <cinttypes>
 #include <cstdio>
-#include <numeric>
 
 #include <faiss/IndexFlat.h>
 
@@ -21,8 +20,14 @@
 #include <faiss/impl/IDSelector.h>
 
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/utils/distances.h>
+#include <faiss/impl/expanded_scanners.h>
+#include <faiss/utils/distances_dispatch.h>
 #include <faiss/utils/extra_distances.h>
+
+#define THE_SIMD_LEVEL SIMDLevel::NONE
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/simd_impl/IVFFlatScanner-inl.h>
+
 #include <faiss/utils/utils.h>
 
 namespace faiss {
@@ -32,18 +37,18 @@ namespace faiss {
  ******************************************/
 
 IndexIVFFlat::IndexIVFFlat(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
+        Index* quantizer_in,
+        size_t d_in,
+        size_t nlist_in,
         MetricType metric,
-        bool own_invlists)
+        bool own_invlists_in)
         : IndexIVF(
-                  quantizer,
-                  d,
-                  nlist,
-                  sizeof(float) * d,
+                  quantizer_in,
+                  d_in,
+                  nlist_in,
+                  sizeof(float) * d_in,
                   metric,
-                  own_invlists) {
+                  own_invlists_in) {
     code_size = sizeof(float) * d;
     by_residual = false;
 }
@@ -61,7 +66,7 @@ void IndexIVFFlat::add_core(
     FAISS_THROW_IF_NOT(is_trained);
     FAISS_THROW_IF_NOT(coarse_idx);
     FAISS_THROW_IF_NOT(!by_residual);
-    assert(invlists);
+    FAISS_THROW_IF_NOT_MSG(invlists, "invlists not initialized");
     direct_map.check_can_add(xids);
 
     int64_t n_add = 0;
@@ -74,7 +79,7 @@ void IndexIVFFlat::add_core(
         int rank = omp_get_thread_num();
 
         // each thread takes care of a subset of lists
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             idx_t list_no = coarse_idx[i];
 
             if (list_no >= 0 && list_no % nt == rank) {
@@ -110,7 +115,7 @@ void IndexIVFFlat::encode_vectors(
         memcpy(codes, x, code_size * n);
     } else {
         size_t coarse_size = coarse_code_size();
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             int64_t list_no = list_nos[i];
             uint8_t* code = codes + i * (code_size + coarse_size);
             const float* xi = x + i * d;
@@ -129,7 +134,7 @@ void IndexIVFFlat::decode_vectors(
         const uint8_t* codes,
         const idx_t* /*listnos*/,
         float* x) const {
-    for (size_t i = 0; i < n; i++) {
+    for (idx_t i = 0; i < n; i++) {
         const uint8_t* code = codes + i * code_size;
         float* xi = x + i * d;
         memcpy(xi, code, code_size);
@@ -138,114 +143,21 @@ void IndexIVFFlat::decode_vectors(
 
 void IndexIVFFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
     size_t coarse_size = coarse_code_size();
-    for (size_t i = 0; i < n; i++) {
+    for (idx_t i = 0; i < n; i++) {
         const uint8_t* code = bytes + i * (code_size + coarse_size);
         float* xi = x + i * d;
         memcpy(xi, code + coarse_size, code_size);
     }
 }
 
-namespace {
-
-template <typename VectorDistance, bool use_sel>
-struct IVFFlatScanner : InvertedListScanner {
-    VectorDistance vd;
-    using C = typename VectorDistance::C;
-
-    IVFFlatScanner(
-            const VectorDistance& vd,
-            bool store_pairs,
-            const IDSelector* sel)
-            : InvertedListScanner(store_pairs, sel), vd(vd) {
-        keep_max = vd.is_similarity;
-        code_size = vd.d * sizeof(float);
-    }
-
-    const float* xi;
-    void set_query(const float* query) override {
-        this->xi = query;
-    }
-
-    void set_list(idx_t list_no, float /* coarse_dis */) override {
-        this->list_no = list_no;
-    }
-
-    float distance_to_code(const uint8_t* code) const override {
-        const float* yj = (float*)code;
-        return vd(xi, yj);
-    }
-
-    size_t scan_codes(
-            size_t list_size,
-            const uint8_t* codes,
-            const idx_t* ids,
-            float* simi,
-            idx_t* idxi,
-            size_t k) const override {
-        const float* list_vecs = (const float*)codes;
-        size_t nup = 0;
-        for (size_t j = 0; j < list_size; j++) {
-            const float* yj = list_vecs + vd.d * j;
-            if (use_sel && !sel->is_member(ids[j])) {
-                continue;
-            }
-            float dis = vd(xi, yj);
-            if (C::cmp(simi[0], dis)) {
-                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                heap_replace_top<C>(k, simi, idxi, dis, id);
-                nup++;
-            }
-        }
-        return nup;
-    }
-
-    void scan_codes_range(
-            size_t list_size,
-            const uint8_t* codes,
-            const idx_t* ids,
-            float radius,
-            RangeQueryResult& res) const override {
-        const float* list_vecs = (const float*)codes;
-        for (size_t j = 0; j < list_size; j++) {
-            const float* yj = list_vecs + vd.d * j;
-            if (use_sel && !sel->is_member(ids[j])) {
-                continue;
-            }
-            float dis = vd(xi, yj);
-            if (C::cmp(radius, dis)) {
-                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                res.add(dis, id);
-            }
-        }
-    }
-};
-
-struct Run_get_InvertedListScanner {
-    using T = InvertedListScanner*;
-
-    template <class VD>
-    InvertedListScanner* f(
-            VD& vd,
-            const IndexIVFFlat* ivf,
-            bool store_pairs,
-            const IDSelector* sel) {
-        if (sel) {
-            return new IVFFlatScanner<VD, true>(vd, store_pairs, sel);
-        } else {
-            return new IVFFlatScanner<VD, false>(vd, store_pairs, sel);
-        }
-    }
-};
-
-} // anonymous namespace
-
 InvertedListScanner* IndexIVFFlat::get_InvertedListScanner(
         bool store_pairs,
         const IDSelector* sel,
         const IVFSearchParameters*) const {
-    Run_get_InvertedListScanner run;
-    return dispatch_VectorDistance(
-            d, metric_type, metric_arg, run, this, store_pairs, sel);
+    return with_VectorDistance(
+            d, metric_type, metric_arg, [&](auto vd) -> InvertedListScanner* {
+                return new IVFFlatScanner<decltype(vd)>(vd, store_pairs, sel);
+            });
 }
 
 void IndexIVFFlat::reconstruct_from_offset(
@@ -260,12 +172,17 @@ void IndexIVFFlat::reconstruct_from_offset(
  ******************************************/
 
 IndexIVFFlatDedup::IndexIVFFlatDedup(
-        Index* quantizer,
-        size_t d,
-        size_t nlist_,
-        MetricType metric_type,
-        bool own_invlists)
-        : IndexIVFFlat(quantizer, d, nlist_, metric_type, own_invlists) {}
+        Index* quantizer_in,
+        size_t d_in,
+        size_t nlist_in,
+        MetricType metric_type_in,
+        bool own_invlists_in)
+        : IndexIVFFlat(
+                  quantizer_in,
+                  d_in,
+                  nlist_in,
+                  metric_type_in,
+                  own_invlists_in) {}
 
 void IndexIVFFlatDedup::train(idx_t n, const float* x) {
     std::unordered_map<uint64_t, idx_t> map;
@@ -298,7 +215,7 @@ void IndexIVFFlatDedup::add_with_ids(
         const float* x,
         const idx_t* xids) {
     FAISS_THROW_IF_NOT(is_trained);
-    assert(invlists);
+    FAISS_THROW_IF_NOT_MSG(invlists, "invlists not initialized");
     FAISS_THROW_IF_NOT_MSG(
             direct_map.no(), "IVFFlatDedup not implemented with direct_map");
     std::unique_ptr<int64_t[]> idx(new int64_t[na]);
@@ -312,7 +229,7 @@ void IndexIVFFlatDedup::add_with_ids(
         int rank = omp_get_thread_num();
 
         // each thread takes care of a subset of lists
-        for (size_t i = 0; i < na; i++) {
+        for (idx_t i = 0; i < na; i++) {
             int64_t list_no = idx[i];
 
             if (list_no < 0 || list_no % nt != rank) {
@@ -371,7 +288,7 @@ void IndexIVFFlatDedup::search_preassigned(
         idx_t* labels,
         bool store_pairs,
         const IVFSearchParameters* params,
-        IndexIVFStats* stats) const {
+        IndexIVFStats* /*stats*/) const {
     FAISS_THROW_IF_NOT_MSG(
             !store_pairs, "store_pairs not supported in IVFDedup");
 
@@ -452,7 +369,7 @@ size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel) {
     std::vector<int64_t> toremove(nlist);
 
 #pragma omp parallel for
-    for (int64_t i = 0; i < nlist; i++) {
+    for (idx_t i = 0; i < static_cast<idx_t>(nlist); i++) {
         int64_t l0 = invlists->list_size(i), l = l0, j = 0;
         InvertedLists::ScopedIds idsi(invlists, i);
         while (j < l) {
@@ -480,7 +397,7 @@ size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel) {
     }
     // this will not run well in parallel on ondisk because of possible shrinks
     int64_t nremove = 0;
-    for (int64_t i = 0; i < nlist; i++) {
+    for (idx_t i = 0; i < static_cast<idx_t>(nlist); i++) {
         if (toremove[i] > 0) {
             nremove += toremove[i];
             invlists->resize(i, invlists->list_size(i) - toremove[i]);

data/vendor/faiss/faiss/IndexIVFFlat.h

@@ -62,6 +62,38 @@ struct IndexIVFFlat : IndexIVF {
     IndexIVFFlat();
 };
 
+template <typename VectorDistance>
+struct IVFFlatScanner : InvertedListScanner {
+    VectorDistance vd;
+    using C = typename VectorDistance::C;
+
+    IVFFlatScanner(
+            const VectorDistance& vd,
+            bool store_pairs,
+            const IDSelector* sel)
+            : InvertedListScanner(store_pairs, sel), vd(vd) {
+        keep_max = vd.is_similarity;
+        code_size = vd.d * sizeof(float);
+    }
+
+    const float* xi = nullptr;
+    void set_query(const float* query) override {
+        this->xi = query;
+    }
+
+    void set_list(idx_t list_no, float /* coarse_dis */) override {
+        this->list_no = list_no;
+    }
+
+    float distance_to_code(const uint8_t* code) const final;
+
+    size_t scan_codes(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            ResultHandler& handler) const override;
+};
+
 struct IndexIVFFlatDedup : IndexIVFFlat {
     /** Maps ids stored in the index to the ids of vectors that are
      *  the same. When a vector is unique, it does not appear in the