faiss 0.1.4 → 0.2.1

data/vendor/faiss/faiss/IndexIVFPQFastScan.h

@@ -15,7 +15,6 @@
 
 namespace faiss {
 
-
 /** Fast scan version of IVFPQ. Works for 4-bit PQ for now.
  *
  * The codes in the inverted lists are not stored sequentially but
@@ -32,13 +31,12 @@ namespace faiss {
  * 13: idem, collect results in reservoir
  */
 
-struct IndexIVFPQFastScan: IndexIVF {
-
-    bool by_residual;              ///< Encode residual or plain vector?
-    ProductQuantizer pq;           ///< produces the codes
+struct IndexIVFPQFastScan : IndexIVF {
+    bool by_residual;    ///< Encode residual or plain vector?
+    ProductQuantizer pq; ///< produces the codes
 
     // size of the kernel
-    int bbs;  // set at build time
+    int bbs; // set at build time
 
     // M rounded up to a multiple of 2
     size_t M2;
@@ -57,88 +55,114 @@ struct IndexIVFPQFastScan: IndexIVF {
     int qbs = 0;
     size_t qbs2 = 0;
 
-    IndexIVFPQFastScan (
-            Index * quantizer, size_t d, size_t nlist,
-            size_t M, size_t nbits_per_idx,
-            MetricType metric = METRIC_L2, int bbs = 32);
+    IndexIVFPQFastScan(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            size_t M,
+            size_t nbits_per_idx,
+            MetricType metric = METRIC_L2,
+            int bbs = 32);
 
-    IndexIVFPQFastScan ();
+    IndexIVFPQFastScan();
 
     // built from an IndexIVFPQ
-    explicit IndexIVFPQFastScan(const IndexIVFPQ & orig, int bbs = 32);
+    explicit IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs = 32);
 
     /// orig's inverted lists (for debugging)
-    InvertedLists * orig_invlists = nullptr;
+    InvertedLists* orig_invlists = nullptr;
 
-    void train_residual (idx_t n, const float *x) override;
+    void train_residual(idx_t n, const float* x) override;
 
     /// build precomputed table, possibly updating use_precomputed_table
-    void precompute_table ();
+    void precompute_table();
 
     /// same as the regular IVFPQ encoder. The codes are not reorganized by
     /// blocks a that point
     void encode_vectors(
-            idx_t n, const float* x,
-            const idx_t *list_nos, uint8_t * codes,
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            uint8_t* codes,
             bool include_listno = false) const override;
 
-    void add_with_ids (
-        idx_t n, const float * x, const idx_t *xids) override;
+    void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
 
     void search(
-            idx_t n, const float* x, idx_t k,
-            float* distances, idx_t* labels) const override;
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
 
     // prepare look-up tables
 
     void compute_LUT(
-        size_t n, const float *x,
-        const idx_t *coarse_ids, const float *coarse_dis,
-        AlignedTable<float> & dis_tables,
-        AlignedTable<float> & biases
-    ) const;
+            size_t n,
+            const float* x,
+            const idx_t* coarse_ids,
+            const float* coarse_dis,
+            AlignedTable<float>& dis_tables,
+            AlignedTable<float>& biases) const;
 
     void compute_LUT_uint8(
-        size_t n, const float *x,
-        const idx_t *coarse_ids, const float *coarse_dis,
-        AlignedTable<uint8_t> & dis_tables,
-        AlignedTable<uint16_t> & biases,
-        float * normalizers
-    ) const;
+            size_t n,
+            const float* x,
+            const idx_t* coarse_ids,
+            const float* coarse_dis,
+            AlignedTable<uint8_t>& dis_tables,
+            AlignedTable<uint16_t>& biases,
+            float* normalizers) const;
 
     // internal search funcs
 
-    template<bool is_max>
+    template <bool is_max>
     void search_dispatch_implem(
-            idx_t n, const float* x, idx_t k,
-            float* distances, idx_t* labels) const;
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const;
 
-    template<class C>
+    template <class C>
     void search_implem_1(
-            idx_t n, const float* x, idx_t k,
-            float* distances, idx_t* labels) const;
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const;
 
-    template<class C>
+    template <class C>
     void search_implem_2(
-            idx_t n, const float* x, idx_t k,
-            float* distances, idx_t* labels) const;
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const;
 
     // implem 10 and 12 are not multithreaded internally, so
     // export search stats
-    template<class C>
+    template <class C>
     void search_implem_10(
-            idx_t n, const float* x, idx_t k,
-            float* distances, idx_t* labels,
-            int impl, size_t *ndis_out, size_t *nlist_out) const;
-
-    template<class C>
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            int impl,
+            size_t* ndis_out,
+            size_t* nlist_out) const;
+
+    template <class C>
     void search_implem_12(
-            idx_t n, const float* x, idx_t k,
-            float* distances, idx_t* labels,
-            int impl, size_t *ndis_out, size_t *nlist_out) const;
-
-
-
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            int impl,
+            size_t* ndis_out,
+            size_t* nlist_out) const;
 };
 
 struct IVFFastScanStats {
@@ -148,13 +172,15 @@ struct IVFFastScanStats {
     uint64_t reservoir_times[4];
 
     double Mcy_at(int i) {
-        return times[i] / (1000*1000.0);
+        return times[i] / (1000 * 1000.0);
     }
 
     double Mcy_reservoir_at(int i) {
-        return reservoir_times[i] / (1000*1000.0);
+        return reservoir_times[i] / (1000 * 1000.0);
+    }
+    IVFFastScanStats() {
+        reset();
     }
-    IVFFastScanStats() {reset();}
     void reset() {
         memset(this, 0, sizeof(*this));
     }
@@ -162,5 +188,4 @@ struct IVFFastScanStats {
 
 FAISS_API extern IVFFastScanStats IVFFastScan_stats;
 
-
 } // namespace faiss

data/vendor/faiss/faiss/IndexIVFPQR.cpp

@@ -12,115 +12,116 @@
 #include <cinttypes>
 
 #include <faiss/utils/Heap.h>
-#include <faiss/utils/utils.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/utils.h>
 
 #include <faiss/impl/FaissAssert.h>
 
-
 namespace faiss {
 
 /*****************************************
  * IndexIVFPQR implementation
  ******************************************/
 
-IndexIVFPQR::IndexIVFPQR (
-            Index * quantizer, size_t d, size_t nlist,
-            size_t M, size_t nbits_per_idx,
-            size_t M_refine, size_t nbits_per_idx_refine):
-    IndexIVFPQ (quantizer, d, nlist, M, nbits_per_idx),
-    refine_pq (d, M_refine, nbits_per_idx_refine),
-    k_factor (4)
-{
+IndexIVFPQR::IndexIVFPQR(
+        Index* quantizer,
+        size_t d,
+        size_t nlist,
+        size_t M,
+        size_t nbits_per_idx,
+        size_t M_refine,
+        size_t nbits_per_idx_refine)
+        : IndexIVFPQ(quantizer, d, nlist, M, nbits_per_idx),
+          refine_pq(d, M_refine, nbits_per_idx_refine),
+          k_factor(4) {
     by_residual = true;
 }
 
-IndexIVFPQR::IndexIVFPQR ():
-    k_factor (1)
-{
+IndexIVFPQR::IndexIVFPQR() : k_factor(1) {
     by_residual = true;
 }
 
-
-
-void IndexIVFPQR::reset()
-{
+void IndexIVFPQR::reset() {
     IndexIVFPQ::reset();
     refine_codes.clear();
 }
 
+void IndexIVFPQR::train_residual(idx_t n, const float* x) {
+    float* residual_2 = new float[n * d];
+    ScopeDeleter<float> del(residual_2);
 
-
-
-void IndexIVFPQR::train_residual (idx_t n, const float *x)
-{
-
-    float * residual_2 = new float [n * d];
-    ScopeDeleter <float> del(residual_2);
-
-    train_residual_o (n, x, residual_2);
+    train_residual_o(n, x, residual_2);
 
     if (verbose)
-        printf ("training %zdx%zd 2nd level PQ quantizer on %" PRId64 " %dD-vectors\n",
-                refine_pq.M, refine_pq.ksub, n, d);
+        printf("training %zdx%zd 2nd level PQ quantizer on %" PRId64
+               " %dD-vectors\n",
+               refine_pq.M,
+               refine_pq.ksub,
+               n,
+               d);
 
     refine_pq.cp.max_points_per_centroid = 1000;
     refine_pq.cp.verbose = verbose;
 
-    refine_pq.train (n, residual_2);
-
+    refine_pq.train(n, residual_2);
 }
 
-
-void IndexIVFPQR::add_with_ids (idx_t n, const float *x, const idx_t *xids) {
-    add_core (n, x, xids, nullptr);
+void IndexIVFPQR::add_with_ids(idx_t n, const float* x, const idx_t* xids) {
+    add_core(n, x, xids, nullptr);
 }
 
-void IndexIVFPQR::add_core (idx_t n, const float *x, const idx_t *xids,
-                                const idx_t *precomputed_idx) {
-
-    float * residual_2 = new float [n * d];
-    ScopeDeleter <float> del(residual_2);
+void IndexIVFPQR::add_core(
+        idx_t n,
+        const float* x,
+        const idx_t* xids,
+        const idx_t* precomputed_idx) {
+    float* residual_2 = new float[n * d];
+    ScopeDeleter<float> del(residual_2);
 
     idx_t n0 = ntotal;
 
-    add_core_o (n, x, xids, residual_2, precomputed_idx);
-
-    refine_codes.resize (ntotal * refine_pq.code_size);
-
-    refine_pq.compute_codes (
-        residual_2, &refine_codes[n0 * refine_pq.code_size], n);
+    add_core_o(n, x, xids, residual_2, precomputed_idx);
 
+    refine_codes.resize(ntotal * refine_pq.code_size);
 
+    refine_pq.compute_codes(
+            residual_2, &refine_codes[n0 * refine_pq.code_size], n);
 }
 #define TIC t0 = get_cycles()
-#define TOC get_cycles () - t0
-
-
-void IndexIVFPQR::search_preassigned (
-        idx_t n, const float *x, idx_t k,
-        const idx_t *idx, const float *L1_dis,
-        float *distances, idx_t *labels,
+#define TOC get_cycles() - t0
+
+void IndexIVFPQR::search_preassigned(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        const idx_t* idx,
+        const float* L1_dis,
+        float* distances,
+        idx_t* labels,
         bool store_pairs,
-        const IVFSearchParameters *params, IndexIVFStats *stats
-        ) const
-{
+        const IVFSearchParameters* params,
+        IndexIVFStats* stats) const {
     uint64_t t0;
     TIC;
     size_t k_coarse = long(k * k_factor);
-    idx_t *coarse_labels = new idx_t [k_coarse * n];
-    ScopeDeleter<idx_t> del1 (coarse_labels);
+    idx_t* coarse_labels = new idx_t[k_coarse * n];
+    ScopeDeleter<idx_t> del1(coarse_labels);
     { // query with quantizer levels 1 and 2.
-        float *coarse_distances = new float [k_coarse * n];
+        float* coarse_distances = new float[k_coarse * n];
         ScopeDeleter<float> del(coarse_distances);
 
-        IndexIVFPQ::search_preassigned (
-                   n, x, k_coarse,
-                   idx, L1_dis, coarse_distances, coarse_labels,
-                   true, params);
+        IndexIVFPQ::search_preassigned(
+                n,
+                x,
+                k_coarse,
+                idx,
+                L1_dis,
+                coarse_distances,
+                coarse_labels,
+                true,
+                params);
     }
 
-
     indexIVFPQ_stats.search_cycles += TOC;
 
     TIC;
@@ -130,91 +131,92 @@ void IndexIVFPQR::search_preassigned (
 #pragma omp parallel reduction(+ : n_refine)
     {
         // tmp buffers
-        float *residual_1 = new float [2 * d];
-        ScopeDeleter<float> del (residual_1);
-        float *residual_2 = residual_1 + d;
+        float* residual_1 = new float[2 * d];
+        ScopeDeleter<float> del(residual_1);
+        float* residual_2 = residual_1 + d;
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            const float *xq = x + i * d;
-            const idx_t * shortlist = coarse_labels + k_coarse * i;
-            float * heap_sim = distances + k * i;
-            idx_t * heap_ids = labels + k * i;
-            maxheap_heapify (k, heap_sim, heap_ids);
+            const float* xq = x + i * d;
+            const idx_t* shortlist = coarse_labels + k_coarse * i;
+            float* heap_sim = distances + k * i;
+            idx_t* heap_ids = labels + k * i;
+            maxheap_heapify(k, heap_sim, heap_ids);
 
             for (int j = 0; j < k_coarse; j++) {
                 idx_t sl = shortlist[j];
 
-                if (sl == -1) continue;
+                if (sl == -1)
+                    continue;
 
                 int list_no = lo_listno(sl);
                 int ofs = lo_offset(sl);
 
-                assert (list_no >= 0 && list_no < nlist);
-                assert (ofs >= 0 && ofs < invlists->list_size (list_no));
+                assert(list_no >= 0 && list_no < nlist);
+                assert(ofs >= 0 && ofs < invlists->list_size(list_no));
 
                 // 1st level residual
-                quantizer->compute_residual (xq, residual_1, list_no);
+                quantizer->compute_residual(xq, residual_1, list_no);
 
                 // 2nd level residual
-                const uint8_t * l2code =
-                    invlists->get_single_code (list_no, ofs);
+                const uint8_t* l2code = invlists->get_single_code(list_no, ofs);
 
-                pq.decode (l2code, residual_2);
+                pq.decode(l2code, residual_2);
                 for (int l = 0; l < d; l++)
                     residual_2[l] = residual_1[l] - residual_2[l];
 
                 // 3rd level residual's approximation
-                idx_t id = invlists->get_single_id (list_no, ofs);
-                assert (0 <= id && id < ntotal);
-                refine_pq.decode (&refine_codes [id * refine_pq.code_size],
-                                  residual_1);
+                idx_t id = invlists->get_single_id(list_no, ofs);
+                assert(0 <= id && id < ntotal);
+                refine_pq.decode(
+                        &refine_codes[id * refine_pq.code_size], residual_1);
 
-                float dis = fvec_L2sqr (residual_1, residual_2, d);
+                float dis = fvec_L2sqr(residual_1, residual_2, d);
 
                 if (dis < heap_sim[0]) {
                     idx_t id_or_pair = store_pairs ? sl : id;
-                    maxheap_replace_top (k, heap_sim, heap_ids, dis, id_or_pair);
+                    maxheap_replace_top(k, heap_sim, heap_ids, dis, id_or_pair);
                 }
-                n_refine ++;
+                n_refine++;
             }
-            maxheap_reorder (k, heap_sim, heap_ids);
+            maxheap_reorder(k, heap_sim, heap_ids);
         }
     }
     indexIVFPQ_stats.nrefine += n_refine;
     indexIVFPQ_stats.refine_cycles += TOC;
 }
 
-void IndexIVFPQR::reconstruct_from_offset (int64_t list_no, int64_t offset,
-                                           float* recons) const
-{
-    IndexIVFPQ::reconstruct_from_offset (list_no, offset, recons);
+void IndexIVFPQR::reconstruct_from_offset(
+        int64_t list_no,
+        int64_t offset,
+        float* recons) const {
+    IndexIVFPQ::reconstruct_from_offset(list_no, offset, recons);
 
-    idx_t id = invlists->get_single_id (list_no, offset);
-    assert (0 <= id && id < ntotal);
+    idx_t id = invlists->get_single_id(list_no, offset);
+    assert(0 <= id && id < ntotal);
 
     std::vector<float> r3(d);
-    refine_pq.decode (&refine_codes [id * refine_pq.code_size], r3.data());
+    refine_pq.decode(&refine_codes[id * refine_pq.code_size], r3.data());
     for (int i = 0; i < d; ++i) {
-      recons[i] += r3[i];
+        recons[i] += r3[i];
     }
 }
 
-void IndexIVFPQR::merge_from (IndexIVF &other_in, idx_t add_id)
-{
-    IndexIVFPQR *other = dynamic_cast<IndexIVFPQR *> (&other_in);
+void IndexIVFPQR::merge_from(IndexIVF& other_in, idx_t add_id) {
+    IndexIVFPQR* other = dynamic_cast<IndexIVFPQR*>(&other_in);
     FAISS_THROW_IF_NOT(other);
 
-    IndexIVF::merge_from (other_in, add_id);
+    IndexIVF::merge_from(other_in, add_id);
 
-    refine_codes.insert (refine_codes.end(),
-                         other->refine_codes.begin(),
-                         other->refine_codes.end());
+    refine_codes.insert(
+            refine_codes.end(),
+            other->refine_codes.begin(),
+            other->refine_codes.end());
     other->refine_codes.clear();
 }
 
 size_t IndexIVFPQR::remove_ids(const IDSelector& /*sel*/) {
-  FAISS_THROW_MSG("not implemented");
-  return 0;
+    FAISS_THROW_MSG("not implemented");
+    return 0;
 }
 
 } // namespace faiss