Enable precomp table to be shared ivfpq

Changes IVFPQ and IVFPQFastScan indices to be able to share the precomputed table amongst other instances. Switches var to a pointer and add necessary functions to set them correctly. Adds a tests to validate the behavior. Signed-off-by: John Mazanec <[email protected]>
junqiu-lei · Apr 25, 2024 · f2dd667 · f2dd667
1 parent f45bad6
commit f2dd667
Show file tree

Hide file tree

Showing 7 changed files with 284 additions and 15 deletions.
diff --git a/faiss/IndexIVFPQ.cpp b/faiss/IndexIVFPQ.cpp
@@ -59,6 +59,29 @@ IndexIVFPQ::IndexIVFPQ(
     polysemous_training = nullptr;
     do_polysemous_training = false;
     polysemous_ht = 0;
+    precomputed_table = new AlignedTable<float>();
+    owns_precomputed_table = true;
+}
+
+IndexIVFPQ::IndexIVFPQ(const IndexIVFPQ& orig) : IndexIVF(orig), pq(orig.pq) {
+    code_size = orig.pq.code_size;
+    invlists->code_size = code_size;
+    is_trained = orig.is_trained;
+    by_residual = orig.by_residual;
+    use_precomputed_table = orig.use_precomputed_table;
+    scan_table_threshold = orig.scan_table_threshold;
+
+    polysemous_training = orig.polysemous_training;
+    do_polysemous_training = orig.do_polysemous_training;
+    polysemous_ht = orig.polysemous_ht;
+    precomputed_table = new AlignedTable<float>(*orig.precomputed_table);
+    owns_precomputed_table = true;
+}
+
+IndexIVFPQ::~IndexIVFPQ() {
+    if (owns_precomputed_table) {
+        delete precomputed_table;
+    }
 }
 
 /****************************************************************
@@ -466,11 +489,23 @@ void IndexIVFPQ::precompute_table() {
             use_precomputed_table,
             quantizer,
             pq,
-            precomputed_table,
+            *precomputed_table,
             by_residual,
             verbose);
 }
 
+void IndexIVFPQ::set_precomputed_table(
+        AlignedTable<float>* _precompute_table,
+        int _use_precomputed_table) {
+    // Clean up old pre-computed table
+    if (owns_precomputed_table) {
+        delete precomputed_table;
+    }
+    owns_precomputed_table = false;
+    precomputed_table = _precompute_table;
+    use_precomputed_table = _use_precomputed_table;
+}
+
 namespace {
 
 #define TIC t0 = get_cycles()
@@ -650,7 +685,7 @@ struct QueryTables {
 
             fvec_madd(
                     pq.M * pq.ksub,
-                    ivfpq.precomputed_table.data() + key * pq.ksub * pq.M,
+                    ivfpq.precomputed_table->data() + key * pq.ksub * pq.M,
                     -2.0,
                     sim_table_2,
                     sim_table);
@@ -679,7 +714,7 @@ struct QueryTables {
                 k >>= cpq.nbits;
 
                 // get corresponding table
-                const float* pc = ivfpq.precomputed_table.data() +
+                const float* pc = ivfpq.precomputed_table->data() +
                         (ki * pq.M + cm * Mf) * pq.ksub;
 
                 if (polysemous_ht == 0) {
@@ -709,7 +744,7 @@ struct QueryTables {
             dis0 = coarse_dis;
 
             const float* s =
-                    ivfpq.precomputed_table.data() + key * pq.ksub * pq.M;
+                    ivfpq.precomputed_table->data() + key * pq.ksub * pq.M;
             for (int m = 0; m < pq.M; m++) {
                 sim_table_ptrs[m] = s;
                 s += pq.ksub;
@@ -729,7 +764,7 @@ struct QueryTables {
                 int ki = k & ((uint64_t(1) << cpq.nbits) - 1);
                 k >>= cpq.nbits;
 
-                const float* pc = ivfpq.precomputed_table.data() +
+                const float* pc = ivfpq.precomputed_table->data() +
                         (ki * pq.M + cm * Mf) * pq.ksub;
 
                 for (int m = m0; m < m0 + Mf; m++) {
@@ -1346,6 +1381,8 @@ IndexIVFPQ::IndexIVFPQ() {
     do_polysemous_training = false;
     polysemous_ht = 0;
     polysemous_training = nullptr;
+    precomputed_table = new AlignedTable<float>();
+    owns_precomputed_table = true;
 }
 
 struct CodeCmp {

diff --git a/faiss/IndexIVFPQ.h b/faiss/IndexIVFPQ.h
@@ -48,7 +48,8 @@ struct IndexIVFPQ : IndexIVF {
 
     /// if use_precompute_table
     /// size nlist * pq.M * pq.ksub
-    AlignedTable<float> precomputed_table;
+    bool owns_precomputed_table;
+    AlignedTable<float>* precomputed_table;
 
     IndexIVFPQ(
             Index* quantizer,
@@ -58,6 +59,10 @@ struct IndexIVFPQ : IndexIVF {
             size_t nbits_per_idx,
             MetricType metric = METRIC_L2);
 
+    IndexIVFPQ(const IndexIVFPQ& orig);
+
+    ~IndexIVFPQ();
+
     void encode_vectors(
             idx_t n,
             const float* x,
@@ -139,6 +144,15 @@ struct IndexIVFPQ : IndexIVF {
     /// build precomputed table
     void precompute_table();
 
+    /**
+     * Initialize the precomputed table
+     * @param precompute_table
+     * @param _use_precomputed_table
+     */
+    void set_precomputed_table(
+            AlignedTable<float>* precompute_table,
+            int _use_precomputed_table);
+
     IndexIVFPQ();
 };
 

diff --git a/faiss/IndexIVFPQFastScan.cpp b/faiss/IndexIVFPQFastScan.cpp
@@ -46,13 +46,26 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(
         : IndexIVFFastScan(quantizer, d, nlist, 0, metric), pq(d, M, nbits) {
     by_residual = false; // set to false by default because it's faster
 
+    precomputed_table = new AlignedTable<float>();
+    owns_precomputed_table = true;
     init_fastscan(M, nbits, nlist, metric, bbs);
 }
 
 IndexIVFPQFastScan::IndexIVFPQFastScan() {
     by_residual = false;
     bbs = 0;
     M2 = 0;
+    precomputed_table = new AlignedTable<float>();
+    owns_precomputed_table = true;
+}
+
+IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQFastScan& orig)
+        : IndexIVFFastScan(orig), pq(orig.pq) {
+    by_residual = orig.by_residual;
+    bbs = orig.bbs;
+    M2 = orig.M2;
+    precomputed_table = new AlignedTable<float>(*orig.precomputed_table);
+    owns_precomputed_table = true;
 }
 
 IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
@@ -71,13 +84,15 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
     ntotal = orig.ntotal;
     is_trained = orig.is_trained;
     nprobe = orig.nprobe;
+    precomputed_table = new AlignedTable<float>();
+    owns_precomputed_table = true;
 
-    precomputed_table.resize(orig.precomputed_table.size());
+    precomputed_table->resize(orig.precomputed_table->size());
 
-    if (precomputed_table.nbytes() > 0) {
-        memcpy(precomputed_table.get(),
-               orig.precomputed_table.data(),
-               precomputed_table.nbytes());
+    if (precomputed_table->nbytes() > 0) {
+        memcpy(precomputed_table->get(),
+               orig.precomputed_table->data(),
+               precomputed_table->nbytes());
     }
 
     for (size_t i = 0; i < nlist; i++) {
@@ -102,6 +117,12 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
     orig_invlists = orig.invlists;
 }
 
+IndexIVFPQFastScan::~IndexIVFPQFastScan() {
+    if (owns_precomputed_table) {
+        delete precomputed_table;
+    }
+}
+
 /*********************************************************
  * Training
  *********************************************************/
@@ -127,11 +148,23 @@ void IndexIVFPQFastScan::precompute_table() {
             use_precomputed_table,
             quantizer,
             pq,
-            precomputed_table,
+            *precomputed_table,
             by_residual,
             verbose);
 }
 
+void IndexIVFPQFastScan::set_precomputed_table(
+        AlignedTable<float>* _precompute_table,
+        int _use_precomputed_table) {
+    // Clean up old pre-computed table
+    if (owns_precomputed_table) {
+        delete precomputed_table;
+    }
+    owns_precomputed_table = false;
+    precomputed_table = _precompute_table;
+    use_precomputed_table = _use_precomputed_table;
+}
+
 /*********************************************************
  * Code management functions
  *********************************************************/
@@ -229,7 +262,7 @@ void IndexIVFPQFastScan::compute_LUT(
                     if (cij >= 0) {
                         fvec_madd_simd(
                                 dim12,
-                                precomputed_table.get() + cij * dim12,
+                                precomputed_table->get() + cij * dim12,
                                 -2,
                                 ip_table.get() + i * dim12,
                                 tab);

diff --git a/faiss/IndexIVFPQFastScan.h b/faiss/IndexIVFPQFastScan.h
@@ -38,7 +38,8 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
     /// precomputed tables management
     int use_precomputed_table = 0;
     /// if use_precompute_table size (nlist, pq.M, pq.ksub)
-    AlignedTable<float> precomputed_table;
+    bool owns_precomputed_table;
+    AlignedTable<float>* precomputed_table;
 
     IndexIVFPQFastScan(
             Index* quantizer,
@@ -51,6 +52,10 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
 
     IndexIVFPQFastScan();
 
+    IndexIVFPQFastScan(const IndexIVFPQFastScan& orig);
+
+    ~IndexIVFPQFastScan();
+
     // built from an IndexIVFPQ
     explicit IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs = 32);
 
@@ -60,6 +65,10 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
 
     /// build precomputed table, possibly updating use_precomputed_table
     void precompute_table();
+    /// Pass in externally a precomputed
+    void set_precomputed_table(
+            AlignedTable<float>* precompute_table,
+            int _use_precomputed_table);
 
     /// same as the regular IVFPQ encoder. The codes are not reorganized by
     /// blocks a that point

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -35,6 +35,7 @@ set(FAISS_TEST_SRC
   test_partitioning.cpp
   test_fastscan_perf.cpp
   test_disable_pq_sdc_tables.cpp
+  test_ivfpq_share_table.cpp
 )
 
 add_executable(faiss_test ${FAISS_TEST_SRC})

diff --git a/tests/test_disable_pq_sdc_tables.cpp b/tests/test_disable_pq_sdc_tables.cpp
@@ -15,7 +15,9 @@
 #include "faiss/index_io.h"
 #include "test_util.h"
 
-pthread_mutex_t temp_file_mutex = PTHREAD_MUTEX_INITIALIZER;
+namespace {
+    pthread_mutex_t temp_file_mutex = PTHREAD_MUTEX_INITIALIZER;
+}
 
 TEST(IO, TestReadHNSWPQ_whenSDCDisabledFlagPassed_thenDisableSDCTable) {
     Tempfilename index_filename(&temp_file_mutex, "/tmp/faiss_TestReadHNSWPQ");