faiss 0.1.4 → 0.2.1

data/vendor/faiss/faiss/IndexIVF.h

@@ -10,30 +10,27 @@
 #ifndef FAISS_INDEX_IVF_H
 #define FAISS_INDEX_IVF_H
 
-
-#include <vector>
-#include <unordered_map>
 #include <stdint.h>
+#include <unordered_map>
+#include <vector>
 
-#include <faiss/Index.h>
-#include <faiss/invlists/InvertedLists.h>
-#include <faiss/invlists/DirectMap.h>
 #include <faiss/Clustering.h>
+#include <faiss/Index.h>
 #include <faiss/impl/platform_macros.h>
+#include <faiss/invlists/DirectMap.h>
+#include <faiss/invlists/InvertedLists.h>
 #include <faiss/utils/Heap.h>
 
-
 namespace faiss {
 
-
 /** Encapsulates a quantizer object for the IndexIVF
  *
  * The class isolates the fields that are independent of the storage
  * of the lists (especially training)
  */
 struct Level1Quantizer {
-    Index * quantizer;        ///< quantizer that maps vectors to inverted lists
-    size_t nlist;             ///< number of possible key values
+    Index* quantizer; ///< quantizer that maps vectors to inverted lists
+    size_t nlist;     ///< number of possible key values
 
     /**
      * = 0: use the quantizer as index in a kmeans training
@@ -41,40 +38,37 @@ struct Level1Quantizer {
      * = 2: kmeans training on a flat index + add the centroids to the quantizer
      */
     char quantizer_trains_alone;
-    bool own_fields;          ///< whether object owns the quantizer
+    bool own_fields; ///< whether object owns the quantizer
 
     ClusteringParameters cp; ///< to override default clustering params
-    Index *clustering_index; ///< to override index used during clustering
+    Index* clustering_index; ///< to override index used during clustering
 
     /// Trains the quantizer and calls train_residual to train sub-quantizers
-    void train_q1 (size_t n, const float *x, bool verbose,
-                   MetricType metric_type);
-
+    void train_q1(
+            size_t n,
+            const float* x,
+            bool verbose,
+            MetricType metric_type);
 
     /// compute the number of bytes required to store list ids
-    size_t coarse_code_size () const;
-    void encode_listno (Index::idx_t list_no, uint8_t *code) const;
-    Index::idx_t decode_listno (const uint8_t *code) const;
+    size_t coarse_code_size() const;
+    void encode_listno(Index::idx_t list_no, uint8_t* code) const;
+    Index::idx_t decode_listno(const uint8_t* code) const;
 
-    Level1Quantizer (Index * quantizer, size_t nlist);
+    Level1Quantizer(Index* quantizer, size_t nlist);
 
-    Level1Quantizer ();
-
-    ~Level1Quantizer ();
+    Level1Quantizer();
 
+    ~Level1Quantizer();
 };
 
-
-
 struct IVFSearchParameters {
-    size_t nprobe;            ///< number of probes at query time
-    size_t max_codes;         ///< max nb of codes to visit to do a query
-    IVFSearchParameters(): nprobe(1), max_codes(0) {}
-    virtual ~IVFSearchParameters () {}
+    size_t nprobe;    ///< number of probes at query time
+    size_t max_codes; ///< max nb of codes to visit to do a query
+    IVFSearchParameters() : nprobe(1), max_codes(0) {}
+    virtual ~IVFSearchParameters() {}
 };
 
-
-
 struct InvertedListScanner;
 struct IndexIVFStats;
 
@@ -98,15 +92,15 @@ struct IndexIVFStats;
  * Sub-classes implement a post-filtering of the index that refines
  * the distance estimation from the query to databse vectors.
  */
-struct IndexIVF: Index, Level1Quantizer {
+struct IndexIVF : Index, Level1Quantizer {
     /// Access to the actual data
-    InvertedLists *invlists;
+    InvertedLists* invlists;
     bool own_invlists;
 
-    size_t code_size;              ///< code size per vector in bytes
+    size_t code_size; ///< code size per vector in bytes
 
-    size_t nprobe;            ///< number of probes at query time
-    size_t max_codes;         ///< max nb of codes to visit to do a query
+    size_t nprobe;    ///< number of probes at query time
+    size_t max_codes; ///< max nb of codes to visit to do a query
 
     /** Parallel mode determines how queries are parallelized with OpenMP
      *
@@ -130,9 +124,12 @@ struct IndexIVF: Index, Level1Quantizer {
      * identifier. The pointer is borrowed: the quantizer should not
      * be deleted while the IndexIVF is in use.
      */
-    IndexIVF (Index * quantizer, size_t d,
-              size_t nlist, size_t code_size,
-              MetricType metric = METRIC_L2);
+    IndexIVF(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            size_t code_size,
+            MetricType metric = METRIC_L2);
 
     void reset() override;
 
@@ -145,6 +142,19 @@ struct IndexIVF: Index, Level1Quantizer {
     /// default implementation that calls encode_vectors
     void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
 
+    /** Implementation of vector addition where the vector assignments are
+     * predefined. The default implementation hands over the code extraction to
+     * encode_vectors.
+     *
+     * @param precomputed_idx    quantization indices for the input vectors
+     * (size n)
+     */
+    virtual void add_core(
+            idx_t n,
+            const float* x,
+            const idx_t* xids,
+            const idx_t* precomputed_idx);
+
     /** Encodes a set of vectors as they would appear in the inverted lists
      *
      * @param list_nos   inverted list ids as returned by the
@@ -154,14 +164,16 @@ struct IndexIVF: Index, Level1Quantizer {
      *                   include the list ids in the code (in this case add
      *                   ceil(log8(nlist)) to the code size)
      */
-    virtual void encode_vectors(idx_t n, const float* x,
-                                const idx_t *list_nos,
-                                uint8_t * codes,
-                                bool include_listno = false) const = 0;
+    virtual void encode_vectors(
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            uint8_t* codes,
+            bool include_listno = false) const = 0;
 
     /// Sub-classes that encode the residuals can train their encoders here
     /// does nothing by default
-    virtual void train_residual (idx_t n, const float *x);
+    virtual void train_residual(idx_t n, const float* x);
 
     /** search a set of vectors, that are pre-quantized by the IVF
      *  quantizer. Fill in the corresponding heaps with the query
@@ -182,36 +194,50 @@ struct IndexIVF: Index, Level1Quantizer {
      * @param params used to override the object's search parameters
      * @param stats  search stats to be updated (can be null)
      */
-    virtual void search_preassigned (
-            idx_t n, const float *x, idx_t k,
-            const idx_t *assign, const float *centroid_dis,
-            float *distances, idx_t *labels,
+    virtual void search_preassigned(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            const idx_t* assign,
+            const float* centroid_dis,
+            float* distances,
+            idx_t* labels,
             bool store_pairs,
-            const IVFSearchParameters *params=nullptr,
-            IndexIVFStats *stats=nullptr
-            ) const;
+            const IVFSearchParameters* params = nullptr,
+            IndexIVFStats* stats = nullptr) const;
 
     /** assign the vectors, then call search_preassign */
-    void search (idx_t n, const float *x, idx_t k,
-                 float *distances, idx_t *labels) const override;
-
-    void range_search (idx_t n, const float* x, float radius,
-                       RangeSearchResult* result) const override;
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
+
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result) const override;
 
     void range_search_preassigned(
-            idx_t nx, const float *x, float radius,
-            const idx_t *keys, const float *coarse_dis,
-            RangeSearchResult *result,
-            bool store_pairs=false,
-            const IVFSearchParameters *params=nullptr,
-            IndexIVFStats *stats=nullptr) const;
+            idx_t nx,
+            const float* x,
+            float radius,
+            const idx_t* keys,
+            const float* coarse_dis,
+            RangeSearchResult* result,
+            bool store_pairs = false,
+            const IVFSearchParameters* params = nullptr,
+            IndexIVFStats* stats = nullptr) const;
 
     /// get a scanner for this index (store_pairs means ignore labels)
-    virtual InvertedListScanner *get_InvertedListScanner (
-        bool store_pairs=false) const;
+    virtual InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs = false) const;
 
-    /** reconstruct a vector. Works only if maintain_direct_map is set to 1 or 2 */
-    void reconstruct (idx_t key, float* recons) const override;
+    /** reconstruct a vector. Works only if maintain_direct_map is set to 1 or 2
+     */
+    void reconstruct(idx_t key, float* recons) const override;
 
     /** Update a subset of vectors.
      *
@@ -221,7 +247,7 @@ struct IndexIVF: Index, Level1Quantizer {
      * @param idx    vector indices to update, size nv
      * @param v      vectors of new values, size nv*d
      */
-    virtual void update_vectors (int nv, const idx_t *idx, const float *v);
+    virtual void update_vectors(int nv, const idx_t* idx, const float* v);
 
     /** Reconstruct a subset of the indexed vectors.
      *
@@ -243,9 +269,13 @@ struct IndexIVF: Index, Level1Quantizer {
      *
      * @param recons      reconstructed vectors size (n, k, d)
      */
-    void search_and_reconstruct (idx_t n, const float *x, idx_t k,
-                                 float *distances, idx_t *labels,
-                                 float *recons) const override;
+    void search_and_reconstruct(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            float* recons) const override;
 
     /** Reconstruct a vector given the location in terms of (inv list index +
      * inv list offset) instead of the id.
@@ -254,9 +284,10 @@ struct IndexIVF: Index, Level1Quantizer {
      * the inv list offset is computed by search_preassigned() with
      * `store_pairs` set.
      */
-    virtual void reconstruct_from_offset (int64_t list_no, int64_t offset,
-                                          float* recons) const;
-
+    virtual void reconstruct_from_offset(
+            int64_t list_no,
+            int64_t offset,
+            float* recons) const;
 
     /// Dataset manipulation functions
 
@@ -265,12 +296,12 @@ struct IndexIVF: Index, Level1Quantizer {
     /** check that the two indexes are compatible (ie, they are
      * trained in the same way and have the same
      * parameters). Otherwise throw. */
-    void check_compatible_for_merge (const IndexIVF &other) const;
+    void check_compatible_for_merge(const IndexIVF& other) const;
 
     /** moves the entries from another dataset to self. On output,
      * other is empty. add_id is added to all moved ids (for
      * sequential ids, this would be this->ntotal */
-    virtual void merge_from (IndexIVF &other, idx_t add_id);
+    virtual void merge_from(IndexIVF& other, idx_t add_id);
 
     /** copy a subset of the entries index to the other index
      *
@@ -279,34 +310,36 @@ struct IndexIVF: Index, Level1Quantizer {
      * if subset_type == 2: copies inverted lists such that a1
      *                      elements are left before and a2 elements are after
      */
-    virtual void copy_subset_to (IndexIVF & other, int subset_type,
-                                 idx_t a1, idx_t a2) const;
+    virtual void copy_subset_to(
+            IndexIVF& other,
+            int subset_type,
+            idx_t a1,
+            idx_t a2) const;
 
     ~IndexIVF() override;
 
-    size_t get_list_size (size_t list_no) const
-    { return invlists->list_size(list_no); }
+    size_t get_list_size(size_t list_no) const {
+        return invlists->list_size(list_no);
+    }
 
     /** intialize a direct map
      *
      * @param new_maintain_direct_map    if true, create a direct map,
      *                                   else clear it
      */
-    void make_direct_map (bool new_maintain_direct_map=true);
-
-    void set_direct_map_type (DirectMap::Type type);
+    void make_direct_map(bool new_maintain_direct_map = true);
 
+    void set_direct_map_type(DirectMap::Type type);
 
     /// replace the inverted lists, old one is deallocated if own_invlists
-    void replace_invlists (InvertedLists *il, bool own=false);
+    void replace_invlists(InvertedLists* il, bool own = false);
 
     /* The standalone codec interface (except sa_decode that is specific) */
-    size_t sa_code_size () const override;
+    size_t sa_code_size() const override;
 
-    void sa_encode (idx_t n, const float *x,
-                          uint8_t *bytes) const override;
+    void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
 
-    IndexIVF ();
+    IndexIVF();
 };
 
 struct RangeQueryResult;
@@ -316,17 +349,16 @@ struct RangeQueryResult;
  * distance_to_code and scan_codes can be called in multiple
  * threads */
 struct InvertedListScanner {
-
     using idx_t = Index::idx_t;
 
     /// from now on we handle this query.
-    virtual void set_query (const float *query_vector) = 0;
+    virtual void set_query(const float* query_vector) = 0;
 
     /// following codes come from this inverted list
-    virtual void set_list (idx_t list_no, float coarse_dis) = 0;
+    virtual void set_list(idx_t list_no, float coarse_dis) = 0;
 
     /// compute a single query-to-code distance
-    virtual float distance_to_code (const uint8_t *code) const = 0;
+    virtual float distance_to_code(const uint8_t* code) const = 0;
 
     /** scan a set of codes, compute distances to current query and
      * update heap of results if necessary.
@@ -339,45 +371,46 @@ struct InvertedListScanner {
      * @param k          heap size
      * @return number of heap updates performed
      */
-    virtual size_t scan_codes (size_t n,
-                               const uint8_t *codes,
-                               const idx_t *ids,
-                               float *distances, idx_t *labels,
-                               size_t k) const = 0;
+    virtual size_t scan_codes(
+            size_t n,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float* distances,
+            idx_t* labels,
+            size_t k) const = 0;
 
     /** scan a set of codes, compute distances to current query and
      * update results if distances are below radius
      *
      * (default implementation fails) */
-    virtual void scan_codes_range (size_t n,
-                                   const uint8_t *codes,
-                                   const idx_t *ids,
-                                   float radius,
-                                   RangeQueryResult &result) const;
-
-    virtual ~InvertedListScanner () {}
-
+    virtual void scan_codes_range(
+            size_t n,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float radius,
+            RangeQueryResult& result) const;
+
+    virtual ~InvertedListScanner() {}
 };
 
-
 struct IndexIVFStats {
-    size_t nq;       // nb of queries run
-    size_t nlist;    // nb of inverted lists scanned
-    size_t ndis;     // nb of distancs computed
-    size_t nheap_updates; // nb of times the heap was updated
+    size_t nq;                // nb of queries run
+    size_t nlist;             // nb of inverted lists scanned
+    size_t ndis;              // nb of distancs computed
+    size_t nheap_updates;     // nb of times the heap was updated
     double quantization_time; // time spent quantizing vectors (in ms)
     double search_time;       // time spent searching lists (in ms)
 
-    IndexIVFStats () {reset (); }
-    void reset ();
-    void add (const IndexIVFStats & other);
+    IndexIVFStats() {
+        reset();
+    }
+    void reset();
+    void add(const IndexIVFStats& other);
 };
 
 // global var that collects them all
 FAISS_API extern IndexIVFStats indexIVF_stats;
 
-
 } // namespace faiss
 
-
 #endif

data/vendor/faiss/faiss/IndexIVFFlat.cpp

@@ -9,332 +9,340 @@
 
 #include <faiss/IndexIVFFlat.h>
 
+#include <omp.h>
+
 #include <cinttypes>
 #include <cstdio>
 
 #include <faiss/IndexFlat.h>
 
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/utils.h>
-#include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/AuxIndexStructures.h>
-
 
 namespace faiss {
 
-
 /*****************************************
  * IndexIVFFlat implementation
  ******************************************/
 
-IndexIVFFlat::IndexIVFFlat (Index * quantizer,
-                            size_t d, size_t nlist, MetricType metric):
-    IndexIVF (quantizer, d, nlist, sizeof(float) * d, metric)
-{
+IndexIVFFlat::IndexIVFFlat(
+        Index* quantizer,
+        size_t d,
+        size_t nlist,
+        MetricType metric)
+        : IndexIVF(quantizer, d, nlist, sizeof(float) * d, metric) {
     code_size = sizeof(float) * d;
 }
 
+void IndexIVFFlat::add_core(
+        idx_t n,
+        const float* x,
+        const int64_t* xids,
+        const int64_t* coarse_idx)
 
-void IndexIVFFlat::add_with_ids (idx_t n, const float * x, const idx_t *xids)
 {
-    add_core (n, x, xids, nullptr);
-}
-
-void IndexIVFFlat::add_core (idx_t n, const float * x, const int64_t *xids,
-                             const int64_t *precomputed_idx)
+    FAISS_THROW_IF_NOT(is_trained);
+    FAISS_THROW_IF_NOT(coarse_idx);
+    assert(invlists);
+    direct_map.check_can_add(xids);
 
-{
-    FAISS_THROW_IF_NOT (is_trained);
-    assert (invlists);
-    direct_map.check_can_add (xids);
-    const int64_t * idx;
-    ScopeDeleter<int64_t> del;
-
-    if (precomputed_idx) {
-        idx = precomputed_idx;
-    } else {
-        int64_t * idx0 = new int64_t [n];
-        del.set (idx0);
-        quantizer->assign (n, x, idx0);
-        idx = idx0;
-    }
     int64_t n_add = 0;
-    for (size_t i = 0; i < n; i++) {
-        idx_t id = xids ? xids[i] : ntotal + i;
-        idx_t list_no = idx [i];
-        size_t offset;
-
-        if (list_no >= 0) {
-            const float *xi = x + i * d;
-            offset = invlists->add_entry (
-                     list_no, id, (const uint8_t*) xi);
-            n_add++;
-        } else {
-            offset = 0;
+
+    DirectMapAdd dm_adder(direct_map, n, xids);
+
+#pragma omp parallel reduction(+ : n_add)
+    {
+        int nt = omp_get_num_threads();
+        int rank = omp_get_thread_num();
+
+        // each thread takes care of a subset of lists
+        for (size_t i = 0; i < n; i++) {
+            idx_t list_no = coarse_idx[i];
+
+            if (list_no >= 0 && list_no % nt == rank) {
+                idx_t id = xids ? xids[i] : ntotal + i;
+                const float* xi = x + i * d;
+                size_t offset =
+                        invlists->add_entry(list_no, id, (const uint8_t*)xi);
+                dm_adder.add(i, list_no, offset);
+                n_add++;
+            } else if (rank == 0 && list_no == -1) {
+                dm_adder.add(i, -1, 0);
+            }
         }
-        direct_map.add_single_id (id, list_no, offset);
     }
 
     if (verbose) {
-        printf("IndexIVFFlat::add_core: added %" PRId64 " / %" PRId64 " vectors\n",
-               n_add, n);
+        printf("IndexIVFFlat::add_core: added %" PRId64 " / %" PRId64
+               " vectors\n",
+               n_add,
+               n);
     }
     ntotal += n;
 }
 
-void IndexIVFFlat::encode_vectors(idx_t n, const float* x,
-                                  const idx_t * list_nos,
-                                  uint8_t * codes,
-                                  bool include_listnos) const
-{
+void IndexIVFFlat::encode_vectors(
+        idx_t n,
+        const float* x,
+        const idx_t* list_nos,
+        uint8_t* codes,
+        bool include_listnos) const {
     if (!include_listnos) {
-        memcpy (codes, x, code_size * n);
+        memcpy(codes, x, code_size * n);
     } else {
-        size_t coarse_size = coarse_code_size ();
+        size_t coarse_size = coarse_code_size();
         for (size_t i = 0; i < n; i++) {
-            int64_t list_no = list_nos [i];
-            uint8_t *code = codes + i * (code_size + coarse_size);
-            const float *xi = x + i * d;
+            int64_t list_no = list_nos[i];
+            uint8_t* code = codes + i * (code_size + coarse_size);
+            const float* xi = x + i * d;
             if (list_no >= 0) {
-                encode_listno (list_no, code);
-                memcpy (code + coarse_size, xi, code_size);
+                encode_listno(list_no, code);
+                memcpy(code + coarse_size, xi, code_size);
             } else {
-                memset (code, 0, code_size + coarse_size);
+                memset(code, 0, code_size + coarse_size);
             }
-
         }
     }
 }
 
-void IndexIVFFlat::sa_decode (idx_t n, const uint8_t *bytes,
-                                      float *x) const
-{
-    size_t coarse_size = coarse_code_size ();
+void IndexIVFFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
+    size_t coarse_size = coarse_code_size();
     for (size_t i = 0; i < n; i++) {
-        const uint8_t *code = bytes + i * (code_size + coarse_size);
-        float *xi = x + i * d;
-        memcpy (xi, code + coarse_size, code_size);
+        const uint8_t* code = bytes + i * (code_size + coarse_size);
+        float* xi = x + i * d;
+        memcpy(xi, code + coarse_size, code_size);
     }
 }
 
-
 namespace {
 
-
-template<MetricType metric, class C>
-struct IVFFlatScanner: InvertedListScanner {
+template <MetricType metric, class C>
+struct IVFFlatScanner : InvertedListScanner {
     size_t d;
     bool store_pairs;
 
-    IVFFlatScanner(size_t d, bool store_pairs):
-        d(d), store_pairs(store_pairs) {}
+    IVFFlatScanner(size_t d, bool store_pairs)
+            : d(d), store_pairs(store_pairs) {}
 
-    const float *xi;
-    void set_query (const float *query) override {
+    const float* xi;
+    void set_query(const float* query) override {
         this->xi = query;
     }
 
     idx_t list_no;
-    void set_list (idx_t list_no, float /* coarse_dis */) override {
+    void set_list(idx_t list_no, float /* coarse_dis */) override {
         this->list_no = list_no;
     }
 
-    float distance_to_code (const uint8_t *code) const override {
-        const float *yj = (float*)code;
-        float dis = metric == METRIC_INNER_PRODUCT ?
-            fvec_inner_product (xi, yj, d) : fvec_L2sqr (xi, yj, d);
+    float distance_to_code(const uint8_t* code) const override {
+        const float* yj = (float*)code;
+        float dis = metric == METRIC_INNER_PRODUCT
+                ? fvec_inner_product(xi, yj, d)
+                : fvec_L2sqr(xi, yj, d);
         return dis;
     }
 
-    size_t scan_codes (size_t list_size,
-                       const uint8_t *codes,
-                       const idx_t *ids,
-                       float *simi, idx_t *idxi,
-                       size_t k) const override
-    {
-        const float *list_vecs = (const float*)codes;
+    size_t scan_codes(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float* simi,
+            idx_t* idxi,
+            size_t k) const override {
+        const float* list_vecs = (const float*)codes;
         size_t nup = 0;
         for (size_t j = 0; j < list_size; j++) {
-            const float * yj = list_vecs + d * j;
-            float dis = metric == METRIC_INNER_PRODUCT ?
-                fvec_inner_product (xi, yj, d) : fvec_L2sqr (xi, yj, d);
-            if (C::cmp (simi[0], dis)) {
-                int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
-                heap_replace_top<C> (k, simi, idxi, dis, id);
+            const float* yj = list_vecs + d * j;
+            float dis = metric == METRIC_INNER_PRODUCT
+                    ? fvec_inner_product(xi, yj, d)
+                    : fvec_L2sqr(xi, yj, d);
+            if (C::cmp(simi[0], dis)) {
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                heap_replace_top<C>(k, simi, idxi, dis, id);
                 nup++;
             }
         }
         return nup;
     }
 
-    void scan_codes_range (size_t list_size,
-                           const uint8_t *codes,
-                           const idx_t *ids,
-                           float radius,
-                           RangeQueryResult & res) const override
-    {
-        const float *list_vecs = (const float*)codes;
+    void scan_codes_range(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float radius,
+            RangeQueryResult& res) const override {
+        const float* list_vecs = (const float*)codes;
         for (size_t j = 0; j < list_size; j++) {
-            const float * yj = list_vecs + d * j;
-            float dis = metric == METRIC_INNER_PRODUCT ?
-                fvec_inner_product (xi, yj, d) : fvec_L2sqr (xi, yj, d);
-            if (C::cmp (radius, dis)) {
-                int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
-                res.add (dis, id);
+            const float* yj = list_vecs + d * j;
+            float dis = metric == METRIC_INNER_PRODUCT
+                    ? fvec_inner_product(xi, yj, d)
+                    : fvec_L2sqr(xi, yj, d);
+            if (C::cmp(radius, dis)) {
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                res.add(dis, id);
             }
         }
     }
-
-
 };
 
-
 } // anonymous namespace
 
-
-
-InvertedListScanner* IndexIVFFlat::get_InvertedListScanner
-     (bool store_pairs) const
-{
+InvertedListScanner* IndexIVFFlat::get_InvertedListScanner(
+        bool store_pairs) const {
     if (metric_type == METRIC_INNER_PRODUCT) {
-        return new IVFFlatScanner<
-            METRIC_INNER_PRODUCT, CMin<float, int64_t> > (d, store_pairs);
+        return new IVFFlatScanner<METRIC_INNER_PRODUCT, CMin<float, int64_t>>(
+                d, store_pairs);
     } else if (metric_type == METRIC_L2) {
-        return new IVFFlatScanner<
-            METRIC_L2, CMax<float, int64_t> >(d, store_pairs);
+        return new IVFFlatScanner<METRIC_L2, CMax<float, int64_t>>(
+                d, store_pairs);
     } else {
         FAISS_THROW_MSG("metric type not supported");
     }
     return nullptr;
 }
 
-
-
-
-void IndexIVFFlat::reconstruct_from_offset (int64_t list_no, int64_t offset,
-                                            float* recons) const
-{
-    memcpy (recons, invlists->get_single_code (list_no, offset), code_size);
+void IndexIVFFlat::reconstruct_from_offset(
+        int64_t list_no,
+        int64_t offset,
+        float* recons) const {
+    memcpy(recons, invlists->get_single_code(list_no, offset), code_size);
 }
 
 /*****************************************
  * IndexIVFFlatDedup implementation
  ******************************************/
 
-IndexIVFFlatDedup::IndexIVFFlatDedup (
-            Index * quantizer, size_t d, size_t nlist_,
-            MetricType metric_type):
-    IndexIVFFlat (quantizer, d, nlist_, metric_type)
-{}
-
+IndexIVFFlatDedup::IndexIVFFlatDedup(
+        Index* quantizer,
+        size_t d,
+        size_t nlist_,
+        MetricType metric_type)
+        : IndexIVFFlat(quantizer, d, nlist_, metric_type) {}
 
-void IndexIVFFlatDedup::train(idx_t n, const float* x)
-{
+void IndexIVFFlatDedup::train(idx_t n, const float* x) {
     std::unordered_map<uint64_t, idx_t> map;
-    float * x2 = new float [n * d];
-    ScopeDeleter<float> del (x2);
+    float* x2 = new float[n * d];
+    ScopeDeleter<float> del(x2);
 
     int64_t n2 = 0;
     for (int64_t i = 0; i < n; i++) {
-        uint64_t hash = hash_bytes((uint8_t *)(x + i * d), code_size);
+        uint64_t hash = hash_bytes((uint8_t*)(x + i * d), code_size);
         if (map.count(hash) &&
-            !memcmp (x2 + map[hash] * d, x + i * d, code_size)) {
+            !memcmp(x2 + map[hash] * d, x + i * d, code_size)) {
             // is duplicate, skip
         } else {
-            map [hash] = n2;
-            memcpy (x2 + n2 * d, x + i * d, code_size);
-            n2 ++;
+            map[hash] = n2;
+            memcpy(x2 + n2 * d, x + i * d, code_size);
+            n2++;
         }
     }
     if (verbose) {
-        printf ("IndexIVFFlatDedup::train: train on %" PRId64 " points after dedup "
-                "(was %" PRId64 " points)\n", n2, n);
+        printf("IndexIVFFlatDedup::train: train on %" PRId64
+               " points after dedup "
+               "(was %" PRId64 " points)\n",
+               n2,
+               n);
     }
-    IndexIVFFlat::train (n2, x2);
+    IndexIVFFlat::train(n2, x2);
 }
 
+void IndexIVFFlatDedup::add_with_ids(
+        idx_t na,
+        const float* x,
+        const idx_t* xids) {
+    FAISS_THROW_IF_NOT(is_trained);
+    assert(invlists);
+    FAISS_THROW_IF_NOT_MSG(
+            direct_map.no(), "IVFFlatDedup not implemented with direct_map");
+    int64_t* idx = new int64_t[na];
+    ScopeDeleter<int64_t> del(idx);
+    quantizer->assign(na, x, idx);
 
+    int64_t n_add = 0, n_dup = 0;
 
-void IndexIVFFlatDedup::add_with_ids(
-           idx_t na, const float* x, const idx_t* xids)
-{
+#pragma omp parallel reduction(+ : n_add, n_dup)
+    {
+        int nt = omp_get_num_threads();
+        int rank = omp_get_thread_num();
 
-    FAISS_THROW_IF_NOT (is_trained);
-    assert (invlists);
-    FAISS_THROW_IF_NOT_MSG (direct_map.no(),
-           "IVFFlatDedup not implemented with direct_map");
-    int64_t * idx = new int64_t [na];
-    ScopeDeleter<int64_t> del (idx);
-    quantizer->assign (na, x, idx);
+        // each thread takes care of a subset of lists
+        for (size_t i = 0; i < na; i++) {
+            int64_t list_no = idx[i];
 
-    int64_t n_add = 0, n_dup = 0;
-    // TODO make a omp loop with this
-    for (size_t i = 0; i < na; i++) {
-        idx_t id = xids ? xids[i] : ntotal + i;
-        int64_t list_no = idx [i];
+            if (list_no < 0 || list_no % nt != rank) {
+                continue;
+            }
 
-        if (list_no < 0) {
-            continue;
-        }
-        const float *xi = x + i * d;
+            idx_t id = xids ? xids[i] : ntotal + i;
+            const float* xi = x + i * d;
 
-        // search if there is already an entry with that id
-        InvertedLists::ScopedCodes codes (invlists, list_no);
+            // search if there is already an entry with that id
+            InvertedLists::ScopedCodes codes(invlists, list_no);
 
-        int64_t n = invlists->list_size (list_no);
-        int64_t offset = -1;
-        for (int64_t o = 0; o < n; o++) {
-            if (!memcmp (codes.get() + o * code_size,
-                         xi, code_size)) {
-                offset = o;
-                break;
+            int64_t n = invlists->list_size(list_no);
+            int64_t offset = -1;
+            for (int64_t o = 0; o < n; o++) {
+                if (!memcmp(codes.get() + o * code_size, xi, code_size)) {
+                    offset = o;
+                    break;
+                }
             }
-        }
 
-        if (offset == -1) { // not found
-            invlists->add_entry (list_no, id, (const uint8_t*) xi);
-        } else {
-            // mark equivalence
-            idx_t id2 = invlists->get_single_id (list_no, offset);
-            std::pair<idx_t, idx_t> pair (id2, id);
-            instances.insert (pair);
-            n_dup ++;
+            if (offset == -1) { // not found
+                invlists->add_entry(list_no, id, (const uint8_t*)xi);
+            } else {
+                // mark equivalence
+                idx_t id2 = invlists->get_single_id(list_no, offset);
+                std::pair<idx_t, idx_t> pair(id2, id);
+
+#pragma omp critical
+                // executed by one thread at a time
+                instances.insert(pair);
+
+                n_dup++;
+            }
+            n_add++;
         }
-        n_add++;
     }
     if (verbose) {
-        printf("IndexIVFFlat::add_with_ids: added %" PRId64 " / %" PRId64 " vectors"
+        printf("IndexIVFFlat::add_with_ids: added %" PRId64 " / %" PRId64
+               " vectors"
                " (out of which %" PRId64 " are duplicates)\n",
-               n_add, na, n_dup);
+               n_add,
+               na,
+               n_dup);
     }
     ntotal += n_add;
 }
 
-void IndexIVFFlatDedup::search_preassigned (
-           idx_t n, const float *x, idx_t k,
-           const idx_t *assign,
-           const float *centroid_dis,
-           float *distances, idx_t *labels,
-           bool store_pairs,
-           const IVFSearchParameters *params,
-           IndexIVFStats *stats) const
-{
-    FAISS_THROW_IF_NOT_MSG (
-           !store_pairs, "store_pairs not supported in IVFDedup");
-
-    IndexIVFFlat::search_preassigned (n, x, k, assign, centroid_dis,
-                                      distances, labels, false,
-                                      params);
-
-    std::vector <idx_t> labels2 (k);
-    std::vector <float> dis2 (k);
+void IndexIVFFlatDedup::search_preassigned(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        const idx_t* assign,
+        const float* centroid_dis,
+        float* distances,
+        idx_t* labels,
+        bool store_pairs,
+        const IVFSearchParameters* params,
+        IndexIVFStats* stats) const {
+    FAISS_THROW_IF_NOT_MSG(
+            !store_pairs, "store_pairs not supported in IVFDedup");
+
+    IndexIVFFlat::search_preassigned(
+            n, x, k, assign, centroid_dis, distances, labels, false, params);
+
+    std::vector<idx_t> labels2(k);
+    std::vector<float> dis2(k);
 
     for (int64_t i = 0; i < n; i++) {
-        idx_t *labels1 = labels + i * k;
-        float *dis1 = distances + i * k;
+        idx_t* labels1 = labels + i * k;
+        float* dis1 = distances + i * k;
         int64_t j = 0;
         for (; j < k; j++) {
-            if (instances.find (labels1[j]) != instances.end ()) {
+            if (instances.find(labels1[j]) != instances.end()) {
                 // a duplicate: special handling
                 break;
             }
@@ -344,11 +352,11 @@ void IndexIVFFlatDedup::search_preassigned (
             int64_t j0 = j;
             int64_t rp = j;
             while (j < k) {
-                auto range = instances.equal_range (labels1[rp]);
+                auto range = instances.equal_range(labels1[rp]);
                 float dis = dis1[rp];
                 labels2[j] = labels1[rp];
                 dis2[j] = dis;
-                j ++;
+                j++;
                 for (auto it = range.first; j < k && it != range.second; ++it) {
                     labels2[j] = it->second;
                     dis2[j] = dis;
@@ -356,21 +364,18 @@ void IndexIVFFlatDedup::search_preassigned (
                 }
                 rp++;
             }
-            memcpy (labels1 + j0, labels2.data() + j0,
-                    sizeof(labels1[0]) * (k - j0));
-            memcpy (dis1 + j0, dis2.data() + j0,
-                    sizeof(dis2[0]) * (k - j0));
+            memcpy(labels1 + j0,
+                   labels2.data() + j0,
+                   sizeof(labels1[0]) * (k - j0));
+            memcpy(dis1 + j0, dis2.data() + j0, sizeof(dis2[0]) * (k - j0));
         }
     }
-
 }
 
-
-size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
-{
+size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel) {
     std::unordered_map<idx_t, idx_t> replace;
-    std::vector<std::pair<idx_t, idx_t> > toadd;
-    for (auto it = instances.begin(); it != instances.end(); ) {
+    std::vector<std::pair<idx_t, idx_t>> toadd;
+    for (auto it = instances.begin(); it != instances.end();) {
         if (sel.is_member(it->first)) {
             // then we erase this entry
             if (!sel.is_member(it->second)) {
@@ -378,8 +383,8 @@ size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
                 if (replace.count(it->first) == 0) {
                     replace[it->first] = it->second;
                 } else { // remember we should add an element
-                    std::pair<idx_t, idx_t> new_entry (
-                          replace[it->first], it->second);
+                    std::pair<idx_t, idx_t> new_entry(
+                            replace[it->first], it->second);
                     toadd.push_back(new_entry);
                 }
             }
@@ -393,32 +398,34 @@ size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
         }
     }
 
-    instances.insert (toadd.begin(), toadd.end());
+    instances.insert(toadd.begin(), toadd.end());
 
     // mostly copied from IndexIVF.cpp
 
-    FAISS_THROW_IF_NOT_MSG (direct_map.no(),
-                    "direct map remove not implemented");
+    FAISS_THROW_IF_NOT_MSG(
+            direct_map.no(), "direct map remove not implemented");
 
     std::vector<int64_t> toremove(nlist);
 
 #pragma omp parallel for
     for (int64_t i = 0; i < nlist; i++) {
-        int64_t l0 = invlists->list_size (i), l = l0, j = 0;
-        InvertedLists::ScopedIds idsi (invlists, i);
+        int64_t l0 = invlists->list_size(i), l = l0, j = 0;
+        InvertedLists::ScopedIds idsi(invlists, i);
         while (j < l) {
-            if (sel.is_member (idsi[j])) {
+            if (sel.is_member(idsi[j])) {
                 if (replace.count(idsi[j]) == 0) {
                     l--;
-                    invlists->update_entry (
-                        i, j,
-                        invlists->get_single_id (i, l),
-                        InvertedLists::ScopedCodes (invlists, i, l).get());
+                    invlists->update_entry(
+                            i,
+                            j,
+                            invlists->get_single_id(i, l),
+                            InvertedLists::ScopedCodes(invlists, i, l).get());
                 } else {
-                    invlists->update_entry (
-                        i, j,
-                        replace[idsi[j]],
-                        InvertedLists::ScopedCodes (invlists, i, j).get());
+                    invlists->update_entry(
+                            i,
+                            j,
+                            replace[idsi[j]],
+                            InvertedLists::ScopedCodes(invlists, i, j).get());
                     j++;
                 }
             } else {
@@ -432,37 +439,28 @@ size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
     for (int64_t i = 0; i < nlist; i++) {
         if (toremove[i] > 0) {
             nremove += toremove[i];
-            invlists->resize(
-                i, invlists->list_size(i) - toremove[i]);
+            invlists->resize(i, invlists->list_size(i) - toremove[i]);
         }
     }
     ntotal -= nremove;
     return nremove;
 }
 
-
 void IndexIVFFlatDedup::range_search(
-        idx_t ,
-        const float* ,
-        float ,
-        RangeSearchResult* ) const
-{
-    FAISS_THROW_MSG ("not implemented");
+        idx_t,
+        const float*,
+        float,
+        RangeSearchResult*) const {
+    FAISS_THROW_MSG("not implemented");
 }
 
-void IndexIVFFlatDedup::update_vectors (int , const idx_t *, const float *)
-{
-    FAISS_THROW_MSG ("not implemented");
+void IndexIVFFlatDedup::update_vectors(int, const idx_t*, const float*) {
+    FAISS_THROW_MSG("not implemented");
 }
 
-
-void IndexIVFFlatDedup::reconstruct_from_offset (
-         int64_t , int64_t , float* ) const
-{
-    FAISS_THROW_MSG ("not implemented");
+void IndexIVFFlatDedup::reconstruct_from_offset(int64_t, int64_t, float*)
+        const {
+    FAISS_THROW_MSG("not implemented");
 }
 
-
-
-
 } // namespace faiss