faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/IndexIVFPQ.cpp

@@ -9,10 +9,10 @@
 
 #include <faiss/IndexIVFPQ.h>
 
-#include <stdint.h>
 #include <cassert>
 #include <cinttypes>
 #include <cmath>
+#include <cstdint>
 #include <cstdio>
 
 #include <algorithm>
@@ -64,74 +64,16 @@ IndexIVFPQ::IndexIVFPQ(
 /****************************************************************
  * training                                                     */
 
-void IndexIVFPQ::train_residual(idx_t n, const float* x) {
-    train_residual_o(n, x, nullptr);
-}
-
-void IndexIVFPQ::train_residual_o(idx_t n, const float* x, float* residuals_2) {
-    const float* x_in = x;
-
-    x = fvecs_maybe_subsample(
-            d,
-            (size_t*)&n,
-            pq.cp.max_points_per_centroid * pq.ksub,
-            x,
-            verbose,
-            pq.cp.seed);
-
-    ScopeDeleter<float> del_x(x_in == x ? nullptr : x);
-
-    const float* trainset;
-    ScopeDeleter<float> del_residuals;
-    if (by_residual) {
-        if (verbose)
-            printf("computing residuals\n");
-        idx_t* assign = new idx_t[n]; // assignement to coarse centroids
-        ScopeDeleter<idx_t> del(assign);
-        quantizer->assign(n, x, assign);
-        float* residuals = new float[n * d];
-        del_residuals.set(residuals);
-        for (idx_t i = 0; i < n; i++)
-            quantizer->compute_residual(
-                    x + i * d, residuals + i * d, assign[i]);
-
-        trainset = residuals;
-    } else {
-        trainset = x;
-    }
-    if (verbose)
-        printf("training %zdx%zd product quantizer on %" PRId64
-               " vectors in %dD\n",
-               pq.M,
-               pq.ksub,
-               n,
-               d);
-    pq.verbose = verbose;
-    pq.train(n, trainset);
+void IndexIVFPQ::train_encoder(idx_t n, const float* x, const idx_t* assign) {
+    pq.train(n, x);
 
     if (do_polysemous_training) {
         if (verbose)
             printf("doing polysemous training for PQ\n");
         PolysemousTraining default_pt;
-        PolysemousTraining* pt = polysemous_training;
-        if (!pt)
-            pt = &default_pt;
-        pt->optimize_pq_for_hamming(pq, n, trainset);
-    }
-
-    // prepare second-level residuals for refine PQ
-    if (residuals_2) {
-        uint8_t* train_codes = new uint8_t[pq.code_size * n];
-        ScopeDeleter<uint8_t> del(train_codes);
-        pq.compute_codes(trainset, train_codes, n);
-
-        for (idx_t i = 0; i < n; i++) {
-            const float* xx = trainset + i * d;
-            float* res = residuals_2 + i * d;
-            pq.decode(train_codes + i * pq.code_size, res);
-            for (int j = 0; j < d; j++)
-                res[j] = xx[j] - res[j];
-        }
+        PolysemousTraining* pt =
+                polysemous_training ? polysemous_training : &default_pt;
+        pt->optimize_pq_for_hamming(pq, n, x);
     }
 
     if (by_residual) {
@@ -139,6 +81,10 @@ void IndexIVFPQ::train_residual_o(idx_t n, const float* x, float* residuals_2) {
     }
 }
 
+idx_t IndexIVFPQ::train_encoder_num_vectors() const {
+    return pq.cp.max_points_per_centroid * pq.ksub;
+}
+
 /****************************************************************
  * IVFPQ as codec                                               */
 
@@ -189,24 +135,25 @@ void IndexIVFPQ::add_core(
         idx_t n,
         const float* x,
         const idx_t* xids,
-        const idx_t* coarse_idx) {
-    add_core_o(n, x, xids, nullptr, coarse_idx);
+        const idx_t* coarse_idx,
+        void* inverted_list_context) {
+    add_core_o(n, x, xids, nullptr, coarse_idx, inverted_list_context);
 }
 
-static float* compute_residuals(
+static std::unique_ptr<float[]> compute_residuals(
         const Index* quantizer,
         idx_t n,
         const float* x,
         const idx_t* list_nos) {
     size_t d = quantizer->d;
-    float* residuals = new float[n * d];
+    std::unique_ptr<float[]> residuals(new float[n * d]);
     // TODO: parallelize?
     for (size_t i = 0; i < n; i++) {
         if (list_nos[i] < 0)
-            memset(residuals + i * d, 0, sizeof(*residuals) * d);
+            memset(residuals.get() + i * d, 0, sizeof(float) * d);
         else
             quantizer->compute_residual(
-                    x + i * d, residuals + i * d, list_nos[i]);
+                    x + i * d, residuals.get() + i * d, list_nos[i]);
     }
     return residuals;
 }
@@ -218,9 +165,9 @@ void IndexIVFPQ::encode_vectors(
         uint8_t* codes,
         bool include_listnos) const {
     if (by_residual) {
-        float* to_encode = compute_residuals(quantizer, n, x, list_nos);
-        ScopeDeleter<float> del(to_encode);
-        pq.compute_codes(to_encode, codes, n);
+        std::unique_ptr<float[]> to_encode =
+                compute_residuals(quantizer, n, x, list_nos);
+        pq.compute_codes(to_encode.get(), codes, n);
     } else {
         pq.compute_codes(x, codes, n);
     }
@@ -266,7 +213,8 @@ void IndexIVFPQ::add_core_o(
         const float* x,
         const idx_t* xids,
         float* residuals_2,
-        const idx_t* precomputed_idx) {
+        const idx_t* precomputed_idx,
+        void* inverted_list_context) {
     idx_t bs = index_ivfpq_add_core_o_bs;
     if (n > bs) {
         for (idx_t i0 = 0; i0 < n; i0 += bs) {
@@ -283,7 +231,8 @@ void IndexIVFPQ::add_core_o(
                     x + i0 * d,
                     xids ? xids + i0 : nullptr,
                     residuals_2 ? residuals_2 + i0 * d : nullptr,
-                    precomputed_idx ? precomputed_idx + i0 : nullptr);
+                    precomputed_idx ? precomputed_idx + i0 : nullptr,
+                    inverted_list_context);
         }
         return;
     }
@@ -295,31 +244,30 @@ void IndexIVFPQ::add_core_o(
     FAISS_THROW_IF_NOT(is_trained);
     double t0 = getmillisecs();
     const idx_t* idx;
-    ScopeDeleter<idx_t> del_idx;
+    std::unique_ptr<idx_t[]> del_idx;
 
     if (precomputed_idx) {
         idx = precomputed_idx;
     } else {
         idx_t* idx0 = new idx_t[n];
-        del_idx.set(idx0);
+        del_idx.reset(idx0);
         quantizer->assign(n, x, idx0);
         idx = idx0;
     }
 
     double t1 = getmillisecs();
-    uint8_t* xcodes = new uint8_t[n * code_size];
-    ScopeDeleter<uint8_t> del_xcodes(xcodes);
+    std::unique_ptr<uint8_t[]> xcodes(new uint8_t[n * code_size]);
 
     const float* to_encode = nullptr;
-    ScopeDeleter<float> del_to_encode;
+    std::unique_ptr<const float[]> del_to_encode;
 
     if (by_residual) {
-        to_encode = compute_residuals(quantizer, n, x, idx);
-        del_to_encode.set(to_encode);
+        del_to_encode = compute_residuals(quantizer, n, x, idx);
+        to_encode = del_to_encode.get();
     } else {
         to_encode = x;
     }
-    pq.compute_codes(to_encode, xcodes, n);
+    pq.compute_codes(to_encode, xcodes.get(), n);
 
     double t2 = getmillisecs();
     // TODO: parallelize?
@@ -335,8 +283,9 @@ void IndexIVFPQ::add_core_o(
             continue;
         }
 
-        uint8_t* code = xcodes + i * code_size;
-        size_t offset = invlists->add_entry(key, id, code);
+        uint8_t* code = xcodes.get() + i * code_size;
+        size_t offset =
+                invlists->add_entry(key, id, code, inverted_list_context);
 
         if (residuals_2) {
             float* res2 = residuals_2 + i * d;
@@ -369,16 +318,14 @@ void IndexIVFPQ::reconstruct_from_offset(
         float* recons) const {
     const uint8_t* code = invlists->get_single_code(list_no, offset);
 
+    pq.decode(code, recons);
     if (by_residual) {
         std::vector<float> centroid(d);
         quantizer->reconstruct(list_no, centroid.data());
 
-        pq.decode(code, recons);
         for (int i = 0; i < d; ++i) {
             recons[i] += centroid[i];
         }
-    } else {
-        pq.decode(code, recons);
     }
 }
 
@@ -804,7 +751,7 @@ struct QueryTables {
     }
 };
 
-// This way of handling the sleector is not optimal since all distances
+// This way of handling the selector is not optimal since all distances
 // are computed even if the id would filter it out.
 template <class C, bool use_sel>
 struct KnnSearchResults {
@@ -937,7 +884,8 @@ struct IVFPQScannerT : QueryTables {
                 float distance_2 = 0;
                 float distance_3 = 0;
                 distance_four_codes<PQDecoder>(
-                        pq,
+                        pq.M,
+                        pq.nbits,
                         sim_table,
                         codes + saved_j[0] * pq.code_size,
                         codes + saved_j[1] * pq.code_size,
@@ -957,24 +905,30 @@ struct IVFPQScannerT : QueryTables {
         }
 
         if (counter >= 1) {
-            float dis =
-                    dis0 +
+            float dis = dis0 +
                     distance_single_code<PQDecoder>(
-                            pq, sim_table, codes + saved_j[0] * pq.code_size);
+                                pq.M,
+                                pq.nbits,
+                                sim_table,
+                                codes + saved_j[0] * pq.code_size);
             res.add(saved_j[0], dis);
         }
         if (counter >= 2) {
-            float dis =
-                    dis0 +
+            float dis = dis0 +
                     distance_single_code<PQDecoder>(
-                            pq, sim_table, codes + saved_j[1] * pq.code_size);
+                                pq.M,
+                                pq.nbits,
+                                sim_table,
+                                codes + saved_j[1] * pq.code_size);
             res.add(saved_j[1], dis);
         }
         if (counter >= 3) {
-            float dis =
-                    dis0 +
+            float dis = dis0 +
                     distance_single_code<PQDecoder>(
-                            pq, sim_table, codes + saved_j[2] * pq.code_size);
+                                pq.M,
+                                pq.nbits,
+                                sim_table,
+                                codes + saved_j[2] * pq.code_size);
             res.add(saved_j[2], dis);
         }
     }
@@ -1090,7 +1044,7 @@ struct IVFPQScannerT : QueryTables {
             const uint8_t* codes,
             SearchResultType& res) const {
         int ht = ivfpq.polysemous_ht;
-        size_t n_hamming_pass = 0, nup = 0;
+        size_t n_hamming_pass = 0;
 
         int code_size = pq.code_size;
 
@@ -1137,7 +1091,8 @@ struct IVFPQScannerT : QueryTables {
                 float distance_2 = dis0;
                 float distance_3 = dis0;
                 distance_four_codes<PQDecoder>(
-                        pq,
+                        pq.M,
+                        pq.nbits,
                         sim_table,
                         codes + saved_j[0] * pq.code_size,
                         codes + saved_j[1] * pq.code_size,
@@ -1165,10 +1120,12 @@ struct IVFPQScannerT : QueryTables {
         for (size_t kk = 0; kk < counter; kk++) {
             n_hamming_pass++;
 
-            float dis =
-                    dis0 +
+            float dis = dis0 +
                     distance_single_code<PQDecoder>(
-                            pq, sim_table, codes + saved_j[kk] * pq.code_size);
+                                pq.M,
+                                pq.nbits,
+                                sim_table,
+                                codes + saved_j[kk] * pq.code_size);
 
             res.add(saved_j[kk], dis);
         }
@@ -1185,7 +1142,10 @@ struct IVFPQScannerT : QueryTables {
 
                 float dis = dis0 +
                         distance_single_code<PQDecoder>(
-                                    pq, sim_table, codes + j * code_size);
+                                    pq.M,
+                                    pq.nbits,
+                                    sim_table,
+                                    codes + j * code_size);
 
                 res.add(j, dis);
             }
@@ -1195,30 +1155,23 @@ struct IVFPQScannerT : QueryTables {
         { indexIVFPQ_stats.n_hamming_pass += n_hamming_pass; }
     }
 
+    template <class SearchResultType>
+    struct Run_scan_list_polysemous_hc {
+        using T = void;
+        template <class HammingComputer, class... Types>
+        void f(const IVFPQScannerT* scanner, Types... args) {
+            scanner->scan_list_polysemous_hc<HammingComputer, SearchResultType>(
+                    args...);
+        }
+    };
+
     template <class SearchResultType>
     void scan_list_polysemous(
             size_t ncode,
             const uint8_t* codes,
             SearchResultType& res) const {
-        switch (pq.code_size) {
-#define HANDLE_CODE_SIZE(cs)                                            \
-    case cs:                                                            \
-        scan_list_polysemous_hc<HammingComputer##cs, SearchResultType>( \
-                ncode, codes, res);                                     \
-        break
-            HANDLE_CODE_SIZE(4);
-            HANDLE_CODE_SIZE(8);
-            HANDLE_CODE_SIZE(16);
-            HANDLE_CODE_SIZE(20);
-            HANDLE_CODE_SIZE(32);
-            HANDLE_CODE_SIZE(64);
-#undef HANDLE_CODE_SIZE
-            default:
-                scan_list_polysemous_hc<
-                        HammingComputerDefault,
-                        SearchResultType>(ncode, codes, res);
-                break;
-        }
+        Run_scan_list_polysemous_hc<SearchResultType> r;
+        dispatch_HammingComputer(pq.code_size, r, this, ncode, codes, res);
     }
 };
 
@@ -1248,6 +1201,7 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
               precompute_mode(precompute_mode),
               sel(sel) {
         this->store_pairs = store_pairs;
+        this->keep_max = is_similarity_metric(METRIC_TYPE);
     }
 
     void set_query(const float* query) override {
@@ -1263,7 +1217,7 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
         assert(precompute_mode == 2);
         float dis = this->dis0 +
                 distance_single_code<PQDecoder>(
-                            this->pq, this->sim_table, code);
+                            this->pq.M, this->pq.nbits, this->sim_table, code);
         return dis;
     }
 

data/vendor/faiss/faiss/IndexIVFPQ.h

@@ -32,8 +32,6 @@ FAISS_API extern size_t precomputed_table_max_bytes;
  * vector is encoded as a product quantizer code.
  */
 struct IndexIVFPQ : IndexIVF {
-    bool by_residual; ///< Encode residual or plain vector?
-
     ProductQuantizer pq; ///< produces the codes
 
     bool do_polysemous_training; ///< reorder PQ centroids after training?
@@ -73,7 +71,8 @@ struct IndexIVFPQ : IndexIVF {
             idx_t n,
             const float* x,
             const idx_t* xids,
-            const idx_t* precomputed_idx) override;
+            const idx_t* precomputed_idx,
+            void* inverted_list_context = nullptr) override;
 
     /// same as add_core, also:
     /// - output 2nd level residuals if residuals_2 != NULL
@@ -83,13 +82,13 @@ struct IndexIVFPQ : IndexIVF {
             const float* x,
             const idx_t* xids,
             float* residuals_2,
-            const idx_t* precomputed_idx = nullptr);
+            const idx_t* precomputed_idx = nullptr,
+            void* inverted_list_context = nullptr);
 
     /// trains the product quantizer
-    void train_residual(idx_t n, const float* x) override;
+    void train_encoder(idx_t n, const float* x, const idx_t* assign) override;
 
-    /// same as train_residual, also output 2nd level residuals
-    void train_residual_o(idx_t n, const float* x, float* residuals_2);
+    idx_t train_encoder_num_vectors() const override;
 
     void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
             const override;

data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp

@@ -44,7 +44,7 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(
         MetricType metric,
         int bbs)
         : IndexIVFFastScan(quantizer, d, nlist, 0, metric), pq(d, M, nbits) {
-    by_residual = false; // set to false by default because it's much faster
+    by_residual = false; // set to false by default because it's faster
 
     init_fastscan(M, nbits, nlist, metric, bbs);
 }
@@ -106,54 +106,22 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
  * Training
  *********************************************************/
 
-void IndexIVFPQFastScan::train_residual(idx_t n, const float* x_in) {
-    const float* x = fvecs_maybe_subsample(
-            d,
-            (size_t*)&n,
-            pq.cp.max_points_per_centroid * pq.ksub,
-            x_in,
-            verbose,
-            pq.cp.seed);
-
-    std::unique_ptr<float[]> del_x;
-    if (x != x_in) {
-        del_x.reset((float*)x);
-    }
-
-    const float* trainset;
-    AlignedTable<float> residuals;
-
-    if (by_residual) {
-        if (verbose)
-            printf("computing residuals\n");
-        std::vector<idx_t> assign(n);
-        quantizer->assign(n, x, assign.data());
-        residuals.resize(n * d);
-        for (idx_t i = 0; i < n; i++) {
-            quantizer->compute_residual(
-                    x + i * d, residuals.data() + i * d, assign[i]);
-        }
-        trainset = residuals.data();
-    } else {
-        trainset = x;
-    }
-
-    if (verbose) {
-        printf("training %zdx%zd product quantizer on "
-               "%" PRId64 " vectors in %dD\n",
-               pq.M,
-               pq.ksub,
-               n,
-               d);
-    }
+void IndexIVFPQFastScan::train_encoder(
+        idx_t n,
+        const float* x,
+        const idx_t* assign) {
     pq.verbose = verbose;
-    pq.train(n, trainset);
+    pq.train(n, x);
 
     if (by_residual && metric_type == METRIC_L2) {
         precompute_table();
     }
 }
 
+idx_t IndexIVFPQFastScan::train_encoder_num_vectors() const {
+    return pq.cp.max_points_per_centroid * pq.ksub;
+}
+
 void IndexIVFPQFastScan::precompute_table() {
     initialize_IVFPQ_precomputed_table(
             use_precomputed_table,
@@ -203,7 +171,7 @@ void IndexIVFPQFastScan::encode_vectors(
  * Look-Up Table functions
  *********************************************************/
 
-void fvec_madd_avx(
+void fvec_madd_simd(
         size_t n,
         const float* a,
         float bf,
@@ -234,12 +202,12 @@ bool IndexIVFPQFastScan::lookup_table_is_3d() const {
 void IndexIVFPQFastScan::compute_LUT(
         size_t n,
         const float* x,
-        const idx_t* coarse_ids,
-        const float* coarse_dis,
+        const CoarseQuantized& cq,
         AlignedTable<float>& dis_tables,
         AlignedTable<float>& biases) const {
     size_t dim12 = pq.ksub * pq.M;
     size_t d = pq.d;
+    size_t nprobe = this->nprobe;
 
     if (by_residual) {
         if (metric_type == METRIC_L2) {
@@ -247,7 +215,7 @@ void IndexIVFPQFastScan::compute_LUT(
 
             if (use_precomputed_table == 1) {
                 biases.resize(n * nprobe);
-                memcpy(biases.get(), coarse_dis, sizeof(float) * n * nprobe);
+                memcpy(biases.get(), cq.dis, sizeof(float) * n * nprobe);
 
                 AlignedTable<float> ip_table(n * dim12);
                 pq.compute_inner_prod_tables(n, x, ip_table.get());
@@ -256,10 +224,10 @@ void IndexIVFPQFastScan::compute_LUT(
                 for (idx_t ij = 0; ij < n * nprobe; ij++) {
                     idx_t i = ij / nprobe;
                     float* tab = dis_tables.get() + ij * dim12;
-                    idx_t cij = coarse_ids[ij];
+                    idx_t cij = cq.ids[ij];
 
                     if (cij >= 0) {
-                        fvec_madd_avx(
+                        fvec_madd_simd(
                                 dim12,
                                 precomputed_table.get() + cij * dim12,
                                 -2,
@@ -281,7 +249,7 @@ void IndexIVFPQFastScan::compute_LUT(
                 for (idx_t ij = 0; ij < n * nprobe; ij++) {
                     idx_t i = ij / nprobe;
                     float* xij = &xrel[ij * d];
-                    idx_t cij = coarse_ids[ij];
+                    idx_t cij = cq.ids[ij];
 
                     if (cij >= 0) {
                         quantizer->compute_residual(x + i * d, xij, cij);
@@ -301,7 +269,7 @@ void IndexIVFPQFastScan::compute_LUT(
             // compute_inner_prod_tables(pq, n, x, dis_tables.get());
 
             biases.resize(n * nprobe);
-            memcpy(biases.get(), coarse_dis, sizeof(float) * n * nprobe);
+            memcpy(biases.get(), cq.dis, sizeof(float) * n * nprobe);
         } else {
             FAISS_THROW_FMT("metric %d not supported", metric_type);
         }
@@ -318,9 +286,28 @@ void IndexIVFPQFastScan::compute_LUT(
     }
 }
 
-void IndexIVFPQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x)
+void IndexIVFPQFastScan::sa_decode(idx_t n, const uint8_t* codes, float* x)
         const {
-    pq.decode(bytes, x, n);
+    size_t coarse_size = coarse_code_size();
+
+#pragma omp parallel if (n > 1)
+    {
+        std::vector<float> residual(d);
+
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            const uint8_t* code = codes + i * (code_size + coarse_size);
+            int64_t list_no = decode_listno(code);
+            float* xi = x + i * d;
+            pq.decode(code + coarse_size, xi);
+            if (by_residual) {
+                quantizer->reconstruct(list_no, residual.data());
+                for (size_t j = 0; j < d; j++) {
+                    xi[j] += residual[j];
+                }
+            }
+        }
+    }
 }
 
 } // namespace faiss

data/vendor/faiss/faiss/IndexIVFPQFastScan.h

@@ -54,7 +54,9 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
     // built from an IndexIVFPQ
     explicit IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs = 32);
 
-    void train_residual(idx_t n, const float* x) override;
+    void train_encoder(idx_t n, const float* x, const idx_t* assign) override;
+
+    idx_t train_encoder_num_vectors() const override;
 
     /// build precomputed table, possibly updating use_precomputed_table
     void precompute_table();
@@ -75,8 +77,7 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
     void compute_LUT(
             size_t n,
             const float* x,
-            const idx_t* coarse_ids,
-            const float* coarse_dis,
+            const CoarseQuantized& cq,
             AlignedTable<float>& dis_tables,
             AlignedTable<float>& biases) const override;