Add explicit IVFPQ scanner distance modes (#5632)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5632

Add an explicit IVFPQ scanner mode that preserves precomputed-table behavior by default and can decode PQ codes to compute residual L2 distances on demand. Keep the existing scanner factory overload compatible.

Reviewed By: mdouze

Differential Revision: D119608525

fbshipit-source-id: b60b2c9771bdfd95d9bee2e81b72c9c3e7abb82c

faiss/IndexIVFPQ.cpp

@@ -526,9 +526,17 @@ InvertedListScanner* IndexIVFPQ::get_InvertedListScanner(
         bool store_pairs,
         const IDSelector* sel,
         const IVFSearchParameters*) const {
+    return get_InvertedListScanner(
+            store_pairs, sel, IndexIVFPQ::ScannerMode::Precomputed);
+}
+
+InvertedListScanner* IndexIVFPQ::get_InvertedListScanner(
+        bool store_pairs,
+        const IDSelector* sel,
+        IndexIVFPQ::ScannerMode mode) const {
     return with_simd_level([&]<SIMDLevel SL>() -> InvertedListScanner* {
         return pq_code_distance::make_IVFPQInvertedListScanner<SL>(
-                *this, store_pairs, sel);
+                *this, store_pairs, sel, mode);
     });
 }
 

faiss/IndexIVFPQ.h

@@ -32,6 +32,11 @@ FAISS_API extern size_t precomputed_table_max_bytes;
  * vector is encoded as a product quantizer code.
  */
 struct IndexIVFPQ : IndexIVF {
+    enum class ScannerMode {
+        Precomputed,
+        OnTheFly,
+    };
+
     ProductQuantizer pq; ///< produces the codes
 
     bool do_polysemous_training; ///< reorder PQ centroids after training?
@@ -144,6 +149,12 @@ struct IndexIVFPQ : IndexIVF {
             const IDSelector* sel,
             const IVFSearchParameters* params) const override;
 
+    /** Returns a scanner with a fixed PQ distance-computation strategy. */
+    InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs,
+            const IDSelector* sel,
+            ScannerMode mode) const;
+
     /// build precomputed table
     void precompute_table();
 

faiss/impl/pq_code_distance/IVFPQScanner_impl.h

@@ -216,41 +216,41 @@ struct IVFPQScannerT : QueryTables {
         }
     }
 
+    float on_the_fly_dis0 = 0;
+
+    void prepare_on_the_fly_distance() {
+        on_the_fly_dis0 = 0;
+        if (!by_residual) {
+            return;
+        }
+        if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
+            ivfpq.quantizer->reconstruct(key, residual_vec);
+            on_the_fly_dis0 = fvec_inner_product_dispatch(residual_vec, qi, d);
+        } else {
+            ivfpq.quantizer->compute_residual(qi, residual_vec, key);
+        }
+    }
+
+    float distance_to_code_on_the_fly(const uint8_t* code) const {
+        pq.decode(code, decoded_vec);
+        if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
+            return on_the_fly_dis0 +
+                    fvec_inner_product_dispatch(decoded_vec, qi, d);
+        }
+        const float* dvec = by_residual ? residual_vec : qi;
+        return fvec_L2sqr_dispatch(decoded_vec, dvec, d);
+    }
+
     /// nothing is precomputed: access residuals on-the-fly
     template <class SearchResultType>
     void scan_on_the_fly_dist(
             size_t ncode,
             const uint8_t* codes,
             SearchResultType& res) const {
-        const float* dvec;
-        float local_dis0 = 0;
-        if (by_residual) {
-            if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
-                ivfpq.quantizer->reconstruct(key, residual_vec);
-                local_dis0 = fvec_inner_product_dispatch(residual_vec, qi, d);
-            } else {
-                ivfpq.quantizer->compute_residual(qi, residual_vec, key);
-            }
-            dvec = residual_vec;
-        } else {
-            dvec = qi;
-            local_dis0 = 0;
-        }
-
         for (size_t j = 0; j < ncode; j++, codes += pq.code_size) {
-            if (res.skip_entry(j)) {
-                continue;
+            if (!res.skip_entry(j)) {
+                res.add(j, distance_to_code_on_the_fly(codes));
             }
-            pq.decode(codes, decoded_vec);
-
-            float dis;
-            if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
-                dis = local_dis0 +
-                        fvec_inner_product_dispatch(decoded_vec, qi, d);
-            } else {
-                dis = fvec_L2sqr_dispatch(decoded_vec, dvec, d);
-            }
-            res.add(j, dis);
         }
     }
 
@@ -438,7 +438,6 @@ struct IVFPQScannerT : QueryTables {
  *
  * precompute_mode is how much we precompute (2 = precompute distance tables,
  * 1 = precompute pointers to distances, 0 = compute distances one by one).
- * Currently only 2 is supported
  *
  * use_sel: store or ignore the IDSelector
  */
@@ -463,20 +462,29 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQCodeDist>,
     }
 
     void set_query(const float* query) override {
-        this->init_query(query);
+        if (precompute_mode == 0) {
+            this->qi = query;
+        } else {
+            this->init_query(query);
+        }
     }
 
     void set_list(idx_t list_no_in, float coarse_dis_in) override {
         this->list_no = list_no_in;
         this->init_list(list_no_in, coarse_dis_in, precompute_mode);
+        if (precompute_mode == 0) {
+            this->prepare_on_the_fly_distance();
+        }
     }
 
     float distance_to_code(const uint8_t* code) const override {
+        if (precompute_mode == 0) {
+            return this->distance_to_code_on_the_fly(code);
+        }
         FAISS_THROW_IF_NOT(precompute_mode == 2);
-        float dis = this->dis0 +
+        return this->dis0 +
                 PQCodeDist::distance_single_code(
-                            this->pq.M, this->pq.nbits, this->sim_table, code);
-        return dis;
+                        this->pq.M, this->pq.nbits, this->sim_table, code);
     }
 
     size_t scan_codes(
@@ -511,27 +519,40 @@ template <SIMDLevel SL>
 InvertedListScanner* make_IVFPQInvertedListScanner(
         const IndexIVFPQ& ivfpq,
         bool store_pairs,
-        const IDSelector* sel);
+        const IDSelector* sel,
+        IndexIVFPQ::ScannerMode mode);
+
+template <SIMDLevel SL>
+inline InvertedListScanner* make_IVFPQInvertedListScanner(
+        const IndexIVFPQ& ivfpq,
+        bool store_pairs,
+        const IDSelector* sel) {
+    return make_IVFPQInvertedListScanner<SL>(
+            ivfpq, store_pairs, sel, IndexIVFPQ::ScannerMode::Precomputed);
+}
 
 // NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
 template <>
 InvertedListScanner* make_IVFPQInvertedListScanner<THE_SIMD_LEVEL>(
         const IndexIVFPQ& ivfpq,
         bool store_pairs,
-        const IDSelector* sel) {
+        const IDSelector* sel,
+        IndexIVFPQ::ScannerMode mode) {
+    const int precompute_mode =
+            mode == IndexIVFPQ::ScannerMode::Precomputed ? 2 : 0;
     auto make = [&]<class PQCodeDist, bool use_sel>() -> InvertedListScanner* {
         if (ivfpq.metric_type == METRIC_INNER_PRODUCT) {
             return new IVFPQScanner<
                     METRIC_INNER_PRODUCT,
                     CMin<float, idx_t>,
                     PQCodeDist,
-                    use_sel>(ivfpq, store_pairs, 2, sel);
+                    use_sel>(ivfpq, store_pairs, precompute_mode, sel);
         } else if (ivfpq.metric_type == METRIC_L2) {
             return new IVFPQScanner<
                     METRIC_L2,
                     CMax<float, idx_t>,
                     PQCodeDist,
-                    use_sel>(ivfpq, store_pairs, 2, sel);
+                    use_sel>(ivfpq, store_pairs, precompute_mode, sel);
         } else {
             FAISS_THROW_MSG("unsupported metric type");
         }

tests/test_ivfpq_indexing.cpp

@@ -7,6 +7,7 @@
 
 #include <cstdio>
 #include <cstdlib>
+#include <memory>
 #include <random>
 
 #include <gtest/gtest.h>
@@ -14,6 +15,46 @@
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexIVFPQ.h>
 
+TEST(IVFPQ, onTheFlyDistanceToCodeMatchesPrecomputedTable) {
+    constexpr int d = 4;
+    constexpr size_t nt = 256;
+    faiss::IndexFlatL2 coarse_quantizer(d);
+    faiss::IndexIVFPQ index(&coarse_quantizer, d, 2, 2, 4);
+
+    std::mt19937 rng(1234);
+    std::uniform_real_distribution<float> distrib(-1.0f, 1.0f);
+    std::vector<float> training(nt * d);
+    for (float& value : training) {
+        value = distrib(rng);
+    }
+    index.train(nt, training.data());
+
+    const std::vector<float> query = {0.2f, -0.4f, 0.6f, -0.8f};
+    const std::vector<float> value = {-0.1f, 0.3f, 0.7f, -0.5f};
+    faiss::idx_t list_no;
+    float coarse_distance;
+    index.quantizer->search(1, query.data(), 1, &coarse_distance, &list_no);
+    ASSERT_GE(list_no, 0);
+
+    std::vector<uint8_t> code(index.code_size);
+    index.encode(list_no, value.data(), code.data());
+    std::unique_ptr<faiss::InvertedListScanner> table_scanner(
+            index.get_InvertedListScanner(false, nullptr, nullptr));
+    std::unique_ptr<faiss::InvertedListScanner> on_the_fly_scanner(
+            index.get_InvertedListScanner(
+                    false, nullptr, faiss::IndexIVFPQ::ScannerMode::OnTheFly));
+
+    for (auto* scanner : {table_scanner.get(), on_the_fly_scanner.get()}) {
+        scanner->set_query(query.data());
+        scanner->set_list(list_no, coarse_distance);
+    }
+
+    EXPECT_NEAR(
+            on_the_fly_scanner->distance_to_code(code.data()),
+            table_scanner->distance_to_code(code.data()),
+            1e-4);
+}
+
 TEST(IVFPQ, accuracy) {
     // dimension of the vectors to index
     int d = 64;