IndexShards with common quantizer (#2682)

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

IndexShards normally sees the indexes as opaque, so there is no way to factrorize the coarse quantizer.
This diff introduces IndexIVFShards that handles IVF indexes with a common quantizer so that the quantization is computed only once.

Reviewed By: alexanderguzhva

Differential Revision: D42781513

fbshipit-source-id: 441316eff4c1ba0468501c456af9194ea5f042d6

c_api/IndexReplicas_c.cpp

@@ -14,8 +14,8 @@ using faiss::IndexReplicas;
 
 DEFINE_DESTRUCTOR(IndexReplicas)
 
-DEFINE_GETTER(IndexReplicas, int, own_fields)
-DEFINE_SETTER(IndexReplicas, int, own_fields)
+DEFINE_GETTER(IndexReplicas, int, own_indices)
+DEFINE_SETTER(IndexReplicas, int, own_indices)
 
 int faiss_IndexReplicas_new(FaissIndexReplicas** p_index, idx_t d) {
     try {

c_api/IndexShards_c.cpp

@@ -14,8 +14,8 @@ using faiss::IndexShards;
 
 DEFINE_DESTRUCTOR(IndexShards)
 
-DEFINE_GETTER(IndexShards, int, own_fields)
-DEFINE_SETTER(IndexShards, int, own_fields)
+DEFINE_GETTER(IndexShards, int, own_indices)
+DEFINE_SETTER(IndexShards, int, own_indices)
 
 DEFINE_GETTER(IndexShards, int, successive_ids)
 DEFINE_SETTER(IndexShards, int, successive_ids)

faiss/AutoTune.cpp

@@ -32,6 +32,7 @@
 #include <faiss/IndexPreTransform.h>
 #include <faiss/IndexRefine.h>
 #include <faiss/IndexScalarQuantizer.h>
+#include <faiss/IndexShardsIVF.h>
 #include <faiss/MetaIndexes.h>
 #include <faiss/VectorTransform.h>
 
@@ -354,7 +355,7 @@ void ParameterSpace::initialize(const Index* index) {
         index = ix->index;
     }
 
-    if (DC(IndexIVF)) {
+    if (DC(IndexIVFInterface)) {
         {
             ParameterRange& pr = add_range("nprobe");
             for (int i = 0; i < 13; i++) {
@@ -461,6 +462,16 @@ void ParameterSpace::set_index_parameter(
         set_index_parameter(ix->index, name, val);
         return;
     }
+    if (DC(IndexShardsIVF)) {
+        // special handling because the nprobe is set at the sub-class level
+        // but other params are set on the class itself
+        if (name.find("quantizer_") == 0 && name != "nprobe" &&
+            name != "quantizer_nprobe") {
+            std::string sub_name = name.substr(strlen("quantizer_"));
+            set_index_parameter(ix->quantizer, sub_name, val);
+            return;
+        }
+    }
     if (DC(ThreadedIndex<Index>)) {
         // call on all sub-indexes
         auto fn = [this, name, val](int /* no */, Index* subIndex) {

faiss/CMakeLists.txt

@@ -44,6 +44,7 @@ set(FAISS_SRC
   IndexRowwiseMinMax.cpp
   IndexScalarQuantizer.cpp
   IndexShards.cpp
+  IndexShardsIVF.cpp
   MatrixStats.cpp
   MetaIndexes.cpp
   VectorTransform.cpp
@@ -133,6 +134,7 @@ set(FAISS_HEADERS
   IndexRowwiseMinMax.h
   IndexScalarQuantizer.h
   IndexShards.h
+  IndexShardsIVF.h
   MatrixStats.h
   MetaIndexes.h
   MetricType.h

faiss/IndexIVF.cpp

@@ -160,7 +160,7 @@ IndexIVF::IndexIVF(
         size_t code_size,
         MetricType metric)
         : Index(d, metric),
-          Level1Quantizer(quantizer, nlist),
+          IndexIVFInterface(quantizer, nlist),
           invlists(new ArrayInvertedLists(nlist, code_size)),
           own_invlists(true),
           code_size(code_size) {
@@ -172,13 +172,7 @@ IndexIVF::IndexIVF(
     }
 }
 
-IndexIVF::IndexIVF()
-        : invlists(nullptr),
-          own_invlists(false),
-          code_size(0),
-          nprobe(1),
-          max_codes(0),
-          parallel_mode(0) {}
+IndexIVF::IndexIVF() {}
 
 void IndexIVF::add(idx_t n, const float* x) {
     add_with_ids(n, x, nullptr);

faiss/IndexIVF.h

@@ -83,6 +83,74 @@ struct InvertedListScanner;
 struct IndexIVFStats;
 struct CodePacker;
 
+struct IndexIVFInterface : Level1Quantizer {
+    size_t nprobe = 1;    ///< number of probes at query time
+    size_t max_codes = 0; ///< max nb of codes to visit to do a query
+
+    explicit IndexIVFInterface(Index* quantizer = nullptr, size_t nlist = 0)
+            : Level1Quantizer(quantizer, nlist) {}
+
+    /** search a set of vectors, that are pre-quantized by the IVF
+     *  quantizer. Fill in the corresponding heaps with the query
+     *  results. The default implementation uses InvertedListScanners
+     *  to do the search.
+     *
+     * @param n      nb of vectors to query
+     * @param x      query vectors, size nx * d
+     * @param assign coarse quantization indices, size nx * nprobe
+     * @param centroid_dis
+     *               distances to coarse centroids, size nx * nprobe
+     * @param distance
+     *               output distances, size n * k
+     * @param labels output labels, size n * k
+     * @param store_pairs store inv list index + inv list offset
+     *                     instead in upper/lower 32 bit of result,
+     *                     instead of ids (used for reranking).
+     * @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,
+            bool store_pairs,
+            const IVFSearchParameters* params = nullptr,
+            IndexIVFStats* stats = nullptr) const = 0;
+
+    /** Range search a set of vectors, that are pre-quantized by the IVF
+     *  quantizer. Fill in the RangeSearchResults results. The default
+     * implementation uses InvertedListScanners to do the search.
+     *
+     * @param n      nb of vectors to query
+     * @param x      query vectors, size nx * d
+     * @param assign coarse quantization indices, size nx * nprobe
+     * @param centroid_dis
+     *               distances to coarse centroids, size nx * nprobe
+     * @param result Output results
+     * @param store_pairs store inv list index + inv list offset
+     *                     instead in upper/lower 32 bit of result,
+     *                     instead of ids (used for reranking).
+     * @param params used to override the object's search parameters
+     * @param stats  search stats to be updated (can be null)
+     */
+    virtual 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 = 0;
+
+    virtual ~IndexIVFInterface() {}
+};
+
 /** Index based on a inverted file (IVF)
  *
  * In the inverted file, the quantizer (an Index instance) provides a
@@ -103,15 +171,12 @@ struct CodePacker;
  * 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, IndexIVFInterface {
     /// Access to the actual data
     InvertedLists* invlists = nullptr;
     bool own_invlists = false;
 
     size_t code_size = 0; ///< code size per vector in bytes
-    size_t nprobe = 1;    ///< number of probes at query time
-    size_t max_codes = 0; ///< max nb of codes to visit to do a query
-
     /** Parallel mode determines how queries are parallelized with OpenMP
      *
      * 0 (default): split over queries
@@ -191,26 +256,7 @@ struct IndexIVF : Index, Level1Quantizer {
     /// does nothing by default
     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
-     *  results. The default implementation uses InvertedListScanners
-     *  to do the search.
-     *
-     * @param n      nb of vectors to query
-     * @param x      query vectors, size nx * d
-     * @param assign coarse quantization indices, size nx * nprobe
-     * @param centroid_dis
-     *               distances to coarse centroids, size nx * nprobe
-     * @param distance
-     *               output distances, size n * k
-     * @param labels output labels, size n * k
-     * @param store_pairs store inv list index + inv list offset
-     *                     instead in upper/lower 32 bit of result,
-     *                     instead of ids (used for reranking).
-     * @param params used to override the object's search parameters
-     * @param stats  search stats to be updated (can be null)
-     */
-    virtual void search_preassigned(
+    void search_preassigned(
             idx_t n,
             const float* x,
             idx_t k,
@@ -220,7 +266,18 @@ struct IndexIVF : Index, Level1Quantizer {
             idx_t* labels,
             bool store_pairs,
             const IVFSearchParameters* params = nullptr,
-            IndexIVFStats* stats = nullptr) const;
+            IndexIVFStats* stats = nullptr) 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 override;
 
     /** assign the vectors, then call search_preassign */
     void search(
@@ -238,17 +295,6 @@ struct IndexIVF : Index, Level1Quantizer {
             RangeSearchResult* result,
             const SearchParameters* params = nullptr) 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;
-
     /** Get a scanner for this index (store_pairs means ignore labels)
      *
      * The default search implementation uses this to compute the distances

faiss/IndexShards.cpp

@@ -23,10 +23,10 @@ namespace faiss {
 namespace {
 
 // add translation to all valid labels
-void translate_labels(long n, idx_t* labels, long translation) {
+void translate_labels(int64_t n, idx_t* labels, int64_t translation) {
     if (translation == 0)
         return;
-    for (long i = 0; i < n; i++) {
+    for (int64_t i = 0; i < n; i++) {
         if (labels[i] < 0)
             continue;
         labels[i] += translation;
@@ -166,11 +166,9 @@ void IndexShardsTemplate<IndexT>::add_with_ids(
 
     if (!ids && !successive_ids) {
         aids.resize(n);
-
         for (idx_t i = 0; i < n; i++) {
             aids[i] = this->ntotal + i;
         }
-
         ids = aids.data();
     }
 
@@ -213,11 +211,11 @@ void IndexShardsTemplate<IndexT>::search(
             !params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
-    long nshard = this->count();
+    int64_t nshard = this->count();
 
     std::vector<distance_t> all_distances(nshard * k * n);
     std::vector<idx_t> all_labels(nshard * k * n);
-    std::vector<long> translations(nshard, 0);
+    std::vector<int64_t> translations(nshard, 0);
 
     // Because we just called runOnIndex above, it is safe to access the
     // sub-index ntotal here
@@ -277,4 +275,5 @@ void IndexShardsTemplate<IndexT>::search(
 template struct IndexShardsTemplate<Index>;
 template struct IndexShardsTemplate<IndexBinary>;
 
+
 } // namespace faiss

faiss/IndexShards.h

@@ -95,7 +95,7 @@ struct IndexShardsTemplate : public ThreadedIndex<IndexT> {
 
     /// Synchronize the top-level index (IndexShards) with data in the
     /// sub-indices
-    void syncWithSubIndexes();
+    virtual void syncWithSubIndexes();
 
    protected:
     /// Called just after an index is added

faiss/IndexShardsIVF.cpp

@@ -0,0 +1,246 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#include <faiss/IndexShardsIVF.h>
+
+#include <cinttypes>
+#include <cstdio>
+#include <functional>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/WorkerThread.h>
+#include <faiss/utils/utils.h>
+
+namespace faiss {
+
+// subroutines
+namespace {
+
+// add translation to all valid labels
+void translate_labels(int64_t n, idx_t* labels, int64_t translation) {
+    if (translation == 0) {
+        return;
+    }
+    for (int64_t i = 0; i < n; i++) {
+        if (labels[i] < 0) {
+            continue;
+        }
+        labels[i] += translation;
+    }
+}
+
+} // anonymous namespace
+
+/************************************************************
+ * IndexShardsIVF
+ ************************************************************/
+
+IndexShardsIVF::IndexShardsIVF(
+        Index* quantizer,
+        size_t nlist,
+        bool threaded,
+        bool successive_ids)
+        : IndexShardsTemplate<Index>(quantizer->d, threaded, successive_ids),
+          Level1Quantizer(quantizer, nlist) {
+    is_trained = quantizer->is_trained && quantizer->ntotal == nlist;
+}
+
+void IndexShardsIVF::addIndex(Index* index) {
+    auto index_ivf = dynamic_cast<IndexIVFInterface*>(index);
+    FAISS_THROW_IF_NOT_MSG(index_ivf, "can only add IndexIVFs");
+    FAISS_THROW_IF_NOT(index_ivf->nlist == nlist);
+    IndexShardsTemplate<Index>::addIndex(index);
+}
+
+void IndexShardsIVF::train(idx_t n, const component_t* x) {
+    if (verbose) {
+        printf("Training level-1 quantizer\n");
+    }
+    train_q1(n, x, verbose, metric_type);
+
+    // set the sub-quantizer codebooks
+    std::vector<float> centroids(nlist * d);
+    quantizer->reconstruct_n(0, nlist, centroids.data());
+
+    // probably not worth running in parallel
+    for (size_t i = 0; i < indices_.size(); i++) {
+        Index* index = indices_[i].first;
+        auto index_ivf = dynamic_cast<IndexIVFInterface*>(index);
+        Index* quantizer = index_ivf->quantizer;
+        if (!quantizer->is_trained) {
+            quantizer->train(nlist, centroids.data());
+        }
+        quantizer->add(nlist, centroids.data());
+        // finish training
+        index->train(n, x);
+    }
+
+    is_trained = true;
+}
+
+void IndexShardsIVF::add_with_ids(
+        idx_t n,
+        const component_t* x,
+        const idx_t* xids) {
+    // IndexIVF exposes add_core that we can use to factorize the
+    bool all_index_ivf = true;
+    for (size_t i = 0; i < indices_.size(); i++) {
+        Index* index = indices_[i].first;
+        all_index_ivf = all_index_ivf && dynamic_cast<IndexIVF*>(index);
+    }
+    if (!all_index_ivf) {
+        IndexShardsTemplate<Index>::add_with_ids(n, x, xids);
+        return;
+    }
+    FAISS_THROW_IF_NOT_MSG(
+            !(successive_ids && xids),
+            "It makes no sense to pass in ids and "
+            "request them to be shifted");
+
+    if (successive_ids) {
+        FAISS_THROW_IF_NOT_MSG(
+                !xids,
+                "It makes no sense to pass in ids and "
+                "request them to be shifted");
+        FAISS_THROW_IF_NOT_MSG(
+                this->ntotal == 0,
+                "when adding to IndexShards with sucessive_ids, "
+                "only add() in a single pass is supported");
+    }
+
+    // perform coarse quantization
+    std::vector<idx_t> Iq(n);
+    std::vector<float> Dq(n);
+    quantizer->search(n, x, 1, Dq.data(), Iq.data());
+
+    // possibly shift ids
+    idx_t nshard = this->count();
+    const idx_t* ids = xids;
+    std::vector<idx_t> aids;
+    if (!ids && !successive_ids) {
+        aids.resize(n);
+
+        for (idx_t i = 0; i < n; i++) {
+            aids[i] = this->ntotal + i;
+        }
+        ids = aids.data();
+    }
+    idx_t d = this->d;
+
+    auto fn = [n, ids, x, nshard, d, Iq](int no, Index* index) {
+        idx_t i0 = (idx_t)no * n / nshard;
+        idx_t i1 = ((idx_t)no + 1) * n / nshard;
+        const float* x0 = x + i0 * d;
+        auto index_ivf = dynamic_cast<IndexIVF*>(index);
+
+        if (index->verbose) {
+            printf("begin add shard %d on %" PRId64 " points\n", no, n);
+        }
+
+        index_ivf->add_core(
+                i1 - i0, x + i0 * d, ids ? ids + i0 : nullptr, Iq.data() + i0);
+
+        if (index->verbose) {
+            printf("end add shard %d on %" PRId64 " points\n", no, i1 - i0);
+        }
+    };
+
+    this->runOnIndex(fn);
+    syncWithSubIndexes();
+}
+
+void IndexShardsIVF::search(
+        idx_t n,
+        const component_t* x,
+        idx_t k,
+        distance_t* distances,
+        idx_t* labels,
+        const SearchParameters* params_in) const {
+    FAISS_THROW_IF_NOT(k > 0);
+    FAISS_THROW_IF_NOT(count() > 0);
+    const IVFSearchParameters* params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
+    }
+
+    auto index0 = dynamic_cast<const IndexIVFInterface*>(at(0));
+    idx_t nprobe = params ? params->nprobe : index0->nprobe;
+
+    // coarse quantization (TODO: support tiling with search_precomputed)
+    std::vector<distance_t> Dq(n * nprobe);
+    std::vector<idx_t> Iq(n * nprobe);
+
+    quantizer->search(n, x, nprobe, Dq.data(), Iq.data());
+
+    int64_t nshard = this->count();
+
+    std::vector<distance_t> all_distances(nshard * k * n);
+    std::vector<idx_t> all_labels(nshard * k * n);
+    std::vector<int64_t> translations(nshard, 0);
+
+    if (successive_ids) {
+        translations[0] = 0;
+        for (int s = 0; s + 1 < nshard; s++) {
+            translations[s + 1] = translations[s] + this->at(s)->ntotal;
+        }
+    }
+
+    auto fn = [&](int no, const Index* indexIn) {
+        if (indexIn->verbose) {
+            printf("begin query shard %d on %" PRId64 " points\n", no, n);
+        }
+
+        auto index = dynamic_cast<const IndexIVFInterface*>(indexIn);
+
+        FAISS_THROW_IF_NOT_MSG(index->nprobe == nprobe, "inconsistent nprobe");
+
+        index->search_preassigned(
+                n,
+                x,
+                k,
+                Iq.data(),
+                Dq.data(),
+                all_distances.data() + no * k * n,
+                all_labels.data() + no * k * n,
+                false);
+
+        translate_labels(
+                n * k, all_labels.data() + no * k * n, translations[no]);
+
+        if (indexIn->verbose) {
+            printf("end query shard %d\n", no);
+        }
+    };
+
+    this->runOnIndex(fn);
+
+    if (this->metric_type == METRIC_L2) {
+        merge_knn_results<idx_t, CMin<distance_t, int>>(
+                n,
+                k,
+                nshard,
+                all_distances.data(),
+                all_labels.data(),
+                distances,
+                labels);
+    } else {
+        merge_knn_results<idx_t, CMax<distance_t, int>>(
+                n,
+                k,
+                nshard,
+                all_distances.data(),
+                all_labels.data(),
+                distances,
+                labels);
+    }
+}
+
+} // namespace faiss

faiss/IndexShardsIVF.h

@@ -0,0 +1,42 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <faiss/IndexIVF.h>
+#include <faiss/IndexShards.h>
+
+namespace faiss {
+
+/**
+ * IndexShards with a common coarse quantizer. All the indexes added should be
+ * IndexIVFInterface indexes so that the search_precomputed can be called.
+ */
+struct IndexShardsIVF : public IndexShards, Level1Quantizer {
+    explicit IndexShardsIVF(
+            Index* quantizer,
+            size_t nlist,
+            bool threaded = false,
+            bool successive_ids = true);
+
+    void addIndex(Index* index) override;
+
+    void add_with_ids(idx_t n, const component_t* x, const idx_t* xids)
+            override;
+
+    void train(idx_t n, const component_t* x) override;
+
+    void search(
+            idx_t n,
+            const component_t* x,
+            idx_t k,
+            distance_t* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+};
+
+} // namespace faiss

faiss/gpu/CMakeLists.txt

@@ -7,7 +7,6 @@
 set(FAISS_GPU_SRC
   GpuAutoTune.cpp
   GpuCloner.cpp
-  GpuClonerOptions.cpp
   GpuDistance.cu
   GpuIcmEncoder.cu
   GpuIndex.cu

faiss/gpu/GpuAutoTune.cpp

@@ -11,6 +11,8 @@
 #include <faiss/IndexPreTransform.h>
 #include <faiss/IndexReplicas.h>
 #include <faiss/IndexShards.h>
+#include <faiss/IndexShardsIVF.h>
+
 #include <faiss/gpu/GpuIndex.h>
 #include <faiss/gpu/GpuIndexFlat.h>
 #include <faiss/gpu/GpuIndexIVFFlat.h>
@@ -33,7 +35,12 @@ using namespace ::faiss;
 
 void GpuParameterSpace::initialize(const Index* index) {
     if (DC(IndexPreTransform)) {
-        index = ix->index;
+        initialize(ix->index);
+        return;
+    }
+    if (DC(IndexShardsIVF)) {
+        ParameterSpace::initialize(index);
+        return;
     }
     if (DC(IndexReplicas)) {
         if (ix->count() == 0)

faiss/gpu/GpuCloner.cpp

@@ -18,6 +18,7 @@
 #include <faiss/IndexPreTransform.h>
 #include <faiss/IndexReplicas.h>
 #include <faiss/IndexScalarQuantizer.h>
+#include <faiss/IndexShardsIVF.h>
 #include <faiss/MetaIndexes.h>
 #include <faiss/gpu/GpuIndex.h>
 #include <faiss/gpu/GpuIndexFlat.h>
@@ -226,8 +227,8 @@ ToGpuClonerMultiple::ToGpuClonerMultiple(
         std::vector<int>& devices,
         const GpuMultipleClonerOptions& options)
         : GpuMultipleClonerOptions(options) {
-    FAISS_ASSERT(provider.size() == devices.size());
-    for (int i = 0; i < provider.size(); i++) {
+    FAISS_THROW_IF_NOT(provider.size() == devices.size());
+    for (size_t i = 0; i < provider.size(); i++) {
         sub_cloners.push_back(ToGpuCloner(provider[i], devices[i], options));
     }
 }
@@ -240,11 +241,11 @@ ToGpuClonerMultiple::ToGpuClonerMultiple(
 void ToGpuClonerMultiple::copy_ivf_shard(
         const IndexIVF* index_ivf,
         IndexIVF* idx2,
-        long n,
-        long i) {
+        idx_t n,
+        idx_t i) {
     if (shard_type == 2) {
-        long i0 = i * index_ivf->ntotal / n;
-        long i1 = (i + 1) * index_ivf->ntotal / n;
+        idx_t i0 = i * index_ivf->ntotal / n;
+        idx_t i1 = (i + 1) * index_ivf->ntotal / n;
 
         if (verbose)
             printf("IndexShards shard %ld indices %ld:%ld\n", i, i0, i1);
@@ -256,14 +257,27 @@ void ToGpuClonerMultiple::copy_ivf_shard(
             printf("IndexShards shard %ld select modulo %ld = %ld\n", i, n, i);
         index_ivf->copy_subset_to(
                 *idx2, InvertedLists::SUBSET_TYPE_ID_MOD, n, i);
+    } else if (shard_type == 4) {
+        idx_t i0 = i * index_ivf->nlist / n;
+        idx_t i1 = (i + 1) * index_ivf->nlist / n;
+        if (verbose) {
+            printf("IndexShards %ld/%ld select lists %d:%d\n",
+                   i,
+                   n,
+                   int(i0),
+                   int(i1));
+        }
+        index_ivf->copy_subset_to(
+                *idx2, InvertedLists::SUBSET_TYPE_INVLIST, i0, i1);
     } else {
         FAISS_THROW_FMT("shard_type %d not implemented", shard_type);
     }
 }
 
 Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
-    long n = sub_cloners.size();
+    idx_t n = sub_cloners.size();
 
+    auto index_ivf = dynamic_cast<const faiss::IndexIVF*>(index);
     auto index_ivfpq = dynamic_cast<const faiss::IndexIVFPQ*>(index);
     auto index_ivfflat = dynamic_cast<const faiss::IndexIVFFlat*>(index);
     auto index_ivfsq =
@@ -275,16 +289,36 @@ Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
             "IndexIVFFlat, IndexIVFScalarQuantizer, "
             "IndexFlat and IndexIVFPQ");
 
+    // decide what coarse quantizer the sub-indexes are going to have
+    const Index* quantizer = nullptr;
+    std::unique_ptr<Index> new_quantizer;
+    if (index_ivf) {
+        quantizer = index_ivf->quantizer;
+        if (common_ivf_quantizer &&
+            !dynamic_cast<const IndexFlat*>(quantizer)) {
+            // then we flatten the coarse quantizer so that everything remains
+            // on GPU
+            new_quantizer.reset(
+                    new IndexFlat(quantizer->d, quantizer->metric_type));
+            std::vector<float> centroids(quantizer->d * quantizer->ntotal);
+            quantizer->reconstruct_n(0, quantizer->ntotal, centroids.data());
+            new_quantizer->add(quantizer->ntotal, centroids.data());
+            quantizer = new_quantizer.get();
+        }
+    }
+
     std::vector<faiss::Index*> shards(n);
 
-    for (long i = 0; i < n; i++) {
+    for (idx_t i = 0; i < n; i++) {
         // make a shallow copy
-        if (reserveVecs)
+        if (reserveVecs) {
             sub_cloners[i].reserveVecs = (reserveVecs + n - 1) / n;
-
+        }
+        // note: const_casts here are harmless because the indexes build here
+        // are short-lived, translated immediately to GPU indexes.
         if (index_ivfpq) {
             faiss::IndexIVFPQ idx2(
-                    index_ivfpq->quantizer,
+                    const_cast<Index*>(quantizer),
                     index_ivfpq->d,
                     index_ivfpq->nlist,
                     index_ivfpq->code_size,
@@ -298,7 +332,7 @@ Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
             shards[i] = sub_cloners[i].clone_Index(&idx2);
         } else if (index_ivfflat) {
             faiss::IndexIVFFlat idx2(
-                    index_ivfflat->quantizer,
+                    const_cast<Index*>(quantizer),
                     index->d,
                     index_ivfflat->nlist,
                     index_ivfflat->metric_type);
@@ -308,7 +342,7 @@ Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
             shards[i] = sub_cloners[i].clone_Index(&idx2);
         } else if (index_ivfsq) {
             faiss::IndexIVFScalarQuantizer idx2(
-                    index_ivfsq->quantizer,
+                    const_cast<Index*>(quantizer),
                     index->d,
                     index_ivfsq->nlist,
                     index_ivfsq->sq.qtype,
@@ -324,40 +358,52 @@ Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
             faiss::IndexFlat idx2(index->d, index->metric_type);
             shards[i] = sub_cloners[i].clone_Index(&idx2);
             if (index->ntotal > 0) {
-                long i0 = index->ntotal * i / n;
-                long i1 = index->ntotal * (i + 1) / n;
+                idx_t i0 = index->ntotal * i / n;
+                idx_t i1 = index->ntotal * (i + 1) / n;
                 shards[i]->add(i1 - i0, index_flat->get_xb() + i0 * index->d);
             }
         }
     }
 
     bool successive_ids = index_flat != nullptr;
-    faiss::IndexShards* res =
-            new faiss::IndexShards(index->d, true, successive_ids);
+    faiss::IndexShards* res;
+    if (common_ivf_quantizer && index_ivf) {
+        this->shard = false;
+        Index* common_quantizer = clone_Index(index_ivf->quantizer);
+        this->shard = true;
+        IndexShardsIVF* idx = new faiss::IndexShardsIVF(
+                common_quantizer, index_ivf->nlist, true, false);
+        idx->own_fields = true;
+        idx->own_indices = true;
+        res = idx;
+    } else {
+        res = new faiss::IndexShards(index->d, true, successive_ids);
+        res->own_indices = true;
+    }
 
     for (int i = 0; i < n; i++) {
         res->add_shard(shards[i]);
     }
-    res->own_fields = true;
     FAISS_ASSERT(index->ntotal == res->ntotal);
     return res;
 }
 
 Index* ToGpuClonerMultiple::clone_Index(const Index* index) {
-    long n = sub_cloners.size();
-    if (n == 1)
+    idx_t n = sub_cloners.size();
+    if (n == 1) {
         return sub_cloners[0].clone_Index(index);
+    }
 
     if (dynamic_cast<const IndexFlat*>(index) ||
-        dynamic_cast<const faiss::IndexIVFFlat*>(index) ||
-        dynamic_cast<const faiss::IndexIVFScalarQuantizer*>(index) ||
-        dynamic_cast<const faiss::IndexIVFPQ*>(index)) {
+        dynamic_cast<const IndexIVFFlat*>(index) ||
+        dynamic_cast<const IndexIVFScalarQuantizer*>(index) ||
+        dynamic_cast<const IndexIVFPQ*>(index)) {
         if (!shard) {
             IndexReplicas* res = new IndexReplicas();
             for (auto& sub_cloner : sub_cloners) {
                 res->addIndex(sub_cloner.clone_Index(index));
             }
-            res->own_fields = true;
+            res->own_indices = true;
             return res;
         } else {
             return clone_Index_to_shards(index);
@@ -374,8 +420,8 @@ Index* ToGpuClonerMultiple::clone_Index(const Index* index) {
         for (int m = 0; m < pq.M; m++) {
             // which GPU(s) will be assigned to this sub-quantizer
 
-            long i0 = m * n / pq.M;
-            long i1 = pq.M <= n ? (m + 1) * n / pq.M : i0 + 1;
+            idx_t i0 = m * n / pq.M;
+            idx_t i1 = pq.M <= n ? (m + 1) * n / pq.M : i0 + 1;
             std::vector<ToGpuCloner> sub_cloners_2;
             sub_cloners_2.insert(
                     sub_cloners_2.begin(),

faiss/gpu/GpuCloner.h

@@ -55,8 +55,8 @@ struct ToGpuClonerMultiple : faiss::Cloner, GpuMultipleClonerOptions {
     void copy_ivf_shard(
             const IndexIVF* index_ivf,
             IndexIVF* idx2,
-            long n,
-            long i);
+            idx_t n,
+            idx_t i);
 
     Index* clone_Index_to_shards(const Index* index);
 

faiss/gpu/GpuClonerOptions.cpp

@@ -1,18 +0,0 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
- *
- * This source code is licensed under the MIT license found in the
- * LICENSE file in the root directory of this source tree.
- */
-
-#include <faiss/gpu/GpuClonerOptions.h>
-
-namespace faiss {
-namespace gpu {
-
-GpuClonerOptions::GpuClonerOptions() {}
-
-GpuMultipleClonerOptions::GpuMultipleClonerOptions() {}
-
-} // namespace gpu
-} // namespace faiss

faiss/gpu/GpuClonerOptions.h

@@ -14,8 +14,6 @@ namespace gpu {
 
 /// set some options on how to copy to GPU
 struct GpuClonerOptions {
-    GpuClonerOptions();
-
     /// how should indices be stored on index types that support indices
     /// (anything but GpuIndexFlat*)?
     IndicesOptions indicesOptions = INDICES_64_BIT;
@@ -41,14 +39,17 @@ struct GpuClonerOptions {
 };
 
 struct GpuMultipleClonerOptions : public GpuClonerOptions {
-    GpuMultipleClonerOptions();
-
     /// Whether to shard the index across GPUs, versus replication
     /// across GPUs
     bool shard = false;
 
     /// IndexIVF::copy_subset_to subset type
     int shard_type = 1;
+
+    /// set to true if an IndexIVF is to be dispatched to multiple GPUs with a
+    /// single common IVF quantizer, ie. only the inverted lists are sharded on
+    /// the sub-indexes (uses an IndexShardsIVF)
+    bool common_ivf_quantizer = false;
 };
 
 } // namespace gpu

faiss/gpu/GpuIndexIVF.cu

@@ -27,10 +27,7 @@ GpuIndexIVF::GpuIndexIVF(
         idx_t nlistIn,
         GpuIndexIVFConfig config)
         : GpuIndex(provider->getResources(), dims, metric, metricArg, config),
-          nlist(nlistIn),
-          nprobe(1),
-          quantizer(nullptr),
-          own_fields(false),
+          IndexIVFInterface(nullptr, nlistIn),
           ivfConfig_(config) {
     // Only IP and L2 are supported for now
     if (!(metric_type == faiss::METRIC_L2 ||
@@ -50,9 +47,7 @@ GpuIndexIVF::GpuIndexIVF(
         idx_t nlistIn,
         GpuIndexIVFConfig config)
         : GpuIndex(provider->getResources(), dims, metric, metricArg, config),
-          nlist(nlistIn),
-          nprobe(1),
-          quantizer(coarseQuantizer),
+          IndexIVFInterface(coarseQuantizer, nlistIn),
           ivfConfig_(config) {
     FAISS_THROW_IF_NOT_MSG(
             quantizer, "expecting a coarse quantizer object; none provided");
@@ -113,11 +108,7 @@ void GpuIndexIVF::init_() {
     verifyIVFSettings_();
 }
 
-GpuIndexIVF::~GpuIndexIVF() {
-    if (own_fields) {
-        delete quantizer;
-    }
-}
+GpuIndexIVF::~GpuIndexIVF() {}
 
 void GpuIndexIVF::verifyIVFSettings_() const {
     // We should always have a quantizer instance
@@ -130,7 +121,7 @@ void GpuIndexIVF::verifyIVFSettings_() const {
         // IVF quantizer should correspond to our set of lists
         FAISS_THROW_IF_NOT_FMT(
                 quantizer->ntotal == nlist,
-                "IVF nlist count (%ld) does not match trained coarse quantizer size (%zu)",
+                "IVF nlist count (%zu) does not match trained coarse quantizer size (%zu)",
                 nlist,
                 quantizer->ntotal);
     } else {
@@ -336,13 +327,15 @@ void GpuIndexIVF::searchImpl_(
 void GpuIndexIVF::search_preassigned(
         idx_t n,
         const float* x,
-        int k,
+        idx_t k,
         const idx_t* assign,
         const float* centroid_dis,
         float* distances,
         idx_t* labels,
         bool store_pairs,
-        const IVFSearchParameters* params) const {
+        const IVFSearchParameters* params,
+        IndexIVFStats* stats) const {
+    FAISS_THROW_IF_NOT_MSG(stats == nullptr, "IVF stats not supported");
     DeviceScope scope(config_.device);
     auto stream = resources_->getDefaultStream(config_.device);
 
@@ -360,7 +353,15 @@ void GpuIndexIVF::search_preassigned(
         return;
     }
 
-    int use_nprobe = getCurrentNProbe_(params);
+    idx_t use_nprobe = params ? params->nprobe : this->nprobe;
+    validateNProbe(use_nprobe);
+
+    size_t max_codes = params ? params->max_codes : this->max_codes;
+    FAISS_THROW_IF_NOT_FMT(
+            max_codes == 0,
+            "GPU IVF index does not currently support "
+            "SearchParametersIVF::max_codes (passed %zu, must be 0)",
+            max_codes);
 
     // Ensure that all data/output buffers are resident on our desired device
     auto vecsDevice = toDeviceTemporary<float, 2>(
@@ -405,6 +406,19 @@ void GpuIndexIVF::search_preassigned(
     fromDevice<idx_t, 2>(outIndicesDevice, labels, stream);
 }
 
+void GpuIndexIVF::range_search_preassigned(
+        idx_t nx,
+        const float* x,
+        float radius,
+        const idx_t* keys,
+        const float* coarse_dis,
+        RangeSearchResult* result,
+        bool store_pairs,
+        const IVFSearchParameters* params,
+        IndexIVFStats* stats) const {
+    FAISS_THROW_MSG("range search not implemented");
+}
+
 bool GpuIndexIVF::addImplRequiresIDs_() const {
     // All IVF indices have storage for IDs
     return true;

faiss/gpu/GpuIndexIVF.h

@@ -33,7 +33,7 @@ struct GpuIndexIVFConfig : public GpuIndexConfig {
 /// Base class of all GPU IVF index types. This (for now) deliberately does not
 /// inherit from IndexIVF, as many of the public data members and functionality
 /// in IndexIVF is not supported in the same manner on the GPU.
-class GpuIndexIVF : public GpuIndex {
+class GpuIndexIVF : public GpuIndex, public IndexIVFInterface {
    public:
     /// Version that auto-constructs a flat coarse quantizer based on the
     /// desired metric
@@ -93,31 +93,30 @@ class GpuIndexIVF : public GpuIndex {
     /// debugging purposes.
     std::vector<idx_t> getListIndices(idx_t listId) const;
 
-    /// Same interface as faiss::IndexIVF, in order to search a set of vectors
-    /// pre-quantized by the IVF quantizer. Does not include IndexIVFStats as
-    /// that can only be obtained on the host via a GPU d2h copy.
-    /// @param n      nb of vectors to query
-    /// @param x      query vectors, size nx * d
-    /// @param assign coarse quantization indices, size nx * nprobe
-    /// @param centroid_dis
-    ///             distances to coarse centroids, size nx * nprobe
-    /// @param distance
-    ///             output distances, size n * k
-    /// @param labels output labels, size n * k
-    /// @param store_pairs store inv list index + inv list offset
-    ///                   instead in upper/lower 32 bit of result,
-    ///                   instead of ids (used for reranking).
-    /// @param params used to override the object's search parameters
     void search_preassigned(
             idx_t n,
             const float* x,
-            int k,
+            idx_t k,
             const idx_t* assign,
             const float* centroid_dis,
             float* distances,
             idx_t* labels,
             bool store_pairs,
-            const SearchParametersIVF* params = nullptr) const;
+            const SearchParametersIVF* params = nullptr,
+            IndexIVFStats* stats = nullptr) const override;
+
+    // not implemented for GPU
+    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 override;
+
 
    protected:
     /// From either the current set nprobe or the SearchParameters if available,
@@ -139,23 +138,6 @@ class GpuIndexIVF : public GpuIndex {
             idx_t* labels,
             const SearchParameters* params) const override;
 
-   public:
-    /// Exposing this like the CPU version for manipulation
-    ClusteringParameters cp;
-
-    /// Exposing this like the CPU version for query
-    idx_t nlist;
-
-    /// Exposing this like the CPU version for manipulation
-    /// FIXME: IndexIVF has size_t for nprobe, not idx_t
-    size_t nprobe;
-
-    /// A user-pluggable coarse quantizer
-    Index* quantizer;
-
-    /// Whether or not we own the coarse quantizer
-    bool own_fields;
-
    protected:
     /// Our configuration options
     const GpuIndexIVFConfig ivfConfig_;

faiss/gpu/test/test_gpu_index.py

@@ -6,7 +6,6 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import math
-import time
 import unittest
 import numpy as np
 import faiss
@@ -14,6 +13,7 @@ from faiss.contrib import datasets
 from faiss.contrib import ivf_tools
 from faiss.contrib.evaluation import knn_intersection_measure
 
+
 class TestIVFSearchPreassigned(unittest.TestCase):
     def test_ivfflat_search_preassigned(self):
         res = faiss.StandardGpuResources()
@@ -522,9 +522,10 @@ class TestInvalidParams(unittest.TestCase):
         idx.nprobe = 3
         self.assertRaises(AssertionError, idx.search, xb[10:20], k)
 
-        k = 5
-        idx.nprobe = 0
-        self.assertRaises(RuntimeError, idx.search, xb[10:20], k)
+        # nprobe is unsigned now, so this is caught before reaching C++
+        # k = 5
+        # idx.nprobe = -3
+        # self.assertRaises(RuntimeError, idx.search, xb[10:20], k)
 
         # valid params
         k = 5

faiss/gpu/test/test_multi_gpu.py

@@ -8,8 +8,11 @@ import unittest
 import numpy as np
 import faiss
 
+from faiss.contrib.datasets import SyntheticDataset
+
 
 class TestShardedFlat(unittest.TestCase):
+
     @unittest.skipIf(faiss.get_num_gpus() < 2, "multiple GPU only test")
     def test_sharded(self):
         d = 32
@@ -49,6 +52,52 @@ class TestShardedFlat(unittest.TestCase):
         else:
             raise AssertionError("errpr: call should fail but isn't failing")
 
+    @unittest.skipIf(faiss.get_num_gpus() < 2, "multiple GPU only test")
+    def do_test_sharded_ivf(self, index_key):
+        ds = SyntheticDataset(32, 8000, 10000, 100)
+        index = faiss.index_factory(ds.d, index_key)
+        if 'HNSW' in index_key:
+            # make a bit more reproducible...
+            faiss.ParameterSpace().set_index_parameter(
+                index, 'quantizer_efSearch', 40)
+        index.train(ds.get_train())
+        index.add(ds.get_database())
+        Dref, Iref = index.search(ds.get_queries(), 10)
+        index.nprobe = 8
+        Dref8, Iref8 = index.search(ds.get_queries(), 10)
+        index.nprobe = 1
+        print("REF checksum", faiss.checksum(Iref))
+
+        co = faiss.GpuMultipleClonerOptions()
+        co.shard = True
+        co.common_ivf_quantizer = True
+        index = faiss.index_cpu_to_all_gpus(index, co, ngpu=2)
+
+        index.quantizer  # make sure there is indeed a quantizer
+        print("QUANT", faiss.downcast_index(index.quantizer))
+        Dnew, Inew = index.search(ds.get_queries(), 10)
+        np.testing.assert_array_equal(Iref, Inew)
+        np.testing.assert_array_almost_equal(Dref, Dnew, decimal=4)
+
+        # the nprobe is taken from the sub-indexes
+        faiss.GpuParameterSpace().set_index_parameter(index, 'nprobe', 8)
+        Dnew8, Inew8 = index.search(ds.get_queries(), 10)
+        np.testing.assert_array_equal(Iref8, Inew8)
+        np.testing.assert_array_almost_equal(Dref8, Dnew8, decimal=4)
+
+        index.reset()
+        index.add(ds.get_database())
+
+        Dnew8, Inew8 = index.search(ds.get_queries(), 10)
+        np.testing.assert_array_equal(Iref8, Inew8)
+        np.testing.assert_array_almost_equal(Dref8, Dnew8, decimal=4)
+
+    def test_sharded_IVFSQ(self):
+        self.do_test_sharded_ivf("IVF128,SQ8")
+
+    def test_sharded_IVF_HNSW(self):
+        self.do_test_sharded_ivf("IVF1000_HNSW,Flat")
+
 
 # This class also has a multi-GPU test within
 class EvalIVFPQAccuracy(unittest.TestCase):

faiss/impl/ThreadedIndex-inl.h

@@ -18,7 +18,7 @@ ThreadedIndex<IndexT>::ThreadedIndex(bool threaded)
 
 template <typename IndexT>
 ThreadedIndex<IndexT>::ThreadedIndex(int d, bool threaded)
-        : IndexT(d), own_fields(false), isThreaded_(threaded) {}
+        : IndexT(d), isThreaded_(threaded) {}
 
 template <typename IndexT>
 ThreadedIndex<IndexT>::~ThreadedIndex() {
@@ -35,7 +35,7 @@ ThreadedIndex<IndexT>::~ThreadedIndex() {
             FAISS_ASSERT(!(bool)p.second);
         }
 
-        if (own_fields) {
+        if (own_indices) {
             delete p.first;
         }
     }
@@ -102,7 +102,7 @@ void ThreadedIndex<IndexT>::removeIndex(IndexT* index) {
             indices_.erase(it);
             onAfterRemoveIndex(index);
 
-            if (own_fields) {
+            if (own_indices) {
                 delete index;
             }
 

faiss/impl/ThreadedIndex.h

@@ -29,7 +29,7 @@ class ThreadedIndex : public IndexT {
     /// WARNING: once an index is added, it becomes unsafe to touch it from any
     /// other thread than that on which is managing it, until we are shut
     /// down. Use runOnIndex to perform work on it instead.
-    void addIndex(IndexT* index);
+    virtual void addIndex(IndexT* index);
 
     /// Remove an index that is managed by ourselves.
     /// This will flush all pending work on that index, and then shut
@@ -52,17 +52,17 @@ class ThreadedIndex : public IndexT {
     }
 
     /// Returns the i-th sub-index
-    IndexT* at(int i) {
+    IndexT* at(size_t i) {
         return indices_[i].first;
     }
 
     /// Returns the i-th sub-index (const version)
-    const IndexT* at(int i) const {
+    const IndexT* at(size_t i) const {
         return indices_[i].first;
     }
 
     /// Whether or not we are responsible for deleting our contained indices
-    bool own_fields;
+    bool own_indices = false;
 
    protected:
     /// Called just after an index is added

faiss/python/swigfaiss.swig

@@ -90,6 +90,7 @@ typedef uint64_t size_t;
 #include <faiss/IndexIVFSpectralHash.h>
 #include <faiss/impl/ThreadedIndex.h>
 #include <faiss/IndexShards.h>
+#include <faiss/IndexShardsIVF.h>
 #include <faiss/IndexReplicas.h>
 #include <faiss/impl/HNSW.h>
 #include <faiss/IndexHNSW.h>
@@ -484,6 +485,7 @@ void gpu_sync_all_devices()
 %include  <faiss/IndexShards.h>
 %template(IndexShards) faiss::IndexShardsTemplate<faiss::Index>;
 %template(IndexBinaryShards) faiss::IndexShardsTemplate<faiss::IndexBinary>;
+%include  <faiss/IndexShardsIVF.h>
 
 %include  <faiss/IndexReplicas.h>
 %template(IndexReplicas) faiss::IndexReplicasTemplate<faiss::Index>;
@@ -578,6 +580,7 @@ void gpu_sync_all_devices()
 %typemap(out) faiss::Index * {
     DOWNCAST2 ( IndexIDMap2, IndexIDMap2TemplateT_faiss__Index_t )
     DOWNCAST2 ( IndexIDMap, IndexIDMapTemplateT_faiss__Index_t )
+    DOWNCAST ( IndexShardsIVF )
     DOWNCAST2 ( IndexShards, IndexShardsTemplateT_faiss__Index_t )
     DOWNCAST2 ( IndexReplicas, IndexReplicasTemplateT_faiss__Index_t )
     DOWNCAST ( IndexIVFPQR )

tests/test_merge.cpp

@@ -148,7 +148,7 @@ int compare_merged(
 // test on IVFFlat with implicit numbering
 TEST(MERGE, merge_flat_no_ids) {
     faiss::IndexShards index_shards(d);
-    index_shards.own_fields = true;
+    index_shards.own_indices = true;
     for (int i = 0; i < nindex; i++) {
         index_shards.add_shard(
                 new faiss::IndexIVFFlat(&cd.quantizer, d, nlist));
@@ -164,7 +164,7 @@ TEST(MERGE, merge_flat_no_ids) {
 // test on IVFFlat, explicit ids
 TEST(MERGE, merge_flat) {
     faiss::IndexShards index_shards(d, false, false);
-    index_shards.own_fields = true;
+    index_shards.own_indices = true;
 
     for (int i = 0; i < nindex; i++) {
         index_shards.add_shard(
@@ -180,7 +180,7 @@ TEST(MERGE, merge_flat) {
 // test on IVFFlat and a VectorTransform
 TEST(MERGE, merge_flat_vt) {
     faiss::IndexShards index_shards(d, false, false);
-    index_shards.own_fields = true;
+    index_shards.own_indices = true;
 
     // here we have to retrain because of the vectorTransform
     faiss::RandomRotationMatrix rot(d, d);
@@ -211,7 +211,7 @@ TEST(MERGE, merge_flat_vt) {
 // put the merged invfile on disk
 TEST(MERGE, merge_flat_ondisk) {
     faiss::IndexShards index_shards(d, false, false);
-    index_shards.own_fields = true;
+    index_shards.own_indices = true;
     Tempfilename filename;
 
     for (int i = 0; i < nindex; i++) {
@@ -234,7 +234,7 @@ TEST(MERGE, merge_flat_ondisk) {
 // now use ondisk specific merge
 TEST(MERGE, merge_flat_ondisk_2) {
     faiss::IndexShards index_shards(d, false, false);
-    index_shards.own_fields = true;
+    index_shards.own_indices = true;
 
     for (int i = 0; i < nindex; i++) {
         index_shards.add_shard(

tests/test_meta_index.py

@@ -10,6 +10,8 @@ import unittest
 
 from common_faiss_tests import Randu10k
 
+from faiss.contrib.datasets import SyntheticDataset
+
 ru = Randu10k()
 
 xb = ru.xb
@@ -130,6 +132,52 @@ class Shards(unittest.TestCase):
             print('%d / %d differences' % (ndiff, nq * k))
             assert (ndiff < nq * k / 1000.)
 
+    def test_shards_ivf(self):
+        ds = SyntheticDataset(32, 1000, 100, 20)
+        ref_index = faiss.index_factory(ds.d, "IVF32,SQ8")
+        ref_index.train(ds.get_train())
+        xb = ds.get_database()
+        ref_index.add(ds.get_database())
 
-if __name__ == '__main__':
-    unittest.main()
+        Dref, Iref = ref_index.search(ds.get_database(), 10)
+        ref_index.reset()
+
+        sharded_index = faiss.IndexShardsIVF(
+            ref_index.quantizer, ref_index.nlist, False, True)
+        for shard in range(3):
+            index_i = faiss.clone_index(ref_index)
+            index_i.add(xb[shard * nb // 3: (shard + 1)* nb // 3])
+            sharded_index.add_shard(index_i)
+
+        Dnew, Inew = sharded_index.search(ds.get_database(), 10)
+
+        np.testing.assert_equal(Inew, Iref)
+        np.testing.assert_allclose(Dnew, Dref)
+
+    def test_shards_ivf_train_add(self):
+        ds = SyntheticDataset(32, 1000, 600, 20)
+        quantizer = faiss.IndexFlatL2(ds.d)
+        sharded_index = faiss.IndexShardsIVF(quantizer, 40, False, False)
+
+        for _ in range(3):
+            sharded_index.add_shard(faiss.index_factory(ds.d, "IVF40,Flat"))
+
+        sharded_index.train(ds.get_train())
+        sharded_index.add(ds.get_database())
+        Dnew, Inew = sharded_index.search(ds.get_queries(), 10)
+
+        index_ref = faiss.IndexIVFFlat(quantizer, ds.d, sharded_index.nlist)
+        index_ref.train(ds.get_train())
+        index_ref.add(ds.get_database())
+        Dref, Iref = index_ref.search(ds.get_queries(), 10)
+        np.testing.assert_equal(Inew, Iref)
+        np.testing.assert_allclose(Dnew, Dref)
+
+        # mess around with the quantizer's centroids
+        centroids = quantizer.reconstruct_n()
+        centroids = centroids[::-1].copy()
+        quantizer.reset()
+        quantizer.add(centroids)
+
+        D2, I2 = sharded_index.search(ds.get_queries(), 10)
+        self.assertFalse(np.all(I2 == Inew))