faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/IndexIVFFastScan.cpp

@@ -17,11 +17,10 @@
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/FastScanDistancePostProcessing.h>
-#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/RaBitQUtils.h>
-#include <faiss/impl/pq4_fast_scan.h>
-#include <faiss/impl/simd_result_handlers.h>
+#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
+#include <faiss/impl/fast_scan/simd_result_handlers.h>
 #include <faiss/invlists/BlockInvertedLists.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/quantize_lut.h>
@@ -29,20 +28,24 @@
 
 namespace faiss {
 
-using namespace simd_result_handlers;
-
 inline size_t roundup(size_t a, size_t b) {
     return (a + b - 1) / b * b;
 }
 
 IndexIVFFastScan::IndexIVFFastScan(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
-        size_t code_size,
+        Index* quantizer_in,
+        size_t d_in,
+        size_t nlist_in,
+        size_t code_size_in,
         MetricType metric,
-        bool own_invlists)
-        : IndexIVF(quantizer, d, nlist, code_size, metric, own_invlists) {
+        bool own_invlists_in)
+        : IndexIVF(
+                  quantizer_in,
+                  d_in,
+                  nlist_in,
+                  code_size_in,
+                  metric,
+                  own_invlists_in) {
     // unlike other indexes, we prefer no residuals for performance reasons.
     by_residual = false;
     FAISS_THROW_IF_NOT(metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT);
@@ -56,35 +59,40 @@ IndexIVFFastScan::IndexIVFFastScan() {
 }
 
 void IndexIVFFastScan::init_fastscan(
-        Quantizer* fine_quantizer,
-        size_t M,
+        Quantizer* fine_quantizer_in,
+        size_t M_in,
         size_t nbits_init,
-        size_t nlist,
+        size_t nlist_in,
         MetricType /* metric */,
         int bbs_2,
-        bool own_invlists) {
+        bool own_invlists_in) {
     FAISS_THROW_IF_NOT(bbs_2 % 32 == 0);
     FAISS_THROW_IF_NOT(nbits_init == 4);
-    FAISS_THROW_IF_NOT(fine_quantizer->d == d);
+    FAISS_THROW_IF_NOT(fine_quantizer_in->d == static_cast<size_t>(d));
 
-    this->fine_quantizer = fine_quantizer;
-    this->M = M;
+    this->fine_quantizer = fine_quantizer_in;
+    this->M = M_in;
     this->nbits = nbits_init;
     this->bbs = bbs_2;
     ksub = (1 << nbits_init);
-    M2 = roundup(M, 2);
+    M2 = roundup(M_in, 2);
     code_size = M2 / 2;
-    FAISS_THROW_IF_NOT(code_size == fine_quantizer->code_size);
+    FAISS_THROW_IF_NOT(code_size == fine_quantizer_in->code_size);
 
     is_trained = false;
-    if (own_invlists) {
-        replace_invlists(new BlockInvertedLists(nlist, get_CodePacker()), true);
+    if (own_invlists_in) {
+        replace_invlists(
+                new BlockInvertedLists(nlist_in, get_CodePacker()), true);
     }
 }
 
 void IndexIVFFastScan::init_code_packer() {
     auto bil = dynamic_cast<BlockInvertedLists*>(invlists);
-    FAISS_THROW_IF_NOT(bil);
+    if (!bil) {
+        // invlists is not block-packed (e.g., when own_invlists=false).
+        // Nothing to do — the caller manages inverted lists externally.
+        return;
+    }
     delete bil->packer; // in case there was one before
     bil->packer = get_CodePacker();
 }
@@ -239,22 +247,22 @@ void estimators_from_tables_generic(
         int64_t* heap_ids,
         const FastScanDistancePostProcessing& context) {
     using accu_t = typename C::T;
-    size_t nscale = context.norm_scaler ? context.norm_scaler->nscale : 0;
+    size_t nscale = context.pq2x4_scale ? 2 : 0;
     for (size_t j = 0; j < ncodes; ++j) {
         BitstringReader bsr(codes + j * index.code_size, index.code_size);
         accu_t dis = bias;
         const dis_t* __restrict dt = dis_table;
 
         for (size_t m = 0; m < index.M - nscale; m++) {
-            uint64_t c = bsr.read(index.nbits);
+            uint64_t c = bsr.read(static_cast<int>(index.nbits));
             dis += dt[c];
             dt += index.ksub;
         }
 
-        if (context.norm_scaler) {
+        if (nscale) {
             for (size_t m = 0; m < nscale; m++) {
-                uint64_t c = bsr.read(index.nbits);
-                dis += context.norm_scaler->scale_one(dt[c]);
+                uint64_t c = bsr.read(static_cast<int>(index.nbits));
+                dis += dt[c] * context.pq2x4_scale;
                 dt += index.ksub;
             }
         }
@@ -285,33 +293,33 @@ void IndexIVFFastScan::compute_LUT_uint8(
     AlignedTable<float> biases_float;
 
     compute_LUT(n, x, cq, dis_tables_float, biases_float, context);
-    size_t nprobe = cq.nprobe;
+    size_t cur_nprobe = cq.nprobe;
     bool lut_is_3d = lookup_table_is_3d();
     size_t dim123 = ksub * M;
     size_t dim123_2 = ksub * M2;
     if (lut_is_3d) {
-        dim123 *= nprobe;
-        dim123_2 *= nprobe;
+        dim123 *= cur_nprobe;
+        dim123_2 *= cur_nprobe;
     }
     dis_tables.resize(n * dim123_2);
     if (biases_float.get()) {
-        biases.resize(n * nprobe);
+        biases.resize(n * cur_nprobe);
     }
 
     // OMP for MSVC requires i to have signed integral type
 #pragma omp parallel for if (n > 100)
-    for (int64_t i = 0; i < n; i++) {
+    for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
         const float* t_in = dis_tables_float.get() + i * dim123;
         const float* b_in = nullptr;
         uint8_t* t_out = dis_tables.get() + i * dim123_2;
         uint16_t* b_out = nullptr;
         if (biases_float.get()) {
-            b_in = biases_float.get() + i * nprobe;
-            b_out = biases.get() + i * nprobe;
+            b_in = biases_float.get() + i * cur_nprobe;
+            b_out = biases.get() + i * cur_nprobe;
         }
 
         quantize_LUT_and_bias(
-                nprobe,
+                cur_nprobe,
                 M,
                 ksub,
                 lut_is_3d,
@@ -358,10 +366,14 @@ void IndexIVFFastScan::search_preassigned(
         bool store_pairs,
         const IVFSearchParameters* params,
         IndexIVFStats* stats) const {
-    size_t nprobe = this->nprobe;
+    size_t cur_nprobe = this->nprobe;
     if (params) {
-        FAISS_THROW_IF_NOT(params->max_codes == 0);
-        nprobe = params->nprobe;
+        // Range-search-only option.
+        FAISS_THROW_IF_NOT_MSG(
+                params->max_empty_result_buckets == 0,
+                "max_empty_result_buckets is a range-search knob and is "
+                "not honored by fastscan knn search");
+        cur_nprobe = params->nprobe;
     }
 
     FAISS_THROW_IF_NOT_MSG(
@@ -370,7 +382,7 @@ void IndexIVFFastScan::search_preassigned(
     FAISS_THROW_IF_NOT(k > 0);
     FastScanDistancePostProcessing empty_context{};
 
-    const CoarseQuantized cq = {nprobe, centroid_dis, assign};
+    const CoarseQuantized cq = {cur_nprobe, centroid_dis, assign};
     search_dispatch_implem(
             n, x, k, distances, labels, cq, empty_context, params);
 }
@@ -381,45 +393,35 @@ void IndexIVFFastScan::range_search(
         float radius,
         RangeSearchResult* result,
         const SearchParameters* params_in) const {
-    size_t nprobe = this->nprobe;
+    size_t cur_nprobe = this->nprobe;
     const IVFSearchParameters* params = nullptr;
     if (params_in) {
         params = dynamic_cast<const IVFSearchParameters*>(params_in);
         FAISS_THROW_IF_NOT_MSG(
                 params, "IndexIVFFastScan params have incorrect type");
-        nprobe = params->nprobe;
+        // k-NN-only options.
+        FAISS_THROW_IF_NOT_MSG(
+                params->max_lists_num == 0,
+                "max_lists_num is a knn knob and is not honored by "
+                "fastscan range search");
+        FAISS_THROW_IF_NOT_MSG(
+                !params->ensure_topk_full,
+                "ensure_topk_full is a knn knob and is not honored by "
+                "fastscan range search");
+        FAISS_THROW_IF_NOT_MSG(
+                params->max_codes == 0,
+                "max_codes is not honored by fastscan range search");
+        cur_nprobe = params->nprobe;
     }
     FastScanDistancePostProcessing empty_context{};
 
-    const CoarseQuantized cq = {nprobe, nullptr, nullptr};
+    const CoarseQuantized cq = {cur_nprobe, nullptr, nullptr};
     range_search_dispatch_implem(
             n, x, radius, *result, cq, empty_context, params);
 }
 
 namespace {
 
-template <class C>
-ResultHandlerCompare<C, true>* make_knn_handler_fixC(
-        int impl,
-        idx_t n,
-        idx_t k,
-        float* distances,
-        idx_t* labels,
-        const IDSelector* sel,
-        const float* normalizers) {
-    using HeapHC = HeapHandler<C, true>;
-    using ReservoirHC = ReservoirHandler<C, true>;
-    using SingleResultHC = SingleResultHandler<C, true>;
-
-    if (k == 1) {
-        return new SingleResultHC(n, 0, distances, labels, sel);
-    } else if (impl % 2 == 0) {
-        return new HeapHC(n, 0, k, distances, labels, sel, normalizers);
-    } else /* if (impl % 2 == 1) */ {
-        return new ReservoirHC(n, 0, k, 2 * k, distances, labels, sel);
-    }
-}
-
 using CoarseQuantized = IndexIVFFastScan::CoarseQuantized;
 
 struct CoarseQuantizedWithBuffer : CoarseQuantized {
@@ -454,8 +456,8 @@ struct CoarseQuantizedWithBuffer : CoarseQuantized {
 
 struct CoarseQuantizedSlice : CoarseQuantizedWithBuffer {
     const size_t i0, i1;
-    CoarseQuantizedSlice(const CoarseQuantized& cq, size_t i0, size_t i1)
-            : CoarseQuantizedWithBuffer(cq), i0(i0), i1(i1) {
+    CoarseQuantizedSlice(const CoarseQuantized& cq, size_t i0_in, size_t i1_in)
+            : CoarseQuantizedWithBuffer(cq), i0(i0_in), i1(i1_in) {
         if (done()) {
             dis += nprobe * i0;
             ids += nprobe * i0;
@@ -473,20 +475,20 @@ struct CoarseQuantizedSlice : CoarseQuantizedWithBuffer {
 int compute_search_nslice(
         const IndexIVFFastScan* index,
         size_t n,
-        size_t nprobe) {
+        size_t cur_nprobe) {
     int nslice;
-    if (n <= omp_get_max_threads()) {
-        nslice = n;
+    if (n <= static_cast<size_t>(omp_get_max_threads())) {
+        nslice = static_cast<int>(n);
     } else if (index->lookup_table_is_3d()) {
         // make sure we don't make too big LUT tables
-        size_t lut_size_per_query = index->M * index->ksub * nprobe *
+        size_t lut_size_per_query = index->M * index->ksub * cur_nprobe *
                 (sizeof(float) + sizeof(uint8_t));
 
         size_t max_lut_size = precomputed_table_max_bytes;
         // how many queries we can handle within mem budget
         size_t nq_ok = std::max(max_lut_size / lut_size_per_query, size_t(1));
-        nslice = roundup(
-                std::max(size_t(n / nq_ok), size_t(1)), omp_get_max_threads());
+        nslice = static_cast<int>(roundup(
+                std::max(size_t(n / nq_ok), size_t(1)), omp_get_max_threads()));
     } else {
         // LUTs unlikely to be a limiting factor
         nslice = omp_get_max_threads();
@@ -496,23 +498,25 @@ int compute_search_nslice(
 
 } // namespace
 
-SIMDResultHandlerToFloat* IndexIVFFastScan::make_knn_handler(
+std::unique_ptr<FastScanCodeScanner> IndexIVFFastScan::make_knn_scanner(
         bool is_max,
-        int impl,
         idx_t n,
         idx_t k,
         float* distances,
         idx_t* labels,
         const IDSelector* sel,
-        const FastScanDistancePostProcessing&,
-        const float* normalizers) const {
-    if (is_max) {
-        return make_knn_handler_fixC<CMax<uint16_t, int64_t>>(
-                impl, n, k, distances, labels, sel, normalizers);
-    } else {
-        return make_knn_handler_fixC<CMin<uint16_t, int64_t>>(
-                impl, n, k, distances, labels, sel, normalizers);
-    }
+        int impl,
+        const FastScanDistancePostProcessing&) const {
+    return make_fast_scan_knn_scanner(
+            is_max,
+            impl,
+            n,
+            0,
+            k,
+            distances,
+            labels,
+            sel,
+            /*with_id_map=*/true);
 }
 
 void IndexIVFFastScan::search_dispatch_implem(
@@ -524,13 +528,12 @@ void IndexIVFFastScan::search_dispatch_implem(
         const CoarseQuantized& cq_in,
         const FastScanDistancePostProcessing& context,
         const IVFSearchParameters* params) const {
-    const idx_t nprobe = params ? params->nprobe : this->nprobe;
+    const idx_t cur_nprobe = params ? params->nprobe : this->nprobe;
     const IDSelector* sel = (params) ? params->sel : nullptr;
     const SearchParameters* quantizer_params =
             params ? params->quantizer_params : nullptr;
 
     bool is_max = !is_similarity_metric(metric_type);
-    using RH = SIMDResultHandlerToFloat;
 
     if (n == 0) {
         return;
@@ -539,8 +542,14 @@ void IndexIVFFastScan::search_dispatch_implem(
     // actual implementation used
     int impl = implem;
 
+    // Early-stop k-NN options require the per-query implementations.
+    const bool any_early_term_knob = params &&
+            (params->max_codes != 0 || params->max_lists_num != 0 ||
+             params->ensure_topk_full);
+
     if (impl == 0) {
-        if (bbs == 32) {
+        // Auto-select the per-query path when early-stop budgets are used.
+        if (bbs == 32 && !any_early_term_knob) {
             impl = 12;
         } else {
             impl = 10;
@@ -557,15 +566,24 @@ void IndexIVFFastScan::search_dispatch_implem(
         impl -= 100;
     }
 
+    if (any_early_term_knob) {
+        FAISS_THROW_IF_NOT_MSG(
+                impl == 10 || impl == 11,
+                "max_codes / max_lists_num / ensure_topk_full are only "
+                "supported by IndexIVFFastScan implem 10/11; set "
+                "index.implem = 10 (or 11 for k>20) explicitly, or leave it "
+                "at the default 0");
+    }
+
     CoarseQuantizedWithBuffer cq(cq_in);
-    cq.nprobe = nprobe;
+    cq.nprobe = cur_nprobe;
 
     if (!cq.done() && !multiple_threads) {
         // we do the coarse quantization here execpt when search is
         // sliced over threads (then it is more efficient to have each thread do
         // its own coarse quantization)
         cq.quantize(quantizer, n, x, quantizer_params);
-        invlists->prefetch_lists(cq.ids, n * cq.nprobe);
+        invlists->prefetch_lists(cq.ids, static_cast<int>(n * cq.nprobe));
     }
 
     if (impl == 1) {
@@ -588,43 +606,38 @@ void IndexIVFFastScan::search_dispatch_implem(
         size_t ndis = 0, nlist_visited = 0;
 
         if (!multiple_threads) {
-            // clang-format off
-            if (impl == 12 || impl == 13) {
-                std::unique_ptr<RH> handler(
-                    static_cast<RH*>(this->make_knn_handler(
-                        is_max,
-                        impl,
-                        n,
-                        k,
-                        distances,
-                        labels,
-                        sel,
-                        context))
-                );
-                search_implem_12(
-                        n, x, *handler.get(),
-                        cq, &ndis, &nlist_visited, context, params);
-            } else if (impl == 14 || impl == 15) {
+            if (impl == 14 || impl == 15) {
                 search_implem_14(
-                        n, x, k, distances, labels,
-                        cq, impl, context, params);
+                        n, x, k, distances, labels, cq, impl, context, params);
             } else {
-                std::unique_ptr<RH> handler(
-                    static_cast<RH*>(this->make_knn_handler(
-                        is_max,
-                        impl,
-                        n,
-                        k,
-                        distances,
-                        labels,
-                        sel,
-                        context))
-                );
-                search_implem_10(
-                        n, x, *handler.get(), cq,
-                        &ndis, &nlist_visited, context, params);
+                auto scanner = make_knn_scanner(
+                        is_max, n, k, distances, labels, sel, impl, context);
+                auto* handler = scanner->handler();
+                if (impl == 12 || impl == 13) {
+                    search_implem_12(
+                            n,
+                            x,
+                            *handler,
+                            cq,
+                            &ndis,
+                            &nlist_visited,
+                            context,
+                            params,
+                            *scanner);
+                } else {
+                    search_implem_10(
+                            n,
+                            x,
+                            k,
+                            *handler,
+                            cq,
+                            &ndis,
+                            &nlist_visited,
+                            context,
+                            params,
+                            *scanner);
+                }
             }
-            // clang-format on
         } else {
             // explicitly slice over threads
             int nslice = compute_search_nslice(this, n, cq.nprobe);
@@ -649,30 +662,43 @@ void IndexIVFFastScan::search_dispatch_implem(
                     // pointer
                     FastScanDistancePostProcessing thread_context = context;
                     if (thread_context.query_factors != nullptr) {
-                        thread_context.query_factors += i0 * nprobe;
+                        thread_context.query_factors += i0 * cur_nprobe;
                     }
 
-                    std::unique_ptr<RH> handler(
-                            static_cast<RH*>(this->make_knn_handler(
-                                    is_max,
-                                    impl,
-                                    i1 - i0,
-                                    k,
-                                    dis_i,
-                                    lab_i,
-                                    sel,
-                                    thread_context)));
-                    // clang-format off
+                    auto scanner = make_knn_scanner(
+                            is_max,
+                            i1 - i0,
+                            k,
+                            dis_i,
+                            lab_i,
+                            sel,
+                            impl,
+                            thread_context);
+                    auto* handler = scanner->handler();
                     if (impl == 12 || impl == 13) {
                         search_implem_12(
-                                i1 - i0, x + i0 * d, *handler.get(),
-                                cq_i, &ndis, &nlist_visited, thread_context, params);
+                                i1 - i0,
+                                x + i0 * d,
+                                *handler,
+                                cq_i,
+                                &ndis,
+                                &nlist_visited,
+                                thread_context,
+                                params,
+                                *scanner);
                     } else {
                         search_implem_10(
-                                i1 - i0, x + i0 * d, *handler.get(),
-                                cq_i, &ndis, &nlist_visited, thread_context, params);
+                                i1 - i0,
+                                x + i0 * d,
+                                k,
+                                *handler,
+                                cq_i,
+                                &ndis,
+                                &nlist_visited,
+                                thread_context,
+                                params,
+                                *scanner);
                     }
-                    // clang-format on
                 }
             }
         }
@@ -702,11 +728,23 @@ void IndexIVFFastScan::range_search_dispatch_implem(
     if (n == 0) {
         return;
     }
+    // FastScan range early-stop budget: enabled only for ordered per-query
+    // scanning below.
+    const bool use_empty_result_early_exit =
+            params && params->max_empty_result_buckets != 0;
+    const int pmode = this->parallel_mode & ~PARALLEL_MODE_NO_HEAP_INIT;
+    FAISS_THROW_IF_NOT_MSG(
+            !use_empty_result_early_exit || pmode == 0,
+            "max_empty_result_buckets supported only for parallel_mode = 0");
+
     // actual implementation used
     int impl = implem;
 
     if (impl == 0) {
-        if (bbs == 32) {
+        if (use_empty_result_early_exit) {
+            // Empty-bucket early stop needs per-query probe order.
+            impl = 10;
+        } else if (bbs == 32) {
             impl = 12;
         } else {
             impl = 10;
@@ -722,28 +760,44 @@ void IndexIVFFastScan::range_search_dispatch_implem(
         impl -= 100;
     }
 
+    FAISS_THROW_IF_NOT_MSG(
+            !use_empty_result_early_exit || impl == 10,
+            "max_empty_result_buckets is only supported by "
+            "IndexIVFFastScan range-search implem 10");
+
     if (!multiple_threads && !cq.done()) {
         cq.quantize(quantizer, n, x, quantizer_params);
-        invlists->prefetch_lists(cq.ids, n * cq.nprobe);
+        invlists->prefetch_lists(cq.ids, static_cast<int>(n * cq.nprobe));
     }
 
     size_t ndis = 0, nlist_visited = 0;
 
     if (!multiple_threads) { // single thread
-        std::unique_ptr<SIMDResultHandlerToFloat> handler;
-        if (is_max) {
-            handler.reset(new RangeHandler<CMax<uint16_t, int64_t>, true>(
-                    rres, radius, 0, sel));
-        } else {
-            handler.reset(new RangeHandler<CMin<uint16_t, int64_t>, true>(
-                    rres, radius, 0, sel));
-        }
+        auto scanner = make_range_scanner(is_max, rres, radius, 0, sel);
+        auto* handler = scanner->handler();
         if (impl == 12) {
             search_implem_12(
-                    n, x, *handler.get(), cq, &ndis, &nlist_visited, context);
+                    n,
+                    x,
+                    *handler,
+                    cq,
+                    &ndis,
+                    &nlist_visited,
+                    context,
+                    nullptr,
+                    *scanner);
         } else if (impl == 10) {
             search_implem_10(
-                    n, x, *handler.get(), cq, &ndis, &nlist_visited, context);
+                    n,
+                    x,
+                    /*k=*/0, // range search has no k
+                    *handler,
+                    cq,
+                    &ndis,
+                    &nlist_visited,
+                    context,
+                    params,
+                    *scanner);
         } else {
             FAISS_THROW_FMT("Range search implem %d not implemented", impl);
         }
@@ -762,35 +816,33 @@ void IndexIVFFastScan::range_search_dispatch_implem(
                 if (!cq_i.done()) {
                     cq_i.quantize_slice(quantizer, x, quantizer_params);
                 }
-                std::unique_ptr<SIMDResultHandlerToFloat> handler;
-                if (is_max) {
-                    handler.reset(new PartialRangeHandler<
-                                  CMax<uint16_t, int64_t>,
-                                  true>(pres, radius, 0, i0, i1, sel));
-                } else {
-                    handler.reset(new PartialRangeHandler<
-                                  CMin<uint16_t, int64_t>,
-                                  true>(pres, radius, 0, i0, i1, sel));
-                }
+                auto scanner = make_partial_range_scanner(
+                        is_max, pres, radius, 0, i0, i1, sel);
+                auto* handler = scanner->handler();
 
                 if (impl == 12 || impl == 13) {
                     search_implem_12(
                             i1 - i0,
                             x + i0 * d,
-                            *handler.get(),
+                            *handler,
                             cq_i,
                             &ndis,
                             &nlist_visited,
-                            context);
+                            context,
+                            nullptr,
+                            *scanner);
                 } else {
                     search_implem_10(
                             i1 - i0,
                             x + i0 * d,
-                            *handler.get(),
+                            /*k=*/0,
+                            *handler,
                             cq_i,
                             &ndis,
                             &nlist_visited,
-                            context);
+                            context,
+                            params,
+                            *scanner);
                 }
             }
             pres.finalize();
@@ -811,7 +863,7 @@ void IndexIVFFastScan::search_implem_1(
         idx_t* labels,
         const CoarseQuantized& cq,
         const FastScanDistancePostProcessing& context,
-        const IVFSearchParameters* params) const {
+        const IVFSearchParameters* /* params */) const {
     FAISS_THROW_IF_NOT(orig_invlists);
 
     size_t dim12 = ksub * M;
@@ -824,7 +876,7 @@ void IndexIVFFastScan::search_implem_1(
     bool single_LUT = !lookup_table_is_3d();
 
     size_t ndis = 0, nlist_visited = 0;
-    size_t nprobe = cq.nprobe;
+    size_t cur_nprobe = cq.nprobe;
 #pragma omp parallel for reduction(+ : ndis, nlist_visited)
     for (idx_t i = 0; i < n; i++) {
         int64_t* heap_ids = labels + i * k;
@@ -835,11 +887,11 @@ void IndexIVFFastScan::search_implem_1(
         if (single_LUT) {
             LUT = dis_tables.get() + i * dim12;
         }
-        for (idx_t j = 0; j < nprobe; j++) {
+        for (size_t j = 0; j < cur_nprobe; j++) {
             if (!single_LUT) {
-                LUT = dis_tables.get() + (i * nprobe + j) * dim12;
+                LUT = dis_tables.get() + (i * cur_nprobe + j) * dim12;
             }
-            idx_t list_no = cq.ids[i * nprobe + j];
+            idx_t list_no = cq.ids[i * cur_nprobe + j];
             if (list_no < 0) {
                 continue;
             }
@@ -850,7 +902,7 @@ void IndexIVFFastScan::search_implem_1(
             InvertedLists::ScopedCodes codes(orig_invlists, list_no);
             InvertedLists::ScopedIds ids(orig_invlists, list_no);
 
-            float bias = biases.get() ? biases[i * nprobe + j] : 0;
+            float bias = biases.get() ? biases[i * cur_nprobe + j] : 0;
 
             estimators_from_tables_generic<C>(
                     *this,
@@ -882,7 +934,7 @@ void IndexIVFFastScan::search_implem_2(
         idx_t* labels,
         const CoarseQuantized& cq,
         const FastScanDistancePostProcessing& context,
-        const IVFSearchParameters* params) const {
+        const IVFSearchParameters* /* params */) const {
     FAISS_THROW_IF_NOT(orig_invlists);
 
     size_t dim12 = ksub * M2;
@@ -895,7 +947,7 @@ void IndexIVFFastScan::search_implem_2(
     bool single_LUT = !lookup_table_is_3d();
 
     size_t ndis = 0, nlist_visited = 0;
-    size_t nprobe = cq.nprobe;
+    size_t cur_nprobe = cq.nprobe;
 
 #pragma omp parallel for reduction(+ : ndis, nlist_visited)
     for (idx_t i = 0; i < n; i++) {
@@ -908,11 +960,11 @@ void IndexIVFFastScan::search_implem_2(
         if (single_LUT) {
             LUT = dis_tables.get() + i * dim12;
         }
-        for (idx_t j = 0; j < nprobe; j++) {
+        for (size_t j = 0; j < cur_nprobe; j++) {
             if (!single_LUT) {
-                LUT = dis_tables.get() + (i * nprobe + j) * dim12;
+                LUT = dis_tables.get() + (i * cur_nprobe + j) * dim12;
             }
-            idx_t list_no = cq.ids[i * nprobe + j];
+            idx_t list_no = cq.ids[i * cur_nprobe + j];
             if (list_no < 0) {
                 continue;
             }
@@ -923,7 +975,7 @@ void IndexIVFFastScan::search_implem_2(
             InvertedLists::ScopedCodes codes(orig_invlists, list_no);
             InvertedLists::ScopedIds ids(orig_invlists, list_no);
 
-            uint16_t bias = biases.get() ? biases[i * nprobe + j] : 0;
+            uint16_t bias = biases.get() ? biases[i * cur_nprobe + j] : 0;
 
             estimators_from_tables_generic<C>(
                     *this,
@@ -962,12 +1014,14 @@ void IndexIVFFastScan::search_implem_2(
 void IndexIVFFastScan::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 */) const {
+        const IVFSearchParameters* params,
+        FastScanCodeScanner& scanner) const {
     size_t dim12 = ksub * M2;
     AlignedTable<uint8_t> dis_tables;
     AlignedTable<uint16_t> biases;
@@ -981,7 +1035,28 @@ void IndexIVFFastScan::search_implem_10(
     int qmap1[1];
     handler.q_map = qmap1;
     handler.begin(skip & 16 ? nullptr : normalizers.get());
-    size_t nprobe = cq.nprobe;
+    size_t cur_nprobe = cq.nprobe;
+
+    // Per-query early-stop options from SearchParametersIVF.
+    const size_t param_max_codes = params ? params->max_codes : 0;
+    const size_t param_max_lists_num = params ? params->max_lists_num : 0;
+    const bool ensure_topk_full = params ? params->ensure_topk_full : false;
+    const size_t cur_max_codes = (param_max_codes == 0)
+            ? std::numeric_limits<size_t>::max()
+            : param_max_codes;
+    const size_t cur_max_lists_num =
+            (param_max_lists_num == 0) ? cur_nprobe : param_max_lists_num;
+    // Effective budgets are the values tested in the probe loop below.
+    // ensure_topk_full raises small budgets to reduce empty result slots.
+    const size_t effective_max_codes = ensure_topk_full
+            ? std::max(cur_max_codes, (size_t)k)
+            : cur_max_codes;
+    const size_t effective_max_lists_num = ensure_topk_full
+            ? std::max(cur_max_lists_num, (size_t)k)
+            : cur_max_lists_num;
+    const bool is_range_search = k == 0;
+    const size_t max_empty_result_buckets =
+            (is_range_search && params) ? params->max_empty_result_buckets : 0;
 
     // Allocate probe_map once and reuse it
     std::vector<int> probe_map;
@@ -989,13 +1064,30 @@ void IndexIVFFastScan::search_implem_10(
 
     for (idx_t i = 0; i < n; i++) {
         const uint8_t* LUT = nullptr;
-        qmap1[0] = i;
+        qmap1[0] = static_cast<int>(i);
 
         if (single_LUT) {
             LUT = dis_tables.get() + i * dim12;
         }
-        for (idx_t j = 0; j < nprobe; j++) {
-            size_t ij = i * nprobe + j;
+        // Per-query counters. For k-NN, the handler count excludes rows
+        // filtered by IDSelector.
+        const size_t scan0 = handler.count_scanned_rows();
+        size_t nscan_q = 0;
+        size_t nlists_visited_q = 0;
+        size_t nempty_result_buckets = 0;
+        for (size_t j = 0; j < cur_nprobe; j++) {
+            if (!is_range_search) {
+                nscan_q = handler.count_scanned_rows() - scan0;
+            }
+            // Early-stop check: apply k-NN max_codes/max_lists_num before
+            // starting the next list. nscan_q excludes IDSelector-filtered
+            // rows.
+            if (nscan_q >= effective_max_codes ||
+                nlists_visited_q >= effective_max_lists_num) {
+                break;
+            }
+            const size_t prev_in_range_num = handler.in_range_num;
+            size_t ij = i * cur_nprobe + j;
             if (!single_LUT) {
                 LUT = dis_tables.get() + ij * dim12;
             }
@@ -1005,10 +1097,22 @@ void IndexIVFFastScan::search_implem_10(
 
             idx_t list_no = cq.ids[ij];
             if (list_no < 0) {
+                // Early-stop check: invalid probes count as empty range
+                // buckets.
+                if (max_empty_result_buckets > 0 &&
+                    ++nempty_result_buckets >= max_empty_result_buckets) {
+                    break;
+                }
                 continue;
             }
             size_t ls = invlists->list_size(list_no);
             if (ls == 0) {
+                // Early-stop check: empty inverted lists count as empty range
+                // buckets.
+                if (max_empty_result_buckets > 0 &&
+                    ++nempty_result_buckets >= max_empty_result_buckets) {
+                    break;
+                }
                 continue;
             }
 
@@ -1023,19 +1127,35 @@ void IndexIVFFastScan::search_implem_10(
             probe_map[0] = static_cast<int>(j);
             handler.set_list_context(list_no, probe_map);
 
-            pq4_accumulate_loop(
+            scanner.accumulate_loop(
                     1,
                     roundup(ls, bbs),
                     bbs,
                     M2,
                     codes.get(),
                     LUT,
-                    handler,
-                    context.norm_scaler,
+                    context.pq2x4_scale,
                     get_block_stride());
 
             ndis += ls;
             nlist_visited++;
+            if (is_range_search) {
+                nscan_q += ls;
+            }
+            nlists_visited_q++;
+
+            if (max_empty_result_buckets > 0) {
+                // Early-stop check: apply the range-search empty-bucket
+                // budget after each visited list; any hit resets the counter.
+                if (handler.in_range_num == prev_in_range_num) {
+                    nempty_result_buckets++;
+                    if (nempty_result_buckets >= max_empty_result_buckets) {
+                        break;
+                    }
+                } else {
+                    nempty_result_buckets = 0;
+                }
+            }
         }
     }
 
@@ -1052,7 +1172,8 @@ void IndexIVFFastScan::search_implem_12(
         size_t* ndis_out,
         size_t* nlist_out,
         const FastScanDistancePostProcessing& context,
-        const IVFSearchParameters* /* params */) const {
+        const IVFSearchParameters* /* params */,
+        FastScanCodeScanner& scanner) const {
     if (n == 0) { // does not work well with reservoir
         return;
     }
@@ -1073,15 +1194,15 @@ void IndexIVFFastScan::search_implem_12(
         int rank;    // this is the rank'th result of the coarse quantizer
     };
     bool single_LUT = !lookup_table_is_3d();
-    size_t nprobe = cq.nprobe;
+    size_t cur_nprobe = cq.nprobe;
 
     std::vector<QC> qcs;
     {
-        int ij = 0;
-        for (int i = 0; i < n; i++) {
-            for (int j = 0; j < nprobe; j++) {
+        size_t ij = 0;
+        for (idx_t i = 0; i < n; i++) {
+            for (size_t j = 0; j < cur_nprobe; j++) {
                 if (cq.ids[ij] >= 0) {
-                    qcs.push_back(QC{i, int(cq.ids[ij]), int(j)});
+                    qcs.push_back(QC{int(i), int(cq.ids[ij]), int(j)});
                 }
                 ij++;
             }
@@ -1093,7 +1214,7 @@ void IndexIVFFastScan::search_implem_12(
 
     // prepare the result handlers
 
-    int actual_qbs2 = this->qbs2 ? this->qbs2 : 11;
+    int actual_qbs2 = static_cast<int>(this->qbs2 ? this->qbs2 : 11);
 
     std::vector<uint16_t> tmp_bias;
     if (biases.get()) {
@@ -1130,7 +1251,7 @@ void IndexIVFFastScan::search_implem_12(
         nlist_visited++;
 
         // re-organize LUTs and biases into the right order
-        int nc = i1 - i0;
+        int nc = static_cast<int>(i1 - i0);
 
         std::vector<int> q_map(nc), lut_entries(nc);
         AlignedTable<uint8_t> LUT(nc * dim12);
@@ -1140,7 +1261,7 @@ void IndexIVFFastScan::search_implem_12(
         for (size_t i = i0; i < i1; i++) {
             const QC& qc = qcs[i];
             q_map[i - i0] = qc.qno;
-            int ij = qc.qno * nprobe + qc.rank;
+            int ij = static_cast<int>(qc.qno * cur_nprobe + qc.rank);
             lut_entries[i - i0] = single_LUT ? qc.qno : ij;
             if (biases.get()) {
                 tmp_bias[i - i0] = biases[ij];
@@ -1148,7 +1269,7 @@ void IndexIVFFastScan::search_implem_12(
         }
         pq4_pack_LUT_qbs_q_map(
                 qbs_for_list,
-                M2,
+                static_cast<int>(M2),
                 dis_tables.get(),
                 lut_entries.data(),
                 LUT.get());
@@ -1176,14 +1297,13 @@ void IndexIVFFastScan::search_implem_12(
         }
         handler.set_list_context(list_no, probe_map);
 
-        pq4_accumulate_loop_qbs(
+        scanner.accumulate_loop_qbs(
                 qbs_for_list,
                 list_size,
-                M2,
+                static_cast<int>(M2),
                 codes.get(),
                 LUT.get(),
-                handler,
-                context.norm_scaler,
+                context.pq2x4_scale,
                 get_block_stride());
         // prepare for next loop
         i0 = i1;
@@ -1229,15 +1349,15 @@ void IndexIVFFastScan::search_implem_14(
         int rank;    // this is the rank'th result of the coarse quantizer
     };
     bool single_LUT = !lookup_table_is_3d();
-    size_t nprobe = cq.nprobe;
+    size_t cur_nprobe = cq.nprobe;
 
     std::vector<QC> qcs;
     {
-        int ij = 0;
-        for (int i = 0; i < n; i++) {
-            for (int j = 0; j < nprobe; j++) {
+        size_t ij = 0;
+        for (idx_t i = 0; i < n; i++) {
+            for (size_t j = 0; j < cur_nprobe; j++) {
                 if (cq.ids[ij] >= 0) {
-                    qcs.push_back(QC{i, int(cq.ids[ij]), int(j)});
+                    qcs.push_back(QC{int(i), int(cq.ids[ij]), int(j)});
                 }
                 ij++;
             }
@@ -1316,25 +1436,24 @@ void IndexIVFFastScan::search_implem_14(
         std::vector<idx_t> local_idx(k * n);
         std::vector<float> local_dis(k * n);
 
-        // prepare the result handlers
-        std::unique_ptr<SIMDResultHandlerToFloat> handler(
-                this->make_knn_handler(
-                        is_max,
-                        impl,
-                        n,
-                        k,
-                        local_dis.data(),
-                        local_idx.data(),
-                        sel,
-                        context));
-        handler->begin(normalizers.get());
-
-        int actual_qbs2 = this->qbs2 ? this->qbs2 : 11;
+        auto scanner = make_knn_scanner(
+                is_max,
+                n,
+                k,
+                local_dis.data(),
+                local_idx.data(),
+                sel,
+                impl,
+                context);
+        SIMDResultHandlerToFloat* handler_ptr = scanner->handler();
+        handler_ptr->begin(normalizers.get());
+
+        int actual_qbs2 = static_cast<int>(this->qbs2 ? this->qbs2 : 11);
 
         std::vector<uint16_t> tmp_bias;
         if (biases.get()) {
             tmp_bias.resize(actual_qbs2);
-            handler->dbias = tmp_bias.data();
+            handler_ptr->dbias = tmp_bias.data();
         }
 
         std::set<int> q_set;
@@ -1345,7 +1464,8 @@ void IndexIVFFastScan::search_implem_14(
         probe_map.reserve(actual_qbs2);
 
 #pragma omp for schedule(dynamic)
-        for (idx_t cluster = 0; cluster < ses.size(); cluster++) {
+        for (idx_t cluster = 0; cluster < static_cast<idx_t>(ses.size());
+             cluster++) {
             size_t i0 = ses[cluster].start;
             size_t i1 = ses[cluster].end;
             size_t list_size = ses[cluster].list_size;
@@ -1353,7 +1473,7 @@ void IndexIVFFastScan::search_implem_14(
             int list_no = qcs[i0].list_no;
 
             // re-organize LUTs and biases into the right order
-            int nc = i1 - i0;
+            int nc = static_cast<int>(i1 - i0);
 
             std::vector<int> q_map(nc), lut_entries(nc);
             AlignedTable<uint8_t> LUT(nc * dim12);
@@ -1364,7 +1484,7 @@ void IndexIVFFastScan::search_implem_14(
                 const QC& qc = qcs[i];
                 q_map[i - i0] = qc.qno;
                 q_set.insert(qc.qno);
-                int ij = qc.qno * nprobe + qc.rank;
+                int ij = static_cast<int>(qc.qno * cur_nprobe + qc.rank);
                 lut_entries[i - i0] = single_LUT ? qc.qno : ij;
                 if (biases.get()) {
                     tmp_bias[i - i0] = biases[ij];
@@ -1372,7 +1492,7 @@ void IndexIVFFastScan::search_implem_14(
             }
             pq4_pack_LUT_qbs_q_map(
                     qbs_for_list,
-                    M2,
+                    static_cast<int>(M2),
                     dis_tables.get(),
                     lut_entries.data(),
                     LUT.get());
@@ -1386,9 +1506,9 @@ void IndexIVFFastScan::search_implem_14(
 
             // prepare the handler
 
-            handler->ntotal = list_size;
-            handler->q_map = q_map.data();
-            handler->id_map = ids.get();
+            handler_ptr->ntotal = list_size;
+            handler_ptr->q_map = q_map.data();
+            handler_ptr->id_map = ids.get();
 
             // Set context information for handlers that need additional data
             // All queries in this batch access the same list_no, but each
@@ -1398,21 +1518,20 @@ void IndexIVFFastScan::search_implem_14(
                 const QC& qc = qcs[i];
                 probe_map[i - i0] = qc.rank;
             }
-            handler->set_list_context(list_no, probe_map);
+            handler_ptr->set_list_context(list_no, probe_map);
 
-            pq4_accumulate_loop_qbs(
+            scanner->accumulate_loop_qbs(
                     qbs_for_list,
                     list_size,
-                    M2,
+                    static_cast<int>(M2),
                     codes.get(),
                     LUT.get(),
-                    *handler.get(),
-                    context.norm_scaler,
+                    context.pq2x4_scale,
                     get_block_stride());
         }
 
         // labels is in-place for HeapHC
-        handler->end();
+        handler_ptr->end();
 
         // merge per-thread results
 #pragma omp single
@@ -1466,7 +1585,7 @@ void IndexIVFFastScan::reconstruct_from_offset(
     for (size_t m = 0; m < M; m++) {
         uint8_t c =
                 pq4_get_packed_element(list_codes.get(), bbs, M2, offset, m);
-        bsw.write(c, nbits);
+        bsw.write(c, static_cast<int>(nbits));
     }
 
     sa_decode(1, code.data(), recons);
@@ -1477,7 +1596,7 @@ void IndexIVFFastScan::reconstruct_orig_invlists() {
     FAISS_THROW_IF_NOT(orig_invlists->list_size(0) == 0);
 
 #pragma omp parallel for if (nlist > 100)
-    for (idx_t list_no = 0; list_no < nlist; list_no++) {
+    for (idx_t list_no = 0; list_no < static_cast<idx_t>(nlist); list_no++) {
         InvertedLists::ScopedCodes codes(invlists, list_no);
         InvertedLists::ScopedIds ids(invlists, list_no);
         size_t list_size = invlists->list_size(list_no);
@@ -1489,7 +1608,7 @@ void IndexIVFFastScan::reconstruct_orig_invlists() {
             for (size_t m = 0; m < M; m++) {
                 uint8_t c =
                         pq4_get_packed_element(codes.get(), bbs, M2, offset, m);
-                bsw.write(c, nbits);
+                bsw.write(c, static_cast<int>(nbits));
             }
 
             // get id
@@ -1516,7 +1635,7 @@ void IndexIVFFastScan::sa_decode(idx_t n, const uint8_t* codes, float* x)
             fine_quantizer->decode(code + coarse_size, xi, 1);
             if (by_residual) {
                 quantizer->reconstruct(list_no, residual.data());
-                for (size_t j = 0; j < d; j++) {
+                for (int j = 0; j < d; j++) {
                     xi[j] += residual[j];
                 }
             }