faiss 0.2.7 → 0.3.1

data/vendor/faiss/faiss/IndexFlat.h

@@ -18,7 +18,9 @@ namespace faiss {
 
 /** Index that stores the full vectors and performs exhaustive search */
 struct IndexFlat : IndexFlatCodes {
-    explicit IndexFlat(idx_t d, MetricType metric = METRIC_L2);
+    explicit IndexFlat(
+            idx_t d, ///< dimensionality of the input vectors
+            MetricType metric = METRIC_L2);
 
     void search(
             idx_t n,
@@ -76,8 +78,25 @@ struct IndexFlatIP : IndexFlat {
 };
 
 struct IndexFlatL2 : IndexFlat {
+    // Special cache for L2 norms.
+    // If this cache is set, then get_distance_computer() returns
+    // a special version that computes the distance using dot products
+    // and l2 norms.
+    std::vector<float> cached_l2norms;
+
+    /**
+     * @param d dimensionality of the input vectors
+     */
     explicit IndexFlatL2(idx_t d) : IndexFlat(d, METRIC_L2) {}
     IndexFlatL2() {}
+
+    // override for l2 norms cache.
+    FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
+
+    // compute L2 norms
+    void sync_l2norms();
+    // clear L2 norms
+    void clear_l2norms();
 };
 
 /// optimized version for 1D "vectors".

data/vendor/faiss/faiss/IndexFlatCodes.cpp

@@ -103,4 +103,15 @@ CodePacker* IndexFlatCodes::get_CodePacker() const {
     return new CodePackerFlat(code_size);
 }
 
+void IndexFlatCodes::permute_entries(const idx_t* perm) {
+    std::vector<uint8_t> new_codes(codes.size());
+
+    for (idx_t i = 0; i < ntotal; i++) {
+        memcpy(new_codes.data() + i * code_size,
+               codes.data() + perm[i] * code_size,
+               code_size);
+    }
+    std::swap(codes, new_codes);
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/IndexFlatCodes.h

@@ -34,7 +34,6 @@ struct IndexFlatCodes : Index {
 
     void reset() override;
 
-    /// reconstruction using the codec interface
     void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
 
     void reconstruct(idx_t key, float* recons) const override;
@@ -59,6 +58,9 @@ struct IndexFlatCodes : Index {
     void check_compatible_for_merge(const Index& otherIndex) const override;
 
     virtual void merge_from(Index& otherIndex, idx_t add_id = 0) override;
+
+    // permute_entries. perm of size ntotal maps new to old positions
+    void permute_entries(const idx_t* perm);
 };
 
 } // namespace faiss

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -20,16 +20,16 @@
 #include <queue>
 #include <unordered_set>
 
-#include <stdint.h>
 #include <sys/stat.h>
 #include <sys/types.h>
+#include <cstdint>
 
 #include <faiss/Index2Layer.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/utils/Heap.h>
+#include <faiss/impl/ResultHandler.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/random.h>
 #include <faiss/utils/sorting.h>
@@ -87,6 +87,23 @@ struct NegativeDistanceComputer : DistanceComputer {
         return -(*basedis)(i);
     }
 
+    void distances_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) override {
+        basedis->distances_batch_4(
+                idx0, idx1, idx2, idx3, dis0, dis1, dis2, dis3);
+        dis0 = -dis0;
+        dis1 = -dis1;
+        dis2 = -dis2;
+        dis3 = -dis3;
+    }
+
     /// compute distance between two stored vectors
     float symmetric_dis(idx_t i, idx_t j) override {
         return -basedis->symmetric_dis(i, j);
@@ -192,9 +209,8 @@ void hnsw_add_vertices(
             {
                 VisitedTable vt(ntotal);
 
-                DistanceComputer* dis =
-                        storage_distance_computer(index_hnsw.storage);
-                ScopeDeleter1<DistanceComputer> del(dis);
+                std::unique_ptr<DistanceComputer> dis(
+                        storage_distance_computer(index_hnsw.storage));
                 int prev_display =
                         verbose && omp_get_thread_num() == 0 ? 0 : -1;
                 size_t counter = 0;
@@ -250,18 +266,10 @@ void hnsw_add_vertices(
  **************************************************************/
 
 IndexHNSW::IndexHNSW(int d, int M, MetricType metric)
-        : Index(d, metric),
-          hnsw(M),
-          own_fields(false),
-          storage(nullptr),
-          reconstruct_from_neighbors(nullptr) {}
+        : Index(d, metric), hnsw(M) {}
 
 IndexHNSW::IndexHNSW(Index* storage, int M)
-        : Index(storage->d, storage->metric_type),
-          hnsw(M),
-          own_fields(false),
-          storage(storage),
-          reconstruct_from_neighbors(nullptr) {}
+        : Index(storage->d, storage->metric_type), hnsw(M), storage(storage) {}
 
 IndexHNSW::~IndexHNSW() {
     if (own_fields) {
@@ -278,18 +286,20 @@ void IndexHNSW::train(idx_t n, const float* x) {
     is_trained = true;
 }
 
-void IndexHNSW::search(
+namespace {
+
+template <class BlockResultHandler>
+void hnsw_search(
+        const IndexHNSW* index,
         idx_t n,
         const float* x,
-        idx_t k,
-        float* distances,
-        idx_t* labels,
-        const SearchParameters* params_in) const {
-    FAISS_THROW_IF_NOT(k > 0);
+        BlockResultHandler& bres,
+        const SearchParameters* params_in) {
     FAISS_THROW_IF_NOT_MSG(
-            storage,
+            index->storage,
             "Please use IndexHNSWFlat (or variants) instead of IndexHNSW directly");
     const SearchParametersHNSW* params = nullptr;
+    const HNSW& hnsw = index->hnsw;
 
     int efSearch = hnsw.efSearch;
     if (params_in) {
@@ -299,61 +309,81 @@ void IndexHNSW::search(
     }
     size_t n1 = 0, n2 = 0, n3 = 0, ndis = 0, nreorder = 0;
 
-    idx_t check_period =
-            InterruptCallback::get_period_hint(hnsw.max_level * d * efSearch);
+    idx_t check_period = InterruptCallback::get_period_hint(
+            hnsw.max_level * index->d * efSearch);
 
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
 
 #pragma omp parallel
         {
-            VisitedTable vt(ntotal);
+            VisitedTable vt(index->ntotal);
+            typename BlockResultHandler::SingleResultHandler res(bres);
 
-            DistanceComputer* dis = storage_distance_computer(storage);
-            ScopeDeleter1<DistanceComputer> del(dis);
+            std::unique_ptr<DistanceComputer> dis(
+                    storage_distance_computer(index->storage));
 
-#pragma omp for reduction(+ : n1, n2, n3, ndis, nreorder)
+#pragma omp for reduction(+ : n1, n2, n3, ndis, nreorder) schedule(guided)
             for (idx_t i = i0; i < i1; i++) {
-                idx_t* idxi = labels + i * k;
-                float* simi = distances + i * k;
-                dis->set_query(x + i * d);
+                res.begin(i);
+                dis->set_query(x + i * index->d);
 
-                maxheap_heapify(k, simi, idxi);
-                HNSWStats stats = hnsw.search(*dis, k, idxi, simi, vt, params);
+                HNSWStats stats = hnsw.search(*dis, res, vt, params);
                 n1 += stats.n1;
                 n2 += stats.n2;
                 n3 += stats.n3;
                 ndis += stats.ndis;
                 nreorder += stats.nreorder;
-                maxheap_reorder(k, simi, idxi);
-
-                if (reconstruct_from_neighbors &&
-                    reconstruct_from_neighbors->k_reorder != 0) {
-                    int k_reorder = reconstruct_from_neighbors->k_reorder;
-                    if (k_reorder == -1 || k_reorder > k)
-                        k_reorder = k;
-
-                    nreorder += reconstruct_from_neighbors->compute_distances(
-                            k_reorder, idxi, x + i * d, simi);
-
-                    // sort top k_reorder
-                    maxheap_heapify(
-                            k_reorder, simi, idxi, simi, idxi, k_reorder);
-                    maxheap_reorder(k_reorder, simi, idxi);
-                }
+                res.end();
             }
         }
         InterruptCallback::check();
     }
 
-    if (is_similarity_metric(metric_type)) {
+    hnsw_stats.combine({n1, n2, n3, ndis, nreorder});
+}
+
+} // anonymous namespace
+
+void IndexHNSW::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params_in) const {
+    FAISS_THROW_IF_NOT(k > 0);
+
+    using RH = HeapBlockResultHandler<HNSW::C>;
+    RH bres(n, distances, labels, k);
+
+    hnsw_search(this, n, x, bres, params_in);
+
+    if (is_similarity_metric(this->metric_type)) {
         // we need to revert the negated distances
         for (size_t i = 0; i < k * n; i++) {
             distances[i] = -distances[i];
         }
     }
+}
 
-    hnsw_stats.combine({n1, n2, n3, ndis, nreorder});
+void IndexHNSW::range_search(
+        idx_t n,
+        const float* x,
+        float radius,
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    using RH = RangeSearchBlockResultHandler<HNSW::C>;
+    RH bres(result, radius);
+
+    hnsw_search(this, n, x, bres, params);
+
+    if (is_similarity_metric(this->metric_type)) {
+        // we need to revert the negated distances
+        for (size_t i = 0; i < result->lims[result->nq]; i++) {
+            result->distances[i] = -result->distances[i];
+        }
+    }
 }
 
 void IndexHNSW::add(idx_t n, const float* x) {
@@ -381,8 +411,8 @@ void IndexHNSW::reconstruct(idx_t key, float* recons) const {
 void IndexHNSW::shrink_level_0_neighbors(int new_size) {
 #pragma omp parallel
     {
-        DistanceComputer* dis = storage_distance_computer(storage);
-        ScopeDeleter1<DistanceComputer> del(dis);
+        std::unique_ptr<DistanceComputer> dis(
+                storage_distance_computer(storage));
 
 #pragma omp for
         for (idx_t i = 0; i < ntotal; i++) {
@@ -429,35 +459,33 @@ void IndexHNSW::search_level_0(
 
     storage_idx_t ntotal = hnsw.levels.size();
 
+    using RH = HeapBlockResultHandler<HNSW::C>;
+    RH bres(n, distances, labels, k);
+
 #pragma omp parallel
     {
         std::unique_ptr<DistanceComputer> qdis(
                 storage_distance_computer(storage));
         HNSWStats search_stats;
         VisitedTable vt(ntotal);
+        RH::SingleResultHandler res(bres);
 
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            idx_t* idxi = labels + i * k;
-            float* simi = distances + i * k;
-
+            res.begin(i);
             qdis->set_query(x + i * d);
-            maxheap_heapify(k, simi, idxi);
 
             hnsw.search_level_0(
                     *qdis.get(),
-                    k,
-                    idxi,
-                    simi,
+                    res,
                     nprobe,
                     nearest + i * nprobe,
                     nearest_d + i * nprobe,
                     search_type,
                     search_stats,
                     vt);
-
+            res.end();
             vt.advance();
-            maxheap_reorder(k, simi, idxi);
         }
 #pragma omp critical
         { hnsw_stats.combine(search_stats); }
@@ -515,8 +543,8 @@ void IndexHNSW::init_level_0_from_entry_points(
     {
         VisitedTable vt(ntotal);
 
-        DistanceComputer* dis = storage_distance_computer(storage);
-        ScopeDeleter1<DistanceComputer> del(dis);
+        std::unique_ptr<DistanceComputer> dis(
+                storage_distance_computer(storage));
         std::vector<float> vec(storage->d);
 
 #pragma omp for schedule(dynamic)
@@ -551,8 +579,8 @@ void IndexHNSW::reorder_links() {
         std::vector<float> distances(M);
         std::vector<size_t> order(M);
         std::vector<storage_idx_t> tmp(M);
-        DistanceComputer* dis = storage_distance_computer(storage);
-        ScopeDeleter1<DistanceComputer> del(dis);
+        std::unique_ptr<DistanceComputer> dis(
+                storage_distance_computer(storage));
 
 #pragma omp for
         for (storage_idx_t i = 0; i < ntotal; i++) {
@@ -614,245 +642,12 @@ void IndexHNSW::link_singletons() {
     }
 }
 
-/**************************************************************
- * ReconstructFromNeighbors implementation
- **************************************************************/
-
-ReconstructFromNeighbors::ReconstructFromNeighbors(
-        const IndexHNSW& index,
-        size_t k,
-        size_t nsq)
-        : index(index), k(k), nsq(nsq) {
-    M = index.hnsw.nb_neighbors(0);
-    FAISS_ASSERT(k <= 256);
-    code_size = k == 1 ? 0 : nsq;
-    ntotal = 0;
-    d = index.d;
-    FAISS_ASSERT(d % nsq == 0);
-    dsub = d / nsq;
-    k_reorder = -1;
-}
-
-void ReconstructFromNeighbors::reconstruct(
-        storage_idx_t i,
-        float* x,
-        float* tmp) const {
-    const HNSW& hnsw = index.hnsw;
-    size_t begin, end;
-    hnsw.neighbor_range(i, 0, &begin, &end);
-
-    if (k == 1 || nsq == 1) {
-        const float* beta;
-        if (k == 1) {
-            beta = codebook.data();
-        } else {
-            int idx = codes[i];
-            beta = codebook.data() + idx * (M + 1);
-        }
-
-        float w0 = beta[0]; // weight of image itself
-        index.storage->reconstruct(i, tmp);
-
-        for (int l = 0; l < d; l++)
-            x[l] = w0 * tmp[l];
-
-        for (size_t j = begin; j < end; j++) {
-            storage_idx_t ji = hnsw.neighbors[j];
-            if (ji < 0)
-                ji = i;
-            float w = beta[j - begin + 1];
-            index.storage->reconstruct(ji, tmp);
-            for (int l = 0; l < d; l++)
-                x[l] += w * tmp[l];
-        }
-    } else if (nsq == 2) {
-        int idx0 = codes[2 * i];
-        int idx1 = codes[2 * i + 1];
-
-        const float* beta0 = codebook.data() + idx0 * (M + 1);
-        const float* beta1 = codebook.data() + (idx1 + k) * (M + 1);
-
-        index.storage->reconstruct(i, tmp);
-
-        float w0;
-
-        w0 = beta0[0];
-        for (int l = 0; l < dsub; l++)
-            x[l] = w0 * tmp[l];
-
-        w0 = beta1[0];
-        for (int l = dsub; l < d; l++)
-            x[l] = w0 * tmp[l];
-
-        for (size_t j = begin; j < end; j++) {
-            storage_idx_t ji = hnsw.neighbors[j];
-            if (ji < 0)
-                ji = i;
-            index.storage->reconstruct(ji, tmp);
-            float w;
-            w = beta0[j - begin + 1];
-            for (int l = 0; l < dsub; l++)
-                x[l] += w * tmp[l];
-
-            w = beta1[j - begin + 1];
-            for (int l = dsub; l < d; l++)
-                x[l] += w * tmp[l];
-        }
-    } else {
-        std::vector<const float*> betas(nsq);
-        {
-            const float* b = codebook.data();
-            const uint8_t* c = &codes[i * code_size];
-            for (int sq = 0; sq < nsq; sq++) {
-                betas[sq] = b + (*c++) * (M + 1);
-                b += (M + 1) * k;
-            }
-        }
-
-        index.storage->reconstruct(i, tmp);
-        {
-            int d0 = 0;
-            for (int sq = 0; sq < nsq; sq++) {
-                float w = *(betas[sq]++);
-                int d1 = d0 + dsub;
-                for (int l = d0; l < d1; l++) {
-                    x[l] = w * tmp[l];
-                }
-                d0 = d1;
-            }
-        }
-
-        for (size_t j = begin; j < end; j++) {
-            storage_idx_t ji = hnsw.neighbors[j];
-            if (ji < 0)
-                ji = i;
-
-            index.storage->reconstruct(ji, tmp);
-            int d0 = 0;
-            for (int sq = 0; sq < nsq; sq++) {
-                float w = *(betas[sq]++);
-                int d1 = d0 + dsub;
-                for (int l = d0; l < d1; l++) {
-                    x[l] += w * tmp[l];
-                }
-                d0 = d1;
-            }
-        }
-    }
-}
-
-void ReconstructFromNeighbors::reconstruct_n(
-        storage_idx_t n0,
-        storage_idx_t ni,
-        float* x) const {
-#pragma omp parallel
-    {
-        std::vector<float> tmp(index.d);
-#pragma omp for
-        for (storage_idx_t i = 0; i < ni; i++) {
-            reconstruct(n0 + i, x + i * index.d, tmp.data());
-        }
-    }
-}
-
-size_t ReconstructFromNeighbors::compute_distances(
-        size_t n,
-        const idx_t* shortlist,
-        const float* query,
-        float* distances) const {
-    std::vector<float> tmp(2 * index.d);
-    size_t ncomp = 0;
-    for (int i = 0; i < n; i++) {
-        if (shortlist[i] < 0)
-            break;
-        reconstruct(shortlist[i], tmp.data(), tmp.data() + index.d);
-        distances[i] = fvec_L2sqr(query, tmp.data(), index.d);
-        ncomp++;
-    }
-    return ncomp;
-}
-
-void ReconstructFromNeighbors::get_neighbor_table(storage_idx_t i, float* tmp1)
-        const {
-    const HNSW& hnsw = index.hnsw;
-    size_t begin, end;
-    hnsw.neighbor_range(i, 0, &begin, &end);
-    size_t d = index.d;
-
-    index.storage->reconstruct(i, tmp1);
-
-    for (size_t j = begin; j < end; j++) {
-        storage_idx_t ji = hnsw.neighbors[j];
-        if (ji < 0)
-            ji = i;
-        index.storage->reconstruct(ji, tmp1 + (j - begin + 1) * d);
-    }
-}
-
-/// called by add_codes
-void ReconstructFromNeighbors::estimate_code(
-        const float* x,
-        storage_idx_t i,
-        uint8_t* code) const {
-    // fill in tmp table with the neighbor values
-    float* tmp1 = new float[d * (M + 1) + (d * k)];
-    float* tmp2 = tmp1 + d * (M + 1);
-    ScopeDeleter<float> del(tmp1);
-
-    // collect coordinates of base
-    get_neighbor_table(i, tmp1);
-
-    for (size_t sq = 0; sq < nsq; sq++) {
-        int d0 = sq * dsub;
-
-        {
-            FINTEGER ki = k, di = d, m1 = M + 1;
-            FINTEGER dsubi = dsub;
-            float zero = 0, one = 1;
-
-            sgemm_("N",
-                   "N",
-                   &dsubi,
-                   &ki,
-                   &m1,
-                   &one,
-                   tmp1 + d0,
-                   &di,
-                   codebook.data() + sq * (m1 * k),
-                   &m1,
-                   &zero,
-                   tmp2,
-                   &dsubi);
-        }
-
-        float min = HUGE_VAL;
-        int argmin = -1;
-        for (size_t j = 0; j < k; j++) {
-            float dis = fvec_L2sqr(x + d0, tmp2 + j * dsub, dsub);
-            if (dis < min) {
-                min = dis;
-                argmin = j;
-            }
-        }
-        code[sq] = argmin;
-    }
-}
-
-void ReconstructFromNeighbors::add_codes(size_t n, const float* x) {
-    if (k == 1) { // nothing to encode
-        ntotal += n;
-        return;
-    }
-    codes.resize(codes.size() + code_size * n);
-#pragma omp parallel for
-    for (int i = 0; i < n; i++) {
-        estimate_code(
-                x + i * index.d,
-                ntotal + i,
-                codes.data() + (ntotal + i) * code_size);
-    }
-    ntotal += n;
-    FAISS_ASSERT(codes.size() == ntotal * code_size);
+void IndexHNSW::permute_entries(const idx_t* perm) {
+    auto flat_storage = dynamic_cast<IndexFlatCodes*>(storage);
+    FAISS_THROW_IF_NOT_MSG(
+            flat_storage, "don't know how to permute this index");
+    flat_storage->permute_entries(perm);
+    hnsw.permute_entries(perm);
 }
 
 /**************************************************************
@@ -864,7 +659,10 @@ IndexHNSWFlat::IndexHNSWFlat() {
 }
 
 IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
-        : IndexHNSW(new IndexFlat(d, metric), M) {
+        : IndexHNSW(
+                  (metric == METRIC_L2) ? new IndexFlatL2(d)
+                                        : new IndexFlat(d, metric),
+                  M) {
     own_fields = true;
     is_trained = true;
 }
@@ -873,10 +671,10 @@ IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
  * IndexHNSWPQ implementation
  **************************************************************/
 
-IndexHNSWPQ::IndexHNSWPQ() {}
+IndexHNSWPQ::IndexHNSWPQ() = default;
 
-IndexHNSWPQ::IndexHNSWPQ(int d, int pq_m, int M)
-        : IndexHNSW(new IndexPQ(d, pq_m, 8), M) {
+IndexHNSWPQ::IndexHNSWPQ(int d, int pq_m, int M, int pq_nbits)
+        : IndexHNSW(new IndexPQ(d, pq_m, pq_nbits), M) {
     own_fields = true;
     is_trained = false;
 }
@@ -896,11 +694,11 @@ IndexHNSWSQ::IndexHNSWSQ(
         int M,
         MetricType metric)
         : IndexHNSW(new IndexScalarQuantizer(d, qtype, metric), M) {
-    is_trained = false;
+    is_trained = this->storage->is_trained;
     own_fields = true;
 }
 
-IndexHNSWSQ::IndexHNSWSQ() {}
+IndexHNSWSQ::IndexHNSWSQ() = default;
 
 /**************************************************************
  * IndexHNSW2Level implementation
@@ -916,7 +714,7 @@ IndexHNSW2Level::IndexHNSW2Level(
     is_trained = false;
 }
 
-IndexHNSW2Level::IndexHNSW2Level() {}
+IndexHNSW2Level::IndexHNSW2Level() = default;
 
 namespace {
 
@@ -935,7 +733,6 @@ int search_from_candidates_2(
         int level,
         int nres_in = 0) {
     int nres = nres_in;
-    int ndis = 0;
     for (int i = 0; i < candidates.size(); i++) {
         idx_t v1 = candidates.ids[i];
         FAISS_ASSERT(v1 >= 0);
@@ -958,7 +755,6 @@ int search_from_candidates_2(
             if (vt.visited[v1] == vt.visno + 1) {
                 // nothing to do
             } else {
-                ndis++;
                 float d = qdis(v1);
                 candidates.push(v1, d);
 
@@ -1030,8 +826,8 @@ void IndexHNSW2Level::search(
 #pragma omp parallel
         {
             VisitedTable vt(ntotal);
-            DistanceComputer* dis = storage_distance_computer(storage);
-            ScopeDeleter1<DistanceComputer> del(dis);
+            std::unique_ptr<DistanceComputer> dis(
+                    storage_distance_computer(storage));
 
             int candidates_size = hnsw.upper_beam;
             MinimaxHeap candidates(candidates_size);