faiss 0.3.2 → 0.3.3

data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -24,7 +24,6 @@
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
-#include <faiss/utils/utils.h>
 
 extern "C" {
 
@@ -152,6 +151,40 @@ void AdditiveQuantizer::train_norm(size_t n, const float* norms) {
     }
 }
 
+void AdditiveQuantizer::compute_codebook_tables() {
+    centroid_norms.resize(total_codebook_size);
+    fvec_norms_L2sqr(
+            centroid_norms.data(), codebooks.data(), d, total_codebook_size);
+    size_t cross_table_size = 0;
+    for (int m = 0; m < M; m++) {
+        size_t K = (size_t)1 << nbits[m];
+        cross_table_size += K * codebook_offsets[m];
+    }
+    codebook_cross_products.resize(cross_table_size);
+    size_t ofs = 0;
+    for (int m = 1; m < M; m++) {
+        FINTEGER ki = (size_t)1 << nbits[m];
+        FINTEGER kk = codebook_offsets[m];
+        FINTEGER di = d;
+        float zero = 0, one = 1;
+        assert(ofs + ki * kk <= cross_table_size);
+        sgemm_("Transposed",
+               "Not transposed",
+               &ki,
+               &kk,
+               &di,
+               &one,
+               codebooks.data() + d * kk,
+               &di,
+               codebooks.data(),
+               &di,
+               &zero,
+               codebook_cross_products.data() + ofs,
+               &ki);
+        ofs += ki * kk;
+    }
+}
+
 namespace {
 
 // TODO
@@ -471,7 +504,6 @@ namespace {
 float accumulate_IPs(
         const AdditiveQuantizer& aq,
         BitstringReader& bs,
-        const uint8_t* codes,
         const float* LUT) {
     float accu = 0;
     for (int m = 0; m < aq.M; m++) {
@@ -483,6 +515,29 @@ float accumulate_IPs(
     return accu;
 }
 
+float compute_norm_from_LUT(const AdditiveQuantizer& aq, BitstringReader& bs) {
+    float accu = 0;
+    std::vector<int> idx(aq.M);
+    const float* c = aq.codebook_cross_products.data();
+    for (int m = 0; m < aq.M; m++) {
+        size_t nbit = aq.nbits[m];
+        int i = bs.read(nbit);
+        size_t K = 1 << nbit;
+        idx[m] = i;
+
+        accu += aq.centroid_norms[aq.codebook_offsets[m] + i];
+
+        for (int l = 0; l < m; l++) {
+            int j = idx[l];
+            accu += 2 * c[j * K + i];
+            c += (1 << aq.nbits[l]) * K;
+        }
+    }
+    // FAISS_THROW_IF_NOT(c == aq.codebook_cross_products.data() +
+    // aq.codebook_cross_products.size());
+    return accu;
+}
+
 } // anonymous namespace
 
 template <>
@@ -491,7 +546,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    return accumulate_IPs(*this, bs, codes, LUT);
+    return accumulate_IPs(*this, bs, LUT);
 }
 
 template <>
@@ -500,7 +555,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    return -accumulate_IPs(*this, bs, codes, LUT);
+    return -accumulate_IPs(*this, bs, LUT);
 }
 
 template <>
@@ -509,7 +564,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    float accu = accumulate_IPs(*this, bs, LUT);
     uint32_t norm_i = bs.read(32);
     float norm2;
     memcpy(&norm2, &norm_i, 4);
@@ -522,7 +577,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    float accu = accumulate_IPs(*this, bs, LUT);
     uint32_t norm_i = bs.read(8);
     float norm2 = decode_qcint(norm_i);
     return norm2 - 2 * accu;
@@ -534,7 +589,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    float accu = accumulate_IPs(*this, bs, LUT);
     uint32_t norm_i = bs.read(4);
     float norm2 = decode_qcint(norm_i);
     return norm2 - 2 * accu;
@@ -546,7 +601,7 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    float accu = accumulate_IPs(*this, bs, LUT);
     uint32_t norm_i = bs.read(8);
     float norm2 = decode_qint8(norm_i, norm_min, norm_max);
     return norm2 - 2 * accu;
@@ -558,10 +613,23 @@ float AdditiveQuantizer::
                 const uint8_t* codes,
                 const float* LUT) const {
     BitstringReader bs(codes, code_size);
-    float accu = accumulate_IPs(*this, bs, codes, LUT);
+    float accu = accumulate_IPs(*this, bs, LUT);
     uint32_t norm_i = bs.read(4);
     float norm2 = decode_qint4(norm_i, norm_min, norm_max);
     return norm2 - 2 * accu;
 }
 
+template <>
+float AdditiveQuantizer::
+        compute_1_distance_LUT<false, AdditiveQuantizer::ST_norm_from_LUT>(
+                const uint8_t* codes,
+                const float* LUT) const {
+    FAISS_THROW_IF_NOT(codebook_cross_products.size() > 0);
+    BitstringReader bs(codes, code_size);
+    float accu = accumulate_IPs(*this, bs, LUT);
+    BitstringReader bs2(codes, code_size);
+    float norm2 = compute_norm_from_LUT(*this, bs2);
+    return norm2 - 2 * accu;
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/impl/AdditiveQuantizer.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -29,6 +29,8 @@ struct AdditiveQuantizer : Quantizer {
     std::vector<float> codebooks; ///< codebooks
 
     // derived values
+    /// codebook #1 is stored in rows codebook_offsets[i]:codebook_offsets[i+1]
+    /// in the codebooks table of size total_codebook_size by d
     std::vector<uint64_t> codebook_offsets;
     size_t tot_bits = 0;            ///< total number of bits (indexes + norms)
     size_t norm_bits = 0;           ///< bits allocated for the norms
@@ -38,9 +40,19 @@ struct AdditiveQuantizer : Quantizer {
     bool verbose = false;    ///< verbose during training?
     bool is_trained = false; ///< is trained or not
 
-    IndexFlat1D qnorm;            ///< store and search norms
-    std::vector<float> norm_tabs; ///< store norms of codebook entries for 4-bit
-                                  ///< fastscan search
+    /// auxiliary data for ST_norm_lsq2x4 and ST_norm_rq2x4
+    /// store norms of codebook entries for 4-bit fastscan
+    std::vector<float> norm_tabs;
+    IndexFlat1D qnorm; ///< store and search norms
+
+    void compute_codebook_tables();
+
+    /// norms of all codebook entries (size total_codebook_size)
+    std::vector<float> centroid_norms;
+
+    /// dot products of all codebook entries with the previous codebooks
+    /// size sum(codebook_offsets[m] * 2^nbits[m], m=0..M-1)
+    std::vector<float> codebook_cross_products;
 
     /// norms and distance matrixes with beam search can get large, so use this
     /// to control for the amount of memory that can be allocated

data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/AuxIndexStructures.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/CodePacker.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/CodePacker.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/DistanceComputer.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/FaissAssert.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/FaissException.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/FaissException.h

@@ -1,6 +1,5 @@
-
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -166,10 +166,10 @@ void HNSW::print_neighbor_stats(int level) const {
 }
 
 void HNSW::fill_with_random_links(size_t n) {
-    int max_level = prepare_level_tab(n);
+    int max_level_2 = prepare_level_tab(n);
     RandomGenerator rng2(456);
 
-    for (int level = max_level - 1; level >= 0; --level) {
+    for (int level = max_level_2 - 1; level >= 0; --level) {
         std::vector<int> elts;
         for (int i = 0; i < n; i++) {
             if (levels[i] > level) {
@@ -210,16 +210,16 @@ int HNSW::prepare_level_tab(size_t n, bool preset_levels) {
         }
     }
 
-    int max_level = 0;
+    int max_level_2 = 0;
     for (int i = 0; i < n; i++) {
         int pt_level = levels[i + n0] - 1;
-        if (pt_level > max_level)
-            max_level = pt_level;
+        if (pt_level > max_level_2)
+            max_level_2 = pt_level;
         offsets.push_back(offsets.back() + cum_nb_neighbors(pt_level + 1));
     }
     neighbors.resize(offsets.back(), -1);
 
-    return max_level;
+    return max_level_2;
 }
 
 /** Enumerate vertices from nearest to farthest from query, keep a
@@ -351,6 +351,8 @@ void add_link(
     }
 }
 
+} // namespace
+
 /// search neighbors on a single level, starting from an entry point
 void search_neighbors_to_add(
         HNSW& hnsw,
@@ -359,10 +361,8 @@ void search_neighbors_to_add(
         int entry_point,
         float d_entry_point,
         int level,
-        VisitedTable& vt) {
-    // selects a version
-    const bool reference_version = false;
-
+        VisitedTable& vt,
+        bool reference_version) {
     // top is nearest candidate
     std::priority_queue<NodeDistFarther> candidates;
 
@@ -385,7 +385,14 @@ void search_neighbors_to_add(
         size_t begin, end;
         hnsw.neighbor_range(currNode, level, &begin, &end);
 
-        // select a version, based on a flag
+        // The reference version is not used, but kept here because:
+        // 1. It is easier to switch back if the optimized version has a problem
+        // 2. It serves as a starting point for new optimizations
+        // 3. It helps understand the code
+        // 4. It ensures the reference version is still compilable if the
+        // optimized version changes
+        // The reference and the optimized versions' results are compared in
+        // test_hnsw.cpp
         if (reference_version) {
             // a reference version
             for (size_t i = begin; i < end; i++) {
@@ -470,8 +477,6 @@ void search_neighbors_to_add(
     vt.advance();
 }
 
-} // namespace
-
 /// Finds neighbors and builds links with them, starting from an entry
 /// point. The own neighbor list is assumed to be locked.
 void HNSW::add_links_starting_from(
@@ -493,17 +498,17 @@ void HNSW::add_links_starting_from(
 
     ::faiss::shrink_neighbor_list(ptdis, link_targets, M, keep_max_size_level0);
 
-    std::vector<storage_idx_t> neighbors;
-    neighbors.reserve(link_targets.size());
+    std::vector<storage_idx_t> neighbors_to_add;
+    neighbors_to_add.reserve(link_targets.size());
     while (!link_targets.empty()) {
         storage_idx_t other_id = link_targets.top().id;
         add_link(*this, ptdis, pt_id, other_id, level, keep_max_size_level0);
-        neighbors.push_back(other_id);
+        neighbors_to_add.push_back(other_id);
         link_targets.pop();
     }
 
     omp_unset_lock(&locks[pt_id]);
-    for (storage_idx_t other_id : neighbors) {
+    for (storage_idx_t other_id : neighbors_to_add) {
         omp_set_lock(&locks[other_id]);
         add_link(*this, ptdis, other_id, pt_id, level, keep_max_size_level0);
         omp_unset_lock(&locks[other_id]);

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -281,4 +281,14 @@ std::priority_queue<HNSW::Node> search_from_candidate_unbounded(
         VisitedTable* vt,
         HNSWStats& stats);
 
+void search_neighbors_to_add(
+        HNSW& hnsw,
+        DistanceComputer& qdis,
+        std::priority_queue<HNSW::NodeDistCloser>& results,
+        int entry_point,
+        float d_entry_point,
+        int level,
+        VisitedTable& vt,
+        bool reference_version = false);
+
 } // namespace faiss

data/vendor/faiss/faiss/impl/IDSelector.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/IDSelector.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/LookupTableScaler.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

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

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/NNDescent.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/impl/NSG.cpp

@@ -1,12 +1,10 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and 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/impl/NSG.h>
 
 #include <algorithm>
@@ -18,14 +16,16 @@
 
 namespace faiss {
 
-namespace nsg {
-
 namespace {
 
+using LockGuard = std::lock_guard<std::mutex>;
+
 // It needs to be smaller than 0
 constexpr int EMPTY_ID = -1;
 
-} // namespace
+} // anonymous namespace
+
+namespace nsg {
 
 DistanceComputer* storage_distance_computer(const Index* storage) {
     if (is_similarity_metric(storage->metric_type)) {
@@ -35,14 +35,8 @@ DistanceComputer* storage_distance_computer(const Index* storage) {
     }
 }
 
-} // namespace nsg
-
-using namespace nsg;
-
-using LockGuard = std::lock_guard<std::mutex>;
-
 struct Neighbor {
-    int id;
+    int32_t id;
     float distance;
     bool flag;
 
@@ -56,7 +50,7 @@ struct Neighbor {
 };
 
 struct Node {
-    int id;
+    int32_t id;
     float distance;
 
     Node() = default;
@@ -65,6 +59,11 @@ struct Node {
     inline bool operator<(const Node& other) const {
         return distance < other.distance;
     }
+
+    // to keep the compiler happy
+    inline bool operator<(int other) const {
+        return id < other;
+    }
 };
 
 inline int insert_into_pool(Neighbor* addr, int K, Neighbor nn) {
@@ -106,6 +105,10 @@ inline int insert_into_pool(Neighbor* addr, int K, Neighbor nn) {
     return right;
 }
 
+} // namespace nsg
+
+using namespace nsg;
+
 NSG::NSG(int R) : R(R), rng(0x0903) {
     L = R + 32;
     C = R + 100;
@@ -253,9 +256,11 @@ void NSG::search_on_graph(
     std::vector<int> init_ids(pool_size);
 
     int num_ids = 0;
-    for (int i = 0; i < init_ids.size() && i < graph.K; i++) {
-        int id = (int)graph.at(ep, i);
-        if (id < 0 || id >= ntotal) {
+    std::vector<index_t> neighbors(graph.K);
+    size_t nneigh = graph.get_neighbors(ep, neighbors.data());
+    for (int i = 0; i < init_ids.size() && i < nneigh; i++) {
+        int id = (int)neighbors[i];
+        if (id >= ntotal) {
             continue;
         }
 
@@ -296,9 +301,10 @@ void NSG::search_on_graph(
             retset[k].flag = false;
             int n = retset[k].id;
 
-            for (int m = 0; m < graph.K; m++) {
-                int id = (int)graph.at(n, m);
-                if (id < 0 || id > ntotal || vt.get(id)) {
+            size_t nneigh_for_n = graph.get_neighbors(n, neighbors.data());
+            for (int m = 0; m < nneigh_for_n; m++) {
+                int id = neighbors[m];
+                if (id > ntotal || vt.get(id)) {
                     continue;
                 }
                 vt.set(id);

data/vendor/faiss/faiss/impl/NSG.h

@@ -1,12 +1,10 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #pragma once
 
 #include <memory>
@@ -40,11 +38,12 @@ namespace faiss {
  */
 
 struct DistanceComputer; // from AuxIndexStructures
-struct Neighbor;
-struct Node;
 
 namespace nsg {
 
+struct Neighbor;
+struct Node;
+
 /***********************************************************
  * Graph structure to store a graph.
  *
@@ -75,7 +74,7 @@ struct Graph {
     }
 
     // release the allocated memory if needed
-    ~Graph() {
+    virtual ~Graph() {
         if (own_fields) {
             delete[] data;
         }
@@ -90,6 +89,17 @@ struct Graph {
     inline node_t& at(int i, int j) {
         return data[i * K + j];
     }
+
+    // get all neighbors of node i (used during search only)
+    virtual size_t get_neighbors(int i, node_t* neighbors) const {
+        for (int j = 0; j < K; j++) {
+            if (data[i * K + j] < 0) {
+                return j;
+            }
+            neighbors[j] = data[i * K + j];
+        }
+        return K;
+    }
 };
 
 DistanceComputer* storage_distance_computer(const Index* storage);
@@ -99,6 +109,8 @@ DistanceComputer* storage_distance_computer(const Index* storage);
 struct NSG {
     /// internal storage of vectors (32 bits: this is expensive)
     using storage_idx_t = int32_t;
+    using Node = nsg::Node;
+    using Neighbor = nsg::Neighbor;
 
     int ntotal = 0; ///< nb of nodes
 
@@ -112,7 +124,7 @@ struct NSG {
 
     int enterpoint; ///< enterpoint
 
-    std::shared_ptr<nsg::Graph<int>> final_graph; ///< NSG graph structure
+    std::shared_ptr<nsg::Graph<int32_t>> final_graph; ///< NSG graph structure
 
     bool is_built = false; ///< NSG is built or not
 

data/vendor/faiss/faiss/impl/PolysemousTraining.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.