faiss 0.3.4 → 0.4.0

data/vendor/faiss/faiss/clone_index.cpp

@@ -19,6 +19,8 @@
 #include <faiss/IndexAdditiveQuantizerFastScan.h>
 #include <faiss/IndexBinary.h>
 #include <faiss/IndexBinaryFlat.h>
+#include <faiss/IndexBinaryHNSW.h>
+#include <faiss/IndexBinaryIVF.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexIVF.h>
@@ -107,6 +109,11 @@ IndexIVF* Cloner::clone_IndexIVF(const IndexIVF* ivf) {
     return nullptr;
 }
 
+IndexBinaryIVF* clone_IndexBinaryIVF(const IndexBinaryIVF* ivf) {
+    TRYCLONE(IndexBinaryIVF, ivf)
+    return nullptr;
+}
+
 IndexRefine* clone_IndexRefine(const IndexRefine* ir) {
     TRYCLONE(IndexRefineFlat, ir)
     TRYCLONE(IndexRefine, ir) {
@@ -131,6 +138,11 @@ IndexHNSW* clone_IndexHNSW(const IndexHNSW* ihnsw) {
     }
 }
 
+IndexBinaryHNSW* clone_IndexBinaryHNSW(const IndexBinaryHNSW* ihnsw) {
+    TRYCLONE(IndexBinaryHNSW, ihnsw)
+    return nullptr;
+}
+
 IndexNNDescent* clone_IndexNNDescent(const IndexNNDescent* innd) {
     TRYCLONE(IndexNNDescentFlat, innd)
     TRYCLONE(IndexNNDescent, innd) {
@@ -323,9 +335,10 @@ Index* Cloner::clone_Index(const Index* index) {
         IndexNSG* res = clone_IndexNSG(insg);
 
         // copy the dynamic allocated graph
-        auto& new_graph = res->nsg.final_graph;
-        auto& old_graph = insg->nsg.final_graph;
-        new_graph = std::make_shared<nsg::Graph<int>>(*old_graph);
+        if (auto& old_graph = insg->nsg.final_graph) {
+            auto& new_graph = res->nsg.final_graph;
+            new_graph = std::make_shared<nsg::Graph<int>>(*old_graph);
+        }
 
         res->own_fields = true;
         res->storage = clone_Index(insg->storage);
@@ -385,6 +398,28 @@ Quantizer* clone_Quantizer(const Quantizer* quant) {
 IndexBinary* clone_binary_index(const IndexBinary* index) {
     if (auto ii = dynamic_cast<const IndexBinaryFlat*>(index)) {
         return new IndexBinaryFlat(*ii);
+    } else if (
+            const IndexBinaryIVF* ivf =
+                    dynamic_cast<const IndexBinaryIVF*>(index)) {
+        IndexBinaryIVF* res = clone_IndexBinaryIVF(ivf);
+        if (ivf->invlists == nullptr) {
+            res->invlists = nullptr;
+        } else {
+            res->invlists = clone_InvertedLists(ivf->invlists);
+            res->own_invlists = true;
+        }
+
+        res->own_fields = true;
+        res->quantizer = clone_binary_index(ivf->quantizer);
+
+        return res;
+    } else if (
+            const IndexBinaryHNSW* ihnsw =
+                    dynamic_cast<const IndexBinaryHNSW*>(index)) {
+        IndexBinaryHNSW* res = clone_IndexBinaryHNSW(ihnsw);
+        res->own_fields = true;
+        res->storage = clone_binary_index(ihnsw->storage);
+        return res;
     } else {
         FAISS_THROW_MSG("cannot clone this type of index");
     }

data/vendor/faiss/faiss/cppcontrib/factory_tools.cpp

@@ -8,8 +8,22 @@
 // -*- c++ -*-
 
 #include <faiss/cppcontrib/factory_tools.h>
+
 #include <map>
 
+#include <faiss/IndexBinaryFlat.h>
+#include <faiss/IndexBinaryHNSW.h>
+#include <faiss/IndexBinaryIVF.h>
+#include <faiss/IndexHNSW.h>
+#include <faiss/IndexIDMap.h>
+#include <faiss/IndexIVFFlat.h>
+#include <faiss/IndexIVFPQFastScan.h>
+#include <faiss/IndexLSH.h>
+#include <faiss/IndexNSG.h>
+#include <faiss/IndexPQFastScan.h>
+#include <faiss/IndexPreTransform.h>
+#include <faiss/IndexRefine.h>
+
 namespace faiss {
 
 namespace {
@@ -122,6 +136,11 @@ std::string reverse_index_factory(const faiss::Index* index) {
             const faiss::IndexHNSW* hnsw_index =
                     dynamic_cast<const faiss::IndexHNSW*>(index)) {
         return "HNSW" + std::to_string(get_hnsw_M(hnsw_index));
+    } else if (
+            const faiss::IndexNSG* nsg_index =
+                    dynamic_cast<const faiss::IndexNSG*>(index)) {
+        return "NSG" + std::to_string(nsg_index->nsg.R) + "," +
+                reverse_index_factory(nsg_index->storage);
     } else if (
             const faiss::IndexRefine* refine_index =
                     dynamic_cast<const faiss::IndexRefine*>(index)) {

data/vendor/faiss/faiss/cppcontrib/factory_tools.h

@@ -9,20 +9,13 @@
 
 #pragma once
 
-#include <faiss/IndexBinaryFlat.h>
-#include <faiss/IndexBinaryHNSW.h>
-#include <faiss/IndexBinaryIVF.h>
-#include <faiss/IndexHNSW.h>
-#include <faiss/IndexIDMap.h>
-#include <faiss/IndexIVFFlat.h>
-#include <faiss/IndexIVFPQFastScan.h>
-#include <faiss/IndexLSH.h>
-#include <faiss/IndexPQFastScan.h>
-#include <faiss/IndexPreTransform.h>
-#include <faiss/IndexRefine.h>
+#include <string>
 
 namespace faiss {
 
+struct Index;
+struct IndexBinary;
+
 std::string reverse_index_factory(const faiss::Index* index);
 std::string reverse_index_factory(const faiss::IndexBinary* index);
 

data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp

@@ -28,7 +28,8 @@ using namespace ::faiss;
  * Parameters to auto-tune on GpuIndex'es
  **********************************************************/
 
-#define DC(classname) auto ix = dynamic_cast<const classname*>(index)
+#define DC(classname) \
+    [[maybe_unused]] auto ix = dynamic_cast<const classname*>(index)
 
 void GpuParameterSpace::initialize(const Index* index) {
     if (DC(IndexPreTransform)) {

data/vendor/faiss/faiss/gpu/GpuIndexCagra.h

@@ -173,8 +173,8 @@ struct GpuIndexCagraConfig : public GpuIndexConfig {
     /// Number of Iterations to run if building with NN_DESCENT
     size_t nn_descent_niter = 20;
 
-    IVFPQBuildCagraConfig* ivf_pq_params = nullptr;
-    IVFPQSearchCagraConfig* ivf_pq_search_params = nullptr;
+    std::shared_ptr<IVFPQBuildCagraConfig> ivf_pq_params{nullptr};
+    std::shared_ptr<IVFPQSearchCagraConfig> ivf_pq_search_params{nullptr};
     float refine_rate = 2.0f;
     bool store_dataset = true;
 };
@@ -245,7 +245,17 @@ struct GpuIndexCagra : public GpuIndex {
             faiss::MetricType metric = faiss::METRIC_L2,
             GpuIndexCagraConfig config = GpuIndexCagraConfig());
 
-    /// Trains CAGRA based on the given vector data
+    /// Trains CAGRA based on the given vector data and add them along with ids.
+    /// NB: The use of the add function here is to build the CAGRA graph on
+    /// the base dataset. Use this function when you want to add vectors with
+    /// ids. Ref: https://github.com/facebookresearch/faiss/issues/4107
+    void add(idx_t n, const float* x) override;
+
+    /// Trains CAGRA based on the given vector data.
+    /// NB: The use of the train function here is to build the CAGRA graph on
+    /// the base dataset and is currently the only function to add the full set
+    /// of vectors (without IDs) to the index. There is no external quantizer to
+    /// be trained here.
     void train(idx_t n, const float* x) override;
 
     /// Initialize ourselves from the given CPU index; will overwrite

data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp

@@ -411,7 +411,7 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
     raftHandles_.emplace(std::make_pair(device, defaultStream));
 #endif
 
-    cudaStream_t asyncCopyStream = 0;
+    cudaStream_t asyncCopyStream = nullptr;
     CUDA_VERIFY(
             cudaStreamCreateWithFlags(&asyncCopyStream, cudaStreamNonBlocking));
 

data/vendor/faiss/faiss/gpu/StandardGpuResources.h

@@ -162,7 +162,7 @@ class StandardGpuResourcesImpl : public GpuResources {
     /**
      * FIXME: Integrating these in a separate code path for now. Ultimately,
      * it would be nice if we use a simple memory resource abstraction
-     * in FAISS so we could plug in whether to use RMM's memory resources
+     * in Faiss so we could plug in whether to use RMM's memory resources
      * or the default.
      *
      * There's enough duplicated logic that it doesn't *seem* to make sense

data/vendor/faiss/faiss/gpu/test/TestGpuIcmEncoder.cpp

@@ -0,0 +1,112 @@
+#include <faiss/gpu/GpuIcmEncoder.h>
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/gpu/test/TestUtils.h>
+#include <faiss/impl/LocalSearchQuantizer.h>
+
+#include <gtest/gtest.h>
+#include <tuple>
+#include <vector>
+
+using faiss::LocalSearchQuantizer;
+using faiss::gpu::GpuIcmEncoder;
+using faiss::gpu::GpuResourcesProvider;
+using faiss::gpu::StandardGpuResources;
+
+struct ShardingTestParams {
+    size_t n;
+    size_t nshards;
+};
+
+class GpuIcmEncoderShardingTest
+        : public ::testing::TestWithParam<ShardingTestParams> {
+   protected:
+    void SetUp() override {
+        params = GetParam();
+
+        lsq.M = 4;
+        lsq.K = 16;
+        lsq.d = 32;
+
+        std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
+        lsq.codebooks.resize(lsq.M * lsq.K * lsq.d);
+        for (auto& v : lsq.codebooks) {
+            v = dist(gen);
+        }
+
+        x.resize(params.n * lsq.d);
+        codes.resize(params.n * lsq.M);
+
+        for (auto& v : x) {
+            v = dist(gen);
+        }
+
+        std::uniform_int_distribution<int32_t> codeDist(0, lsq.K - 1);
+        for (auto& c : codes) {
+            c = codeDist(gen);
+        }
+    }
+
+    LocalSearchQuantizer lsq;
+    std::vector<float> x;
+    std::vector<int32_t> codes;
+    std::mt19937 gen;
+    ShardingTestParams params;
+    static constexpr size_t ils_iters = 4;
+};
+
+TEST_P(GpuIcmEncoderShardingTest, DataShardingCorrectness) {
+    std::vector<StandardGpuResources> resources(params.nshards);
+    std::vector<GpuResourcesProvider*> provs;
+    std::vector<int> devices;
+
+    for (size_t i = 0; i < params.nshards; ++i) {
+        resources[i].noTempMemory();
+        provs.push_back(&resources[i]);
+        devices.push_back(0); // use GPU 0 for testing all shards
+    }
+
+    GpuIcmEncoder encoder(&lsq, provs, devices);
+    encoder.set_binary_term();
+
+    gen.seed(42);
+    EXPECT_NO_THROW(
+            encoder.encode(codes.data(), x.data(), gen, params.n, ils_iters));
+
+    for (auto c : codes) {
+        EXPECT_GE(c, 0);
+        EXPECT_LT(c, lsq.K);
+    }
+}
+
+std::vector<ShardingTestParams> GetShardingTestCases() {
+    return {
+            {1, 8},
+
+            {5, 4},
+
+            {10, 2},
+            {10, 3},
+            {10, 5},
+            {10, 8},
+
+            {20, 8},
+    };
+}
+
+INSTANTIATE_TEST_SUITE_P(
+        MultiGpuShardingTests,
+        GpuIcmEncoderShardingTest,
+        ::testing::ValuesIn(GetShardingTestCases()),
+        [](const ::testing::TestParamInfo<ShardingTestParams>& info) {
+            return "n" + std::to_string(info.param.n) + "_shards" +
+                    std::to_string(info.param.nshards);
+        });
+
+int main(int argc, char** argv) {
+    testing::InitGoogleTest(&argc, argv);
+
+    // just run with a fixed test seed
+    faiss::gpu::setTestSeed(100);
+
+    return RUN_ALL_TESTS();
+}

data/vendor/faiss/faiss/impl/HNSW.cpp

@@ -8,7 +8,6 @@
 #include <faiss/impl/HNSW.h>
 
 #include <cstddef>
-#include <string>
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/DistanceComputer.h>
@@ -32,6 +31,7 @@ namespace faiss {
  **************************************************************/
 
 int HNSW::nb_neighbors(int layer_no) const {
+    FAISS_THROW_IF_NOT(layer_no + 1 < cum_nneighbor_per_level.size());
     return cum_nneighbor_per_level[layer_no + 1] -
             cum_nneighbor_per_level[layer_no];
 }
@@ -590,15 +590,22 @@ int search_from_candidates(
         HNSWStats& stats,
         int level,
         int nres_in,
-        const SearchParametersHNSW* params) {
+        const SearchParameters* params) {
     int nres = nres_in;
     int ndis = 0;
 
     // can be overridden by search params
-    bool do_dis_check = params ? params->check_relative_distance
-                               : hnsw.check_relative_distance;
-    int efSearch = params ? params->efSearch : hnsw.efSearch;
-    const IDSelector* sel = params ? params->sel : nullptr;
+    bool do_dis_check = hnsw.check_relative_distance;
+    int efSearch = hnsw.efSearch;
+    const IDSelector* sel = nullptr;
+    if (params) {
+        if (const SearchParametersHNSW* hnsw_params =
+                    dynamic_cast<const SearchParametersHNSW*>(params)) {
+            do_dis_check = hnsw_params->check_relative_distance;
+            efSearch = hnsw_params->efSearch;
+        }
+        sel = params->sel;
+    }
 
     C::T threshold = res.threshold;
     for (int i = 0; i < candidates.size(); i++) {
@@ -920,15 +927,22 @@ HNSWStats HNSW::search(
         DistanceComputer& qdis,
         ResultHandler<C>& res,
         VisitedTable& vt,
-        const SearchParametersHNSW* params) const {
+        const SearchParameters* params) const {
     HNSWStats stats;
     if (entry_point == -1) {
         return stats;
     }
     int k = extract_k_from_ResultHandler(res);
 
-    bool bounded_queue =
-            params ? params->bounded_queue : this->search_bounded_queue;
+    bool bounded_queue = this->search_bounded_queue;
+    int efSearch = this->efSearch;
+    if (params) {
+        if (const SearchParametersHNSW* hnsw_params =
+                    dynamic_cast<const SearchParametersHNSW*>(params)) {
+            bounded_queue = hnsw_params->bounded_queue;
+            efSearch = hnsw_params->efSearch;
+        }
+    }
 
     //  greedy search on upper levels
     storage_idx_t nearest = entry_point;
@@ -940,7 +954,7 @@ HNSWStats HNSW::search(
         stats.combine(local_stats);
     }
 
-    int ef = std::max(params ? params->efSearch : efSearch, k);
+    int ef = std::max(efSearch, k);
     if (bounded_queue) { // this is the most common branch
         MinimaxHeap candidates(ef);
 
@@ -980,9 +994,17 @@ void HNSW::search_level_0(
         int search_type,
         HNSWStats& search_stats,
         VisitedTable& vt,
-        const SearchParametersHNSW* params) const {
+        const SearchParameters* params) const {
     const HNSW& hnsw = *this;
-    auto efSearch = params ? params->efSearch : hnsw.efSearch;
+
+    auto efSearch = hnsw.efSearch;
+    if (params) {
+        if (const SearchParametersHNSW* hnsw_params =
+                    dynamic_cast<const SearchParametersHNSW*>(params)) {
+            efSearch = hnsw_params->efSearch;
+        }
+    }
+
     int k = extract_k_from_ResultHandler(res);
 
     if (search_type == 1) {
@@ -1062,7 +1084,7 @@ void HNSW::permute_entries(const idx_t* map) {
     // swap everyone
     std::swap(levels, new_levels);
     std::swap(offsets, new_offsets);
-    std::swap(neighbors, new_neighbors);
+    neighbors = std::move(new_neighbors);
 }
 
 /**************************************************************

data/vendor/faiss/faiss/impl/HNSW.h

@@ -15,6 +15,7 @@
 
 #include <faiss/Index.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/maybe_owned_vector.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/random.h>
@@ -121,7 +122,7 @@ struct HNSW {
 
     /// neighbors[offsets[i]:offsets[i+1]] is the list of neighbors of vector i
     /// for all levels. this is where all storage goes.
-    std::vector<storage_idx_t> neighbors;
+    MaybeOwnedVector<storage_idx_t> neighbors;
 
     /// entry point in the search structure (one of the points with maximum
     /// level
@@ -200,7 +201,7 @@ struct HNSW {
             DistanceComputer& qdis,
             ResultHandler<C>& res,
             VisitedTable& vt,
-            const SearchParametersHNSW* params = nullptr) const;
+            const SearchParameters* params = nullptr) const;
 
     /// search only in level 0 from a given vertex
     void search_level_0(
@@ -212,7 +213,7 @@ struct HNSW {
             int search_type,
             HNSWStats& search_stats,
             VisitedTable& vt,
-            const SearchParametersHNSW* params = nullptr) const;
+            const SearchParameters* params = nullptr) const;
 
     void reset();
 
@@ -264,7 +265,7 @@ int search_from_candidates(
         HNSWStats& stats,
         int level,
         int nres_in = 0,
-        const SearchParametersHNSW* params = nullptr);
+        const SearchParameters* params = nullptr);
 
 HNSWStats greedy_update_nearest(
         const HNSW& hnsw,

data/vendor/faiss/faiss/impl/NNDescent.cpp

@@ -400,7 +400,7 @@ void NNDescent::build(DistanceComputer& qdis, const int n, bool verbose) {
     init_graph(qdis);
     nndescent(qdis, verbose);
 
-    final_graph.resize(ntotal * K);
+    final_graph.resize(uint64_t(ntotal) * K);
 
     // Store the neighbor link structure into final_graph
     // Clear the old graph