diff --git a/c_api/IndexReplicas_c.cpp b/c_api/IndexReplicas_c.cpp index 37b42a50f..820b80323 100644 --- a/c_api/IndexReplicas_c.cpp +++ b/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 { diff --git a/c_api/IndexShards_c.cpp b/c_api/IndexShards_c.cpp index e4d2c47d4..56b520b38 100644 --- a/c_api/IndexShards_c.cpp +++ b/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) diff --git a/faiss/AutoTune.cpp b/faiss/AutoTune.cpp index 3cec38cb4..f9aa4273c 100644 --- a/faiss/AutoTune.cpp +++ b/faiss/AutoTune.cpp @@ -32,6 +32,7 @@ #include #include #include +#include #include #include @@ -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)) { // call on all sub-indexes auto fn = [this, name, val](int /* no */, Index* subIndex) { diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 16eb9e9c9..bcfd71045 100644 --- a/faiss/CMakeLists.txt +++ b/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 diff --git a/faiss/IndexIVF.cpp b/faiss/IndexIVF.cpp index c04665649..19e18e466 100644 --- a/faiss/IndexIVF.cpp +++ b/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); diff --git a/faiss/IndexIVF.h b/faiss/IndexIVF.h index 9af54fd20..ade8b5113 100644 --- a/faiss/IndexIVF.h +++ b/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 diff --git a/faiss/IndexShards.cpp b/faiss/IndexShards.cpp index a20dfbe59..8d5c41ccc 100644 --- a/faiss/IndexShards.cpp +++ b/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::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::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 all_distances(nshard * k * n); std::vector all_labels(nshard * k * n); - std::vector translations(nshard, 0); + std::vector 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::search( template struct IndexShardsTemplate; template struct IndexShardsTemplate; + } // namespace faiss diff --git a/faiss/IndexShards.h b/faiss/IndexShards.h index 479302301..278d5724a 100644 --- a/faiss/IndexShards.h +++ b/faiss/IndexShards.h @@ -95,7 +95,7 @@ struct IndexShardsTemplate : public ThreadedIndex { /// 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 diff --git a/faiss/IndexShardsIVF.cpp b/faiss/IndexShardsIVF.cpp new file mode 100644 index 000000000..9a68f3eb7 --- /dev/null +++ b/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 + +#include +#include +#include + +#include +#include +#include +#include + +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(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(index); + FAISS_THROW_IF_NOT_MSG(index_ivf, "can only add IndexIVFs"); + FAISS_THROW_IF_NOT(index_ivf->nlist == nlist); + IndexShardsTemplate::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 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(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(index); + } + if (!all_index_ivf) { + IndexShardsTemplate::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 Iq(n); + std::vector 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 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(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(params_in); + FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type"); + } + + auto index0 = dynamic_cast(at(0)); + idx_t nprobe = params ? params->nprobe : index0->nprobe; + + // coarse quantization (TODO: support tiling with search_precomputed) + std::vector Dq(n * nprobe); + std::vector Iq(n * nprobe); + + quantizer->search(n, x, nprobe, Dq.data(), Iq.data()); + + int64_t nshard = this->count(); + + std::vector all_distances(nshard * k * n); + std::vector all_labels(nshard * k * n); + std::vector 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(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>( + n, + k, + nshard, + all_distances.data(), + all_labels.data(), + distances, + labels); + } else { + merge_knn_results>( + n, + k, + nshard, + all_distances.data(), + all_labels.data(), + distances, + labels); + } +} + +} // namespace faiss diff --git a/faiss/IndexShardsIVF.h b/faiss/IndexShardsIVF.h new file mode 100644 index 000000000..8b17e22b8 --- /dev/null +++ b/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 +#include + +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 diff --git a/faiss/gpu/CMakeLists.txt b/faiss/gpu/CMakeLists.txt index 5a8c4555c..6fb694764 100644 --- a/faiss/gpu/CMakeLists.txt +++ b/faiss/gpu/CMakeLists.txt @@ -7,7 +7,6 @@ set(FAISS_GPU_SRC GpuAutoTune.cpp GpuCloner.cpp - GpuClonerOptions.cpp GpuDistance.cu GpuIcmEncoder.cu GpuIndex.cu diff --git a/faiss/gpu/GpuAutoTune.cpp b/faiss/gpu/GpuAutoTune.cpp index b91b2b2a3..4a7ba5ffc 100644 --- a/faiss/gpu/GpuAutoTune.cpp +++ b/faiss/gpu/GpuAutoTune.cpp @@ -11,6 +11,8 @@ #include #include #include +#include + #include #include #include @@ -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) diff --git a/faiss/gpu/GpuCloner.cpp b/faiss/gpu/GpuCloner.cpp index 5afb491ee..80e1fa20c 100644 --- a/faiss/gpu/GpuCloner.cpp +++ b/faiss/gpu/GpuCloner.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -226,8 +227,8 @@ ToGpuClonerMultiple::ToGpuClonerMultiple( std::vector& 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(index); auto index_ivfpq = dynamic_cast(index); auto index_ivfflat = dynamic_cast(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 new_quantizer; + if (index_ivf) { + quantizer = index_ivf->quantizer; + if (common_ivf_quantizer && + !dynamic_cast(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 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 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(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(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(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(index) || - dynamic_cast(index) || - dynamic_cast(index) || - dynamic_cast(index)) { + dynamic_cast(index) || + dynamic_cast(index) || + dynamic_cast(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 sub_cloners_2; sub_cloners_2.insert( sub_cloners_2.begin(), diff --git a/faiss/gpu/GpuCloner.h b/faiss/gpu/GpuCloner.h index c692ec69b..584f1163f 100644 --- a/faiss/gpu/GpuCloner.h +++ b/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); diff --git a/faiss/gpu/GpuClonerOptions.cpp b/faiss/gpu/GpuClonerOptions.cpp deleted file mode 100644 index be50a2d06..000000000 --- a/faiss/gpu/GpuClonerOptions.cpp +++ /dev/null @@ -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 - -namespace faiss { -namespace gpu { - -GpuClonerOptions::GpuClonerOptions() {} - -GpuMultipleClonerOptions::GpuMultipleClonerOptions() {} - -} // namespace gpu -} // namespace faiss diff --git a/faiss/gpu/GpuClonerOptions.h b/faiss/gpu/GpuClonerOptions.h index 40809b268..dd3e6c6e8 100644 --- a/faiss/gpu/GpuClonerOptions.h +++ b/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 diff --git a/faiss/gpu/GpuIndexIVF.cu b/faiss/gpu/GpuIndexIVF.cu index 7261ba943..f2ed32360 100644 --- a/faiss/gpu/GpuIndexIVF.cu +++ b/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( @@ -405,6 +406,19 @@ void GpuIndexIVF::search_preassigned( fromDevice(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; diff --git a/faiss/gpu/GpuIndexIVF.h b/faiss/gpu/GpuIndexIVF.h index 98bfe7596..8367df06b 100644 --- a/faiss/gpu/GpuIndexIVF.h +++ b/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 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_; diff --git a/faiss/gpu/test/test_gpu_index.py b/faiss/gpu/test/test_gpu_index.py index 3ab0b355c..fcc4f0813 100755 --- a/faiss/gpu/test/test_gpu_index.py +++ b/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 diff --git a/faiss/gpu/test/test_multi_gpu.py b/faiss/gpu/test/test_multi_gpu.py index 5aac4d5b0..568d444a6 100644 --- a/faiss/gpu/test/test_multi_gpu.py +++ b/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): diff --git a/faiss/impl/ThreadedIndex-inl.h b/faiss/impl/ThreadedIndex-inl.h index e1c91c16c..1dbbb6727 100644 --- a/faiss/impl/ThreadedIndex-inl.h +++ b/faiss/impl/ThreadedIndex-inl.h @@ -18,7 +18,7 @@ ThreadedIndex::ThreadedIndex(bool threaded) template ThreadedIndex::ThreadedIndex(int d, bool threaded) - : IndexT(d), own_fields(false), isThreaded_(threaded) {} + : IndexT(d), isThreaded_(threaded) {} template ThreadedIndex::~ThreadedIndex() { @@ -35,7 +35,7 @@ ThreadedIndex::~ThreadedIndex() { FAISS_ASSERT(!(bool)p.second); } - if (own_fields) { + if (own_indices) { delete p.first; } } @@ -102,7 +102,7 @@ void ThreadedIndex::removeIndex(IndexT* index) { indices_.erase(it); onAfterRemoveIndex(index); - if (own_fields) { + if (own_indices) { delete index; } diff --git a/faiss/impl/ThreadedIndex.h b/faiss/impl/ThreadedIndex.h index 2f5b8a36d..eaea6d933 100644 --- a/faiss/impl/ThreadedIndex.h +++ b/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 diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index b1785b5a8..404644b65 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -90,6 +90,7 @@ typedef uint64_t size_t; #include #include #include +#include #include #include #include @@ -484,6 +485,7 @@ void gpu_sync_all_devices() %include %template(IndexShards) faiss::IndexShardsTemplate; %template(IndexBinaryShards) faiss::IndexShardsTemplate; +%include %include %template(IndexReplicas) faiss::IndexReplicasTemplate; @@ -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 ) diff --git a/tests/test_merge.cpp b/tests/test_merge.cpp index 493fb3ccd..7e23f15f7 100644 --- a/tests/test_merge.cpp +++ b/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( diff --git a/tests/test_meta_index.py b/tests/test_meta_index.py index 370dfab7d..d53cad48f 100644 --- a/tests/test_meta_index.py +++ b/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))