From 82a823513d8a7a7e497de11e7da0cba2aae7721d Mon Sep 17 00:00:00 2001 From: Saumitr Pathak Date: Sat, 11 Apr 2026 01:04:37 -0700 Subject: [PATCH] Add iterator callbacks for distance computation and heap admission (#5082) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5082 Add two virtual callback methods to FAISS's InvertedListsIterator that fire during the iterate_codes hot loop: - on_distance_computed(vid, distance): called after computing the distance for each scored vector, before the heap comparison. Enables callers to observe every distance computation (e.g. for building a reserve pool of all scored vectors sorted by distance). - on_heap_changed(new_id, evicted_id): called when a vector displaces the current worst entry in the top-K heap. Enables callers to react to heap admission events (e.g. maintaining an O(K) PK mapping instead of O(N)). Both methods have default no-op implementations, so existing code is unaffected. The callbacks are invoked in both the min-heap (L2) and max-heap (IP) branches of iterate_codes. Zero overhead for non-users: the callbacks are guarded by a boolean flag (has_search_callbacks_, default false) that iterate_codes checks before each virtual dispatch. A correctly-predicted branch on modern x86 costs ~0 cycles due to speculative execution, compared to ~3-5 cycles per virtual call through a vtable. Since the flag is always false for non-callback users, the branch predictor learns it immediately and never mispredicts. Alternative approaches considered and why this design was chosen: - CRTP / templated iterate_codes: would achieve true zero overhead via static dispatch and if-constexpr elimination, but requires changing FAISS's public API and template-instantiating iterate_codes — too invasive for upstream. - Marking derived classes `final`: only enables devirtualization when the compiler sees the concrete type, which it cannot through InvertedListsIterator base pointers in iterate_codes. - Raw function pointers with null checks: equivalent performance to the guard flag approach, but less clean API semantics for FAISS's object-oriented design. The guard flag achieves the best tradeoff: zero overhead for non-users, clean virtual override semantics for users, minimal code change, backward compatible. Reviewed By: mnorris11 Differential Revision: D100352032 fbshipit-source-id: 42cfd24b4fbca1af9ccf07a2bb29528f35173659 --- faiss/IndexIVF.cpp | 14 ++++ faiss/invlists/InvertedLists.h | 18 +++++ tests/test_ivf_index.cpp | 121 +++++++++++++++++++++++++++++++++ 3 files changed, 153 insertions(+) diff --git a/faiss/IndexIVF.cpp b/faiss/IndexIVF.cpp index eaac7eb20..3d1d251e8 100644 --- a/faiss/IndexIVF.cpp +++ b/faiss/IndexIVF.cpp @@ -1432,11 +1432,19 @@ size_t InvertedListScanner::iterate_codes( size_t nup = 0; list_size = 0; + const bool has_cb = it->has_search_callbacks_; + if (!keep_max) { for (; it->is_available(); it->next()) { auto id_and_codes = it->get_id_and_codes(); float dis = distance_to_code(id_and_codes.second); + if (has_cb) { + it->on_distance_computed(id_and_codes.first, dis); + } if (dis < simi[0]) { + if (has_cb) { + it->on_heap_changed(id_and_codes.first, idxi[0]); + } maxheap_replace_top(k, simi, idxi, dis, id_and_codes.first); nup++; } @@ -1446,7 +1454,13 @@ size_t InvertedListScanner::iterate_codes( for (; it->is_available(); it->next()) { auto id_and_codes = it->get_id_and_codes(); float dis = distance_to_code(id_and_codes.second); + if (has_cb) { + it->on_distance_computed(id_and_codes.first, dis); + } if (dis > simi[0]) { + if (has_cb) { + it->on_heap_changed(id_and_codes.first, idxi[0]); + } minheap_replace_top(k, simi, idxi, dis, id_and_codes.first); nup++; } diff --git a/faiss/invlists/InvertedLists.h b/faiss/invlists/InvertedLists.h index 743bcad62..7f1933143 100644 --- a/faiss/invlists/InvertedLists.h +++ b/faiss/invlists/InvertedLists.h @@ -28,6 +28,24 @@ struct InvertedListsIterator { virtual bool is_available() const = 0; virtual void next() = 0; virtual std::pair get_id_and_codes() = 0; + + /// When true, iterate_codes will invoke on_distance_computed() and + /// on_heap_changed() via virtual dispatch. When false (the default), + /// iterate_codes skips the callbacks entirely — the guard branch is + /// perfectly predicted and costs ~0 cycles, so non-callback users + /// pay no overhead. Derived classes that override the callbacks + /// should set this to true in their constructor. + bool has_search_callbacks_ = false; + + /// Called from iterate_codes after distance computation for the vector + /// returned by the most recent get_id_and_codes(). Default: no-op. + /// Only invoked when has_search_callbacks_ is true. + virtual void on_distance_computed(idx_t /* vid */, float /* distance */) {} + + /// Called from iterate_codes when a vector replaces the current worst + /// in the top-K heap. evicted_id is the displaced entry. Default: no-op. + /// Only invoked when has_search_callbacks_ is true. + virtual void on_heap_changed(idx_t /* new_id */, idx_t /* evicted_id */) {} }; /** Table of inverted lists diff --git a/tests/test_ivf_index.cpp b/tests/test_ivf_index.cpp index b085949cf..38c58d844 100644 --- a/tests/test_ivf_index.cpp +++ b/tests/test_ivf_index.cpp @@ -368,3 +368,124 @@ TEST(IVF, search1_out_of_range_key) { EXPECT_THROW(idx.search1(xq.data(), handler), faiss::FaissException); } + +// Iterator that enables search callbacks and tracks invocations. +class CallbackTrackingIterator : public TestInvertedListIterator { + public: + CallbackTrackingIterator( + size_t list_no, + TestContext* context, + size_t& distance_count, + size_t& heap_count) + : TestInvertedListIterator(list_no, context), + distance_count_{distance_count}, + heap_count_{heap_count} { + has_search_callbacks_ = true; + } + + void on_distance_computed(faiss::idx_t id, float distance) override { + EXPECT_GE(id, 0) << "vector ID should be non-negative"; + EXPECT_GE(distance, 0.0f) << "L2 distance should be non-negative"; + distance_count_++; + } + + void on_heap_changed(faiss::idx_t new_id, faiss::idx_t evicted_id) + override { + EXPECT_GE(new_id, 0) << "new heap entry ID should be non-negative"; + (void)evicted_id; // may be -1 when heap not yet full + heap_count_++; + } + + private: + size_t& distance_count_; + size_t& heap_count_; +}; + +// InvertedLists that uses CallbackTrackingIterator. +class CallbackTrackingInvertedLists : public TestInvertedLists { + public: + CallbackTrackingInvertedLists( + size_t nlist_in, + size_t code_size_in, + size_t& distance_count, + size_t& heap_count) + : TestInvertedLists(nlist_in, code_size_in), + distance_count_{distance_count}, + heap_count_{heap_count} {} + + faiss::InvertedListsIterator* get_iterator(size_t list_no, void* context) + const override { + auto testContext = (TestContext*)context; + testContext->lists_probed.insert(list_no); + return new CallbackTrackingIterator( + list_no, testContext, distance_count_, heap_count_); + } + + private: + size_t& distance_count_; + size_t& heap_count_; +}; + +// Test: on_distance_computed and on_heap_changed fire during search +// when has_search_callbacks_ is true. +TEST(IVF, search_callbacks) { + constexpr int d = 8; + constexpr int nb = 200; + constexpr int nlist = 4; + + std::mt19937 rng(42); + std::uniform_real_distribution<> distrib; + + omp_set_num_threads(1); + + faiss::IndexFlatL2 quantizer(d); + faiss::IndexIVFFlat index(&quantizer, d, nlist); + + size_t distance_count = 0; + size_t heap_count = 0; + CallbackTrackingInvertedLists invlists( + nlist, index.code_size, distance_count, heap_count); + index.replace_invlists(&invlists); + + // Train + constexpr size_t nt = 100; + std::vector trainvecs(nt * d); + for (size_t i = 0; i < nt * d; i++) { + trainvecs[i] = distrib(rng); + } + index.train(nt, trainvecs.data()); + + // Populate via context + TestContext context; + std::vector database(nb * d); + for (size_t i = 0; i < nb * d; i++) { + database[i] = distrib(rng); + } + std::vector coarse_idx(nb); + index.quantizer->assign(nb, database.data(), coarse_idx.data()); + std::vector xids(nb, 42); + index.add_core( + nb, database.data(), xids.data(), coarse_idx.data(), &context); + + // Search + constexpr faiss::idx_t k = 5; + constexpr size_t nprobe = 2; + std::vector query(d); + for (int i = 0; i < d; i++) { + query[i] = distrib(rng); + } + std::vector distances(k); + std::vector labels(k); + faiss::SearchParametersIVF params; + params.inverted_list_context = &context; + params.nprobe = nprobe; + + index.search(1, query.data(), k, distances.data(), labels.data(), ¶ms); + + EXPECT_GT(distance_count, 0) + << "on_distance_computed should fire for scored vectors"; + EXPECT_GT(heap_count, 0) + << "on_heap_changed should fire when vectors enter the heap"; + EXPECT_GE(distance_count, heap_count) + << "not every distance computation leads to a heap change"; +}