From f58fd4c8ca0fdf6c230c062ed70b93586ae9ec28 Mon Sep 17 00:00:00 2001 From: Michael Norris Date: Wed, 15 Oct 2025 20:39:45 -0700 Subject: [PATCH] Add InvertedListScanner for IVFPQFastScan (#4537) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4537 Description - Adds InvertedListScanner for IVFPQFastScan. Right now IVFPQFastScan does not support a Scanner because the codes are PQFS packed. This is required from Unicorn to simplify the integration layer. We can implement a scanner based on the search_implem_10. We can't do distance_to_code but we can do the whole scan. (that is what is done in this diff in the new `IVFPQFastScanScanner.scan_codes`). Summary of changes (v10) - - High level: there's 2 heaps at play during a `scan_codes` call. One is the uint16_t heap inside the HeapHandler, which tracks the best results for a given iteration of nprobe. The other tracks the results across iterations and is in float: heap_distances. At the end of each call, we find the best results from the current iteration (`curr_dists`) then `heap_addn` those to the output distances `distances`). - Moved `make_knn_handler` to the IndexIVFFastScan.h declaration so it can be used outside just IndexIVFFastScan.cpp. It is moved out of the anonymous namespace in the .cpp file too. - Add the IVFPQFastScanScanner to IndexIVFPQFastScan. This impl for scan_codes is directly translated from search_implem_10 in in IndexIVFPQFastScan, because both search_implem_10 and scan_codes process 1 query and 1 invlist at a time (technically, search_implem_10 processes `n` queries at a time, just not batched. search_implem_10 has a simple loop over `n`, so search_implem_10 can be factored out to 1 query for this Scanner easily.). The only changes are removal of extraneous stuff from search_implem_10. - Changes to test_lowlevel_ivf.cpp to test this. Refactored the test because IVFPQFastScanScanner does not behave like other scanners. v15 vs v16 - v15 matches search_implem_10 exactly. The unit tests pass with bit-exact match. They were not passing before due to the biases and LUT not matching exactly to search_implem_10. Problems: 1. It requires quantizer search. Otherwise the biases and LUT will not match search_implem_10. The LUT is built from several inverted lists. 2. It requires mutable state across scan_codes calls. I called them lut_counter and bias_counter. (I think even though this matches results, these problems make this option a no-go. It is mostly here to show the investigation.) v16 does not use quantizer search and has no mutable state. - cq.dis == distance from query (xi) to centroid represented by list_no. - cq.ids == list_no Both are of length 1 since we just have 1 list which is set explicitly. Then we can call compute_LUT as normal. The LUTs and biases are completely different from an existing approach, so results will not match an existing approach. Match percentage between search_implem_10 and various by_residual approaches (I called it `recall` in the test): by_residual_IP: 0.9745 by_residual_L2: 0.936 Thus I set a baseline of 0.9 in the test. Reference impl: https://github.com/facebookresearch/faiss/blob/514b44fca8542bafe8640adcbf1cccce1900f74c/faiss/IndexIVFFastScan.cpp#L918-L992 ntotal vs number of blocks? - - index.bbs is often 32 https://github.com/facebookresearch/faiss/wiki/Fast-accumulation-of-PQ-and-AQ-codes-(FastScan)#implementation. This is the number of codes per block. ~~Why do we need this io_simd_dis and io_simd_ids?~~ - Old, but keeping it struck out. This describes the uint16_t heap kept over multiple scans. ~~Other scanners can scan over multiple lists with `set_query` then `set_list`, then calling `scan_codes`. Why? Because they take the `float* distances`. The heap can operate on these floats to re-use values from last iteration.~~ ~~IVFPQFastScan is different. Codes are pq4 packed. The HeapHandler has to get initialized with values from scanning the prior list.~~ ~~If the user is only setting 1 list and 1 query for the lifetime of this scanner (which I believe is Unicorn's use case), it does not matter because there's no need to use prior results and the user can follow the test named IVFPQFS_null_params. But if the user wants to use like the other Scanners, they can follow the other 2 tests in test_lowlevel_ivf.cpp that pass the heap size.~~ Reviewed By: ddrcoder, mdouze Differential Revision: D80114737 fbshipit-source-id: a28392d444bc4dbe5d12e0fd84884e691cf00cf0 --- faiss/IndexIVF.h | 6 +- faiss/IndexIVFFastScan.cpp | 42 +++++----- faiss/IndexIVFFastScan.h | 11 ++- faiss/IndexIVFPQFastScan.cpp | 130 ++++++++++++++++++++++++++++- faiss/IndexIVFPQFastScan.h | 5 ++ faiss/impl/simd_result_handlers.h | 44 ++++++++-- tests/test_lowlevel_ivf.cpp | 133 ++++++++++++++++++++++++++---- 7 files changed, 324 insertions(+), 47 deletions(-) diff --git a/faiss/IndexIVF.h b/faiss/IndexIVF.h index 304b9d1bd..05f87e07e 100644 --- a/faiss/IndexIVF.h +++ b/faiss/IndexIVF.h @@ -497,12 +497,12 @@ struct InvertedListScanner { /// compute a single query-to-code distance virtual float distance_to_code(const uint8_t* code) const = 0; - /** scan a set of codes, compute distances to current query and + /** scan a set of codes, compute distances to current query, and * update heap of results if necessary. Default implementation * calls distance_to_code. * - * @param n number of codes to scan - * @param codes codes to scan (n * code_size) + * @param n number of codes to scan + * @param codes codes to scan (n * code_size) * @param ids corresponding ids (ignored if store_pairs) * @param distances heap distances (size k) * @param labels heap labels (size k) diff --git a/faiss/IndexIVFFastScan.cpp b/faiss/IndexIVFFastScan.cpp index ef93ef031..b3d28c96e 100644 --- a/faiss/IndexIVFFastScan.cpp +++ b/faiss/IndexIVFFastScan.cpp @@ -379,7 +379,8 @@ ResultHandlerCompare* make_knn_handler_fixC( idx_t k, float* distances, idx_t* labels, - const IDSelector* sel) { + const IDSelector* sel, + const float* normalizers) { using HeapHC = HeapHandler; using ReservoirHC = ReservoirHandler; using SingleResultHC = SingleResultHandler; @@ -387,29 +388,12 @@ ResultHandlerCompare* make_knn_handler_fixC( 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); + 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); } } -SIMDResultHandlerToFloat* make_knn_handler( - bool is_max, - int impl, - idx_t n, - idx_t k, - float* distances, - idx_t* labels, - const IDSelector* sel) { - if (is_max) { - return make_knn_handler_fixC>( - impl, n, k, distances, labels, sel); - } else { - return make_knn_handler_fixC>( - impl, n, k, distances, labels, sel); - } -} - using CoarseQuantized = IndexIVFFastScan::CoarseQuantized; struct CoarseQuantizedWithBuffer : CoarseQuantized { @@ -443,7 +427,7 @@ struct CoarseQuantizedWithBuffer : CoarseQuantized { }; struct CoarseQuantizedSlice : CoarseQuantizedWithBuffer { - size_t i0, i1; + const size_t i0, i1; CoarseQuantizedSlice(const CoarseQuantized& cq, size_t i0, size_t i1) : CoarseQuantizedWithBuffer(cq), i0(i0), i1(i1) { if (done()) { @@ -486,6 +470,24 @@ int compute_search_nslice( } // namespace +SIMDResultHandlerToFloat* IndexIVFFastScan::make_knn_handler( + bool is_max, + int impl, + idx_t n, + idx_t k, + float* distances, + idx_t* labels, + const IDSelector* sel, + const float* normalizers) const { + if (is_max) { + return make_knn_handler_fixC>( + impl, n, k, distances, labels, sel, normalizers); + } else { + return make_knn_handler_fixC>( + impl, n, k, distances, labels, sel, normalizers); + } +} + void IndexIVFFastScan::search_dispatch_implem( idx_t n, const float* x, diff --git a/faiss/IndexIVFFastScan.h b/faiss/IndexIVFFastScan.h index 770ec5b9e..af5ff6378 100644 --- a/faiss/IndexIVFFastScan.h +++ b/faiss/IndexIVFFastScan.h @@ -99,7 +99,7 @@ struct IndexIVFFastScan : IndexIVF { // compact way of conveying coarse quantization results struct CoarseQuantized { - size_t nprobe; + size_t nprobe = 0; const float* dis = nullptr; const idx_t* ids = nullptr; }; @@ -147,6 +147,15 @@ struct IndexIVFFastScan : IndexIVF { const SearchParameters* params = nullptr) const override; // internal search funcs + SIMDResultHandlerToFloat* make_knn_handler( + bool is_max, + int impl, + idx_t n, + idx_t k, + float* distances, + idx_t* labels, + const IDSelector* sel, + const float* normalizers = nullptr) const; // dispatch to implementations and parallelize void search_dispatch_implem( diff --git a/faiss/IndexIVFPQFastScan.cpp b/faiss/IndexIVFPQFastScan.cpp index 0e876bdb3..0ddb34b7d 100644 --- a/faiss/IndexIVFPQFastScan.cpp +++ b/faiss/IndexIVFPQFastScan.cpp @@ -7,6 +7,7 @@ #include +#include #include #include @@ -14,6 +15,7 @@ #include #include +#include #include #include @@ -213,7 +215,7 @@ void IndexIVFPQFastScan::compute_LUT( AlignedTable& biases) const { size_t dim12 = pq.ksub * pq.M; size_t d = pq.d; - size_t nprobe = this->nprobe; + size_t nprobe = cq.nprobe; if (by_residual) { if (metric_type == METRIC_L2) { @@ -292,4 +294,130 @@ void IndexIVFPQFastScan::compute_LUT( } } +/********************************************************* + * InvertedListScanner for IVFPQFS + *********************************************************/ + +namespace { + +struct IVFPQFastScanScanner : InvertedListScanner { + static constexpr int impl = 10; // based on search_implem_10 + static constexpr size_t nq = 1; // 1 query at a time. + const IndexIVFPQFastScan& index; + AlignedTable dis_tables; + AlignedTable biases; + std::array normalizers{}; + const float* xi = nullptr; + + IVFPQFastScanScanner( + const IndexIVFPQFastScan& index, + bool store_pairs, + const IDSelector* sel) + : InvertedListScanner(store_pairs, sel), index(index) { + this->keep_max = is_similarity_metric(index.metric_type); + } + + 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; + IndexIVFFastScan::CoarseQuantized cq{ + .nprobe = 1, // 1 due to explicitly passing in list_no + .dis = &coarse_dis, // dis from query to list_no centroid. + .ids = &list_no, // id of the current list we are scanning + }; + index.compute_LUT_uint8(1, xi, cq, dis_tables, biases, &normalizers[0]); + } + + float distance_to_code(const uint8_t* /* code */) const override { + // It's not really possible to implement a distance_to_code since codes + // for 32 database vectors are intermixed. + FAISS_THROW_MSG("not implemented"); + } + + // Based on IVFFastScan search_implem_10, since it also deals with 1 query + // at a time. + size_t scan_codes( + size_t ntotal, + const uint8_t* codes, + const idx_t* ids, + float* distances, + idx_t* labels, + size_t k) const override { + // initialize the current iteration heap to the worst possible value of + // the prior loop + std::vector curr_dists(k, distances[0]); + std::vector curr_labels(k, labels[0]); + + std::unique_ptr handler( + index.make_knn_handler( + !keep_max, + impl, + nq, + k, + curr_dists.data(), + curr_labels.data(), + sel, + &normalizers[0])); + + // This does not quite match search_implem_10, but it is fine because + // the scanner operates on a single query at a time, and this value is + // used as the query index. For a single query, the value is always 0. + int qmap1[1] = {0}; + + handler->q_map = qmap1; + handler->begin(&normalizers[0]); + + const uint8_t* LUT = dis_tables.get(); + handler->dbias = biases.get(); + + handler->ntotal = ntotal; + handler->id_map = ids; + + pq4_accumulate_loop( + 1, + roundup(ntotal, index.bbs), + index.bbs, + static_cast(index.M2), + codes, + LUT, + *handler, + nullptr); + + // The handler is for the results of this iteration. + // Then we need a second heap to combine across iterations. + handler->end(); + if (keep_max) { + minheap_addn( + k, + distances, + labels, + curr_dists.data(), + curr_labels.data(), + k); + } else { + maxheap_addn( + k, + distances, + labels, + curr_dists.data(), + curr_labels.data(), + k); + } + + return handler->num_updates(); + } +}; + +} // anonymous namespace + +InvertedListScanner* IndexIVFPQFastScan::get_InvertedListScanner( + bool store_pairs, + const IDSelector* sel, + const IVFSearchParameters*) const { + return new IVFPQFastScanScanner(*this, store_pairs, sel); +} + } // namespace faiss diff --git a/faiss/IndexIVFPQFastScan.h b/faiss/IndexIVFPQFastScan.h index 531976dd4..1ae7145ea 100644 --- a/faiss/IndexIVFPQFastScan.h +++ b/faiss/IndexIVFPQFastScan.h @@ -81,6 +81,11 @@ struct IndexIVFPQFastScan : IndexIVFFastScan { const CoarseQuantized& cq, AlignedTable& dis_tables, AlignedTable& biases) const override; + + InvertedListScanner* get_InvertedListScanner( + bool store_pairs, + const IDSelector* sel, + const IVFSearchParameters*) const override; }; } // namespace faiss diff --git a/faiss/impl/simd_result_handlers.h b/faiss/impl/simd_result_handlers.h index baa640d86..34476a798 100644 --- a/faiss/impl/simd_result_handlers.h +++ b/faiss/impl/simd_result_handlers.h @@ -46,7 +46,7 @@ struct SIMDResultHandler { virtual ~SIMDResultHandler() {} }; -/* Result handler that will return float resutls eventually */ +/* Result handler that will return float results eventually */ struct SIMDResultHandlerToFloat : SIMDResultHandler { size_t nq; // number of queries size_t ntotal; // ignore excess elements after ntotal @@ -70,6 +70,12 @@ struct SIMDResultHandlerToFloat : SIMDResultHandler { virtual void end() { normalizers = nullptr; } + + // Number of updates made to the underlying data structure. + // For example: number of heap updates. + virtual size_t num_updates() { + return 0; + } }; FAISS_API extern bool simd_result_handlers_accept_virtual; @@ -318,8 +324,8 @@ struct HeapHandler : ResultHandlerCompare { std::vector iids; float* dis; int64_t* ids; - - int64_t k; // number of results to keep + size_t k; // number of results to keep + size_t nup = 0; // number of heap updates HeapHandler( size_t nq, @@ -327,14 +333,30 @@ struct HeapHandler : ResultHandlerCompare { int64_t k, float* dis, int64_t* ids, - const IDSelector* sel_in) + const IDSelector* sel_in, + const float* normalizers = nullptr) : RHC(nq, ntotal, sel_in), - idis(nq * k), - iids(nq * k), + idis(nq * k, threshold_idis(dis, normalizers)), + iids(nq * k, -1), dis(dis), ids(ids), - k(k) { - heap_heapify(k * nq, idis.data(), iids.data()); + k(k) {} + + static uint16_t threshold_idis(float* dis_in, const float* normalizers) { + if (dis_in[0] == std::numeric_limits::max()) { + return std::numeric_limits::max(); + } + if (dis_in[0] == std::numeric_limits::lowest()) { + return 0; + } + if (normalizers) { + float one_a = 1 / normalizers[0], b = normalizers[1]; + float f = (dis_in[0] - b) / one_a; + f = C::is_max ? std::ceil(f) : std::floor(f); + return std::clamp( + f, 0, std::numeric_limits::max()); + } + return C::neutral(); } void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final { @@ -372,6 +394,7 @@ struct HeapHandler : ResultHandlerCompare { if (C::cmp(heap_dis[0], dis_2)) { heap_replace_top( k, heap_dis, heap_ids, dis_2, real_idx); + nup++; } } } @@ -384,6 +407,7 @@ struct HeapHandler : ResultHandlerCompare { if (C::cmp(heap_dis[0], dis_2)) { int64_t idx = this->adjust_id(b, j); heap_replace_top(k, heap_dis, heap_ids, dis_2, idx); + nup++; } } } @@ -408,6 +432,10 @@ struct HeapHandler : ResultHandlerCompare { } } } + + size_t num_updates() override { + return nup; + } }; /** Simple top-N implementation using a reservoir. diff --git a/tests/test_lowlevel_ivf.cpp b/tests/test_lowlevel_ivf.cpp index 4422c8a8b..c0c99b0ad 100644 --- a/tests/test_lowlevel_ivf.cpp +++ b/tests/test_lowlevel_ivf.cpp @@ -5,7 +5,6 @@ * LICENSE file in the root directory of this source tree. */ -#include #include #include @@ -20,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -64,17 +64,31 @@ std::unique_ptr make_trained_index( return index; } -std::vector search_index(Index* index, const float* xq) { +std::pair, std::vector> search_index( + Index* index, + const float* xq) { std::vector I(k * nq); std::vector D(k * nq); index->search(nq, xq, k, D.data(), I.data()); - return I; + return {I, D}; } /************************************************************* * Test functions for a given index type *************************************************************/ +void test_get_InvertedListScanner( + IndexIVF* index_ivf, + const IndexPreTransform* index_pt, + std::vector codes, + std::unique_ptr scanner, + std::vector xq, + std::vector ref_I, + std::vector ref_D, + MetricType metric, + bool distance_to_code_supported = true, + float accuracy_requirement = 1.0); + void test_lowlevel_access(const char* index_key, MetricType metric) { std::unique_ptr index = make_trained_index(index_key, metric); @@ -152,8 +166,30 @@ void test_lowlevel_access(const char* index_key, MetricType metric) { // sample some example queries and get reference search results. auto xq = make_data(nq); - auto ref_I = search_index(index.get(), xq.data()); + auto res = search_index(index.get(), xq.data()); + test_get_InvertedListScanner( + index_ivf, + index_pt, + codes, + nullptr, + xq, + res.first, + res.second, + metric); +} + +void test_get_InvertedListScanner( + IndexIVF* index_ivf, + const IndexPreTransform* index_pt, + std::vector codes, + std::unique_ptr scanner, + std::vector xq, + std::vector ref_I, + std::vector ref_D, + MetricType metric, + bool distance_to_code_supported, + float accuracy_requirement) { // handle preprocessing const float* xqt = xq.data(); std::unique_ptr del_xqt; @@ -167,22 +203,27 @@ void test_lowlevel_access(const char* index_key, MetricType metric) { // quantize the queries to get the inverted list ids to visit. int nprobe = index_ivf->nprobe; + const InvertedLists* il = index_ivf->invlists; std::vector q_lists(nq * nprobe); std::vector q_dis(nq * nprobe); index_ivf->quantizer->search(nq, xqt, nprobe, q_dis.data(), q_lists.data()); - // object that does the scanning and distance computations. - std::unique_ptr scanner( - index_ivf->get_InvertedListScanner()); - + if (scanner == nullptr) { + // standard flow + scanner = std::unique_ptr( + index_ivf->get_InvertedListScanner()); + } + float recall = 0.0; for (int i = 0; i < nq; i++) { std::vector I(k, -1); - float default_dis = metric == METRIC_L2 ? HUGE_VAL : -HUGE_VAL; + float default_dis = metric == METRIC_L2 + ? std::numeric_limits::max() + : std::numeric_limits::lowest(); std::vector D(k, default_dis); - scanner->set_query(xqt + i * dt); + scanner->set_query(xqt + i * index_ivf->d); for (int j = 0; j < nprobe; j++) { int list_no = q_lists[i * nprobe + j]; @@ -202,7 +243,7 @@ void test_lowlevel_access(const char* index_key, MetricType metric) { I.data(), k); - if (j == 0) { + if (distance_to_code_supported && j == 0) { // all results so far come from list_no, so let's check if // the distance function works for (int jj = 0; jj < k; jj++) { @@ -229,9 +270,57 @@ void test_lowlevel_access(const char* index_key, MetricType metric) { // check that we have the same results as the reference search for (int j = 0; j < k; j++) { - EXPECT_EQ(I[j], ref_I[i * k + j]); + if (I[j] != ref_I[i * k + j]) { + if (accuracy_requirement == 1.0) { + EXPECT_EQ(D[j], ref_D[i * k + j]); + } + if (D[j] == ref_D[i * k + j]) { + recall += 1.0; + } + } else { + recall += 1.0; + } } } + recall = recall / (k * nq); + EXPECT_GE(recall, accuracy_requirement); +} + +void test_ivfpqfs_scanner(MetricType metric, bool by_residual = false) { + auto index = + std::unique_ptr(index_factory(d, "IVF32,PQ4x4fs", metric)); + ParameterSpace().set_index_parameter(index.get(), "nprobe", 4); + IndexIVFPQFastScan* index_ivf = static_cast( + ivflib::extract_index_ivf(index.get())); + index_ivf->by_residual = by_residual; + // implem_10 also processes one query at a time, so compare with that. + index_ivf->implem = 10; + auto xt = make_data(nt); + index_ivf->train(nt, xt.data()); + auto xb = make_data(nb); + index_ivf->add(nb, xb.data()); + // Initialize scanner with context in params so heap results can persist + // across multiple inverted list scans. + auto scanner = std::unique_ptr( + index_ivf->get_InvertedListScanner(true, nullptr, nullptr)); + + // ref data + auto xq = make_data(nq); + auto res = search_index(index_ivf, xq.data()); + + // distance_to_code_supported = false because codes are intermixed. + test_get_InvertedListScanner( + index_ivf, + nullptr, + {}, + std::move(scanner), + xq, + res.first, + res.second, + metric, + false, + by_residual ? 0.9 : 1.0 // recall is not 100% for by_residual + ); } } // anonymous namespace @@ -276,6 +365,22 @@ TEST(TestLowLevelIVF, IVFRQ) { test_lowlevel_access("IVF32,RQ16x8", METRIC_L2); } +TEST(TestLowLevelIVF, IVFPQFS_L2) { + test_ivfpqfs_scanner(METRIC_L2); +} + +TEST(TestLowLevelIVF, IVFPQFS_IP) { + test_ivfpqfs_scanner(METRIC_INNER_PRODUCT); +} + +TEST(TestLowLevelIVF, IVFPQFSr_L2) { + test_ivfpqfs_scanner(METRIC_L2, true); +} + +TEST(TestLowLevelIVF, IVFPQFSr_IP) { + test_ivfpqfs_scanner(METRIC_INNER_PRODUCT, true); +} + /************************************************************* * Same for binary (a bit simpler) *************************************************************/ @@ -459,7 +564,7 @@ void test_threaded_search(const char* index_key, MetricType metric) { // sample some example queries and get reference search results. auto xq = make_data(nq); - auto ref_I = search_index(index.get(), xq.data()); + auto res = search_index(index.get(), xq.data()); // handle preprocessing const float* xqt = xq.data(); @@ -570,7 +675,7 @@ void test_threaded_search(const char* index_key, MetricType metric) { // check that we have the same results as the reference search for (int j = 0; j < k; j++) { - EXPECT_EQ(I[j], ref_I[i * k + j]); + EXPECT_EQ(I[j], res.first[i * k + j]); } } }