faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/IndexLSH.cpp

@@ -22,18 +22,22 @@ namespace faiss {
  * IndexLSH
  ***************************************************************/
 
-IndexLSH::IndexLSH(idx_t d, int nbits, bool rotate_data, bool train_thresholds)
-        : IndexFlatCodes((nbits + 7) / 8, d),
-          nbits(nbits),
-          rotate_data(rotate_data),
-          train_thresholds(train_thresholds),
-          rrot(d, nbits) {
-    is_trained = !train_thresholds;
-
-    if (rotate_data) {
+IndexLSH::IndexLSH(
+        idx_t d_in,
+        int nbits_in,
+        bool rotate_data_in,
+        bool train_thresholds_in)
+        : IndexFlatCodes((nbits_in + 7) / 8, d_in),
+          nbits(nbits_in),
+          rotate_data(rotate_data_in),
+          train_thresholds(train_thresholds_in),
+          rrot(static_cast<int>(d_in), nbits_in) {
+    is_trained = !train_thresholds_in;
+
+    if (rotate_data_in) {
         rrot.init(5);
     } else {
-        FAISS_THROW_IF_NOT(d >= nbits);
+        FAISS_THROW_IF_NOT(d_in >= nbits_in);
     }
 }
 
@@ -45,13 +49,18 @@ const float* IndexLSH::apply_preprocess(idx_t n, const float* x) const {
         // also applies bias if exists
         xt = rrot.apply(n, x);
     } else if (d != nbits) {
-        assert(nbits < d);
+        FAISS_THROW_IF_NOT_FMT(
+                nbits < d,
+                "nbits (%d) must be less than d (%d)",
+                nbits,
+                (int)d);
         xt = new float[nbits * n];
         float* xp = xt;
         for (idx_t i = 0; i < n; i++) {
             const float* xl = x + i * d;
-            for (int j = 0; j < nbits; j++)
+            for (int j = 0; j < nbits; j++) {
                 *xp++ = xl[j];
+            }
         }
     }
 
@@ -62,9 +71,11 @@ const float* IndexLSH::apply_preprocess(idx_t n, const float* x) const {
         }
 
         float* xp = xt;
-        for (idx_t i = 0; i < n; i++)
-            for (int j = 0; j < nbits; j++)
+        for (idx_t i = 0; i < n; i++) {
+            for (int j = 0; j < nbits; j++) {
                 *xp++ -= thresholds[j];
+            }
+        }
     }
 
     return xt ? xt : x;
@@ -80,18 +91,22 @@ void IndexLSH::train(idx_t n, const float* x) {
 
         std::unique_ptr<float[]> transposed_x(new float[n * nbits]);
 
-        for (idx_t i = 0; i < n; i++)
-            for (idx_t j = 0; j < nbits; j++)
+        for (idx_t i = 0; i < n; i++) {
+            for (idx_t j = 0; j < nbits; j++) {
                 transposed_x[j * n + i] = xt[i * nbits + j];
+            }
+        }
 
         for (idx_t i = 0; i < nbits; i++) {
             float* xi = transposed_x.get() + i * n;
-            // std::nth_element
-            std::sort(xi, xi + n);
-            if (n % 2 == 1)
-                thresholds[i] = xi[n / 2];
-            else
-                thresholds[i] = (xi[n / 2 - 1] + xi[n / 2]) / 2;
+            // Use nth_element (O(n)) instead of sort (O(n log n))
+            std::nth_element(xi, xi + n / 2, xi + n);
+            float median = xi[n / 2];
+            if (n % 2 == 0) {
+                std::nth_element(xi, xi + n / 2 - 1, xi + n);
+                median = (median + xi[n / 2 - 1]) / 2;
+            }
+            thresholds[i] = median;
         }
     }
     is_trained = true;
@@ -122,20 +137,24 @@ void IndexLSH::search(
     hammings_knn_hc(&res, qcodes.get(), codes.data(), ntotal, code_size, true);
 
     // convert distances to floats
-    for (int i = 0; i < k * n; i++)
+    for (int i = 0; i < k * n; i++) {
         distances[i] = idistances[i];
+    }
 }
 
 void IndexLSH::transfer_thresholds(LinearTransform* vt) {
-    if (!train_thresholds)
+    if (!train_thresholds) {
         return;
+    }
     FAISS_THROW_IF_NOT(nbits == vt->d_out);
     if (!vt->have_bias) {
         vt->b.resize(nbits, 0);
         vt->have_bias = true;
     }
-    for (int i = 0; i < nbits; i++)
+    FAISS_THROW_IF_NOT(!vt->b.empty());
+    for (int i = 0; i < nbits; i++) {
         vt->b[i] -= thresholds[i];
+    }
     train_thresholds = false;
     thresholds.clear();
 }

data/vendor/faiss/faiss/IndexLattice.cpp

@@ -9,25 +9,26 @@
 
 #include <faiss/IndexLattice.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h> // for the bitstring routines
 
 namespace faiss {
 
-IndexLattice::IndexLattice(idx_t d, int nsq, int scale_nbit, int r2)
-        : IndexFlatCodes(0, d, METRIC_L2),
-          nsq(nsq),
-          dsq(d / nsq),
-          zn_sphere_codec(dsq, r2),
-          scale_nbit(scale_nbit) {
-    FAISS_THROW_IF_NOT(d % nsq == 0);
+IndexLattice::IndexLattice(idx_t d_in, int nsq_in, int scale_nbit_in, int r2)
+        : IndexFlatCodes(0, d_in, METRIC_L2),
+          nsq(nsq_in),
+          dsq(d_in / nsq_in),
+          zn_sphere_codec(static_cast<int>(dsq), r2),
+          scale_nbit(scale_nbit_in) {
+    FAISS_THROW_IF_NOT(d_in % nsq_in == 0);
 
     lattice_nbit = 0;
     while (!(((uint64_t)1 << lattice_nbit) >= zn_sphere_codec.nv)) {
         lattice_nbit++;
     }
 
-    int total_nbit = (lattice_nbit + scale_nbit) * nsq;
+    int total_nbit = (lattice_nbit + scale_nbit_in) * nsq_in;
 
     code_size = (total_nbit + 7) / 8;
 
@@ -44,17 +45,19 @@ void IndexLattice::train(idx_t n, const float* x) {
         maxs[sq] = -1;
     }
 
-    for (idx_t i = 0; i < n; i++) {
-        for (int sq = 0; sq < nsq; sq++) {
-            float norm2 = fvec_norm_L2sqr(x + i * d + sq * dsq, dsq);
-            if (norm2 > maxs[sq]) {
-                maxs[sq] = norm2;
-            }
-            if (norm2 < mins[sq]) {
-                mins[sq] = norm2;
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (idx_t i = 0; i < n; i++) {
+            for (int sq = 0; sq < nsq; sq++) {
+                float norm2 = fvec_norm_L2sqr<SL>(x + i * d + sq * dsq, dsq);
+                if (norm2 > maxs[sq]) {
+                    maxs[sq] = norm2;
+                }
+                if (norm2 < mins[sq]) {
+                    mins[sq] = norm2;
+                }
             }
         }
-    }
+    });
 
     for (int sq = 0; sq < nsq; sq++) {
         mins[sq] = sqrtf(mins[sq]);
@@ -69,32 +72,34 @@ size_t IndexLattice::sa_code_size() const {
     return code_size;
 }
 
-void IndexLattice::sa_encode(idx_t n, const float* x, uint8_t* codes) const {
+void IndexLattice::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
     const float* mins = trained.data();
     const float* maxs = mins + nsq;
     int64_t sc = int64_t(1) << scale_nbit;
 
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for
-    for (idx_t i = 0; i < n; i++) {
-        BitstringWriter wr(codes + i * code_size, code_size);
-        const float* xi = x + i * d;
-        for (int j = 0; j < nsq; j++) {
-            float nj = (sqrtf(fvec_norm_L2sqr(xi, dsq)) - mins[j]) * sc /
-                    (maxs[j] - mins[j]);
-            if (nj < 0) {
-                nj = 0;
-            }
-            if (nj >= sc) {
-                nj = sc - 1;
+        for (idx_t i = 0; i < n; i++) {
+            BitstringWriter wr(bytes + i * code_size, code_size);
+            const float* xi = x + i * d;
+            for (int j = 0; j < nsq; j++) {
+                float nj = (sqrtf(fvec_norm_L2sqr<SL>(xi, dsq)) - mins[j]) *
+                        sc / (maxs[j] - mins[j]);
+                if (nj < 0) {
+                    nj = 0;
+                }
+                if (nj >= sc) {
+                    nj = sc - 1;
+                }
+                wr.write((int64_t)nj, scale_nbit);
+                wr.write(zn_sphere_codec.encode(xi), lattice_nbit);
+                xi += dsq;
             }
-            wr.write((int64_t)nj, scale_nbit);
-            wr.write(zn_sphere_codec.encode(xi), lattice_nbit);
-            xi += dsq;
         }
-    }
+    });
 }
 
-void IndexLattice::sa_decode(idx_t n, const uint8_t* codes, float* x) const {
+void IndexLattice::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
     const float* mins = trained.data();
     const float* maxs = mins + nsq;
     float sc = int64_t(1) << scale_nbit;
@@ -102,7 +107,7 @@ void IndexLattice::sa_decode(idx_t n, const uint8_t* codes, float* x) const {
 
 #pragma omp parallel for
     for (idx_t i = 0; i < n; i++) {
-        BitstringReader rd(codes + i * code_size, code_size);
+        BitstringReader rd(bytes + i * code_size, code_size);
         float* xi = x + i * d;
         for (int j = 0; j < nsq; j++) {
             float norm =
@@ -110,7 +115,7 @@ void IndexLattice::sa_decode(idx_t n, const uint8_t* codes, float* x) const {
                     mins[j];
             norm /= r;
             zn_sphere_codec.decode(rd.read(lattice_nbit), xi);
-            for (int l = 0; l < dsq; l++) {
+            for (size_t l = 0; l < dsq; l++) {
                 xi[l] *= norm;
             }
             xi += dsq;

data/vendor/faiss/faiss/IndexNNDescent.cpp

@@ -9,6 +9,7 @@
 
 #include <faiss/IndexNNDescent.h>
 
+#include <atomic>
 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
@@ -19,7 +20,7 @@
 #include <faiss/IndexFlat.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/utils/distances.h>
+#include <faiss/impl/VisitedTable.h>
 
 extern "C" {
 
@@ -65,17 +66,17 @@ DistanceComputer* storage_distance_computer(const Index* storage) {
  * IndexNNDescent implementation
  **************************************************************/
 
-IndexNNDescent::IndexNNDescent(int d, int K, MetricType metric)
-        : Index(d, metric),
-          nndescent(d, K),
+IndexNNDescent::IndexNNDescent(int d_in, int K, MetricType metric)
+        : Index(d_in, metric),
+          nndescent(d_in, K),
           own_fields(false),
           storage(nullptr) {}
 
-IndexNNDescent::IndexNNDescent(Index* storage, int K)
-        : Index(storage->d, storage->metric_type),
-          nndescent(storage->d, K),
+IndexNNDescent::IndexNNDescent(Index* storage_in, int K)
+        : Index(storage_in->d, storage_in->metric_type),
+          nndescent(storage_in->d, K),
           own_fields(false),
-          storage(storage) {}
+          storage(storage_in) {}
 
 IndexNNDescent::~IndexNNDescent() {
     if (own_fields) {
@@ -118,28 +119,43 @@ void IndexNNDescent::search(
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
 
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 #pragma omp parallel
         {
-            VisitedTable vt(ntotal);
-
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(storage));
+            std::unique_ptr<DistanceComputer> dis;
+            std::unique_ptr<VisitedTable> vt;
+            try {
+                vt = std::make_unique<VisitedTable>(ntotal);
+                dis.reset(storage_distance_computer(storage));
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 
 #pragma omp for
             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);
-
-                nndescent.search(*dis, k, idxi, simi, vt);
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
+                }
+                try {
+                    idx_t* idxi = labels + i * k;
+                    float* simi = distances + i * k;
+                    dis->set_query(x + i * d);
+
+                    nndescent.search(
+                            *dis, static_cast<int>(k), idxi, simi, *vt);
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
+                }
             }
         }
+        omp_rethrow_if_exception(ex);
         InterruptCallback::check();
     }
 
     if (metric_type == METRIC_INNER_PRODUCT) {
         // we need to revert the negated distances
-        for (size_t i = 0; i < k * n; i++) {
+        for (idx_t i = 0; i < k * n; i++) {
             distances[i] = -distances[i];
         }
     }
@@ -162,7 +178,7 @@ void IndexNNDescent::add(idx_t n, const float* x) {
     ntotal = storage->ntotal;
 
     std::unique_ptr<DistanceComputer> dis(storage_distance_computer(storage));
-    nndescent.build(*dis, ntotal, verbose);
+    nndescent.build(*dis, static_cast<int>(ntotal), verbose);
 }
 
 void IndexNNDescent::reset() {
@@ -183,8 +199,8 @@ IndexNNDescentFlat::IndexNNDescentFlat() {
     is_trained = true;
 }
 
-IndexNNDescentFlat::IndexNNDescentFlat(int d, int M, MetricType metric)
-        : IndexNNDescent(new IndexFlat(d, metric), M) {
+IndexNNDescentFlat::IndexNNDescentFlat(int d_in, int M, MetricType metric)
+        : IndexNNDescent(new IndexFlat(d_in, metric), M) {
     own_fields = true;
     is_trained = true;
 }

data/vendor/faiss/faiss/IndexNNDescent.h

@@ -30,8 +30,8 @@ struct IndexNNDescent : Index {
     NNDescent nndescent;
 
     // the sequential storage
-    bool own_fields;
-    Index* storage;
+    bool own_fields = false;
+    Index* storage = nullptr;
 
     explicit IndexNNDescent(
             int d = 0,

data/vendor/faiss/faiss/IndexNSG.cpp

@@ -9,6 +9,7 @@
 
 #include <faiss/IndexNSG.h>
 
+#include <atomic>
 #include <cinttypes>
 #include <memory>
 
@@ -16,7 +17,7 @@
 #include <faiss/IndexNNDescent.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
-#include <faiss/utils/distances.h>
+#include <faiss/impl/VisitedTable.h>
 
 namespace faiss {
 
@@ -26,14 +27,15 @@ using namespace nsg;
  * IndexNSG implementation
  **************************************************************/
 
-IndexNSG::IndexNSG(int d, int R, MetricType metric) : Index(d, metric), nsg(R) {
+IndexNSG::IndexNSG(int d_in, int R, MetricType metric)
+        : Index(d_in, metric), nsg(R) {
     nndescent_L = GK + 50;
 }
 
-IndexNSG::IndexNSG(Index* storage, int R)
-        : Index(storage->d, storage->metric_type),
+IndexNSG::IndexNSG(Index* storage_in, int R)
+        : Index(storage_in->d, storage_in->metric_type),
           nsg(R),
-          storage(storage),
+          storage(storage_in),
           build_type(1) {
     nndescent_L = GK + 50;
 }
@@ -72,30 +74,45 @@ void IndexNSG::search(
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
 
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 #pragma omp parallel
         {
-            VisitedTable vt(ntotal);
-
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(storage));
+            std::unique_ptr<DistanceComputer> dis;
+            std::unique_ptr<VisitedTable> vt;
+            try {
+                vt = std::make_unique<VisitedTable>(
+                        ntotal, nsg.use_visited_hashset);
+                dis.reset(storage_distance_computer(storage));
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 
 #pragma omp for
             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);
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
+                }
+                try {
+                    idx_t* idxi = labels + i * k;
+                    float* simi = distances + i * k;
+                    dis->set_query(x + i * d);
 
-                nsg.search(*dis, k, idxi, simi, vt);
+                    nsg.search(*dis, static_cast<int>(k), idxi, simi, *vt);
 
-                vt.advance();
+                    vt->advance();
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
+                }
             }
         }
+        omp_rethrow_if_exception(ex);
         InterruptCallback::check();
     }
 
     if (is_similarity_metric(metric_type)) {
         // we need to revert the negated distances
-        for (size_t i = 0; i < k * n; i++) {
+        for (idx_t i = 0; i < k * n; i++) {
             distances[i] = -distances[i];
         }
     }
@@ -114,7 +131,7 @@ void IndexNSG::build(idx_t n, const float* x, idx_t* knn_graph, int gk) {
     // check the knn graph
     check_knn_graph(knn_graph, n, gk);
 
-    const nsg::Graph<idx_t> knng(knn_graph, n, gk);
+    const nsg::Graph<idx_t> knng(knn_graph, static_cast<int>(n), gk);
     nsg.build(storage, n, knng, verbose);
     is_built = true;
 }
@@ -222,7 +239,7 @@ void IndexNSG::add(idx_t n, const float* x) {
         printf("  nsg building\n");
     }
 
-    const nsg::Graph<idx_t> knn_graph(knng.data(), n, GK);
+    const nsg::Graph<idx_t> knn_graph(knng.data(), static_cast<int>(n), GK);
     nsg.build(storage, n, knn_graph, verbose);
     is_built = true;
 }
@@ -273,8 +290,8 @@ IndexNSGFlat::IndexNSGFlat() {
     is_trained = true;
 }
 
-IndexNSGFlat::IndexNSGFlat(int d, int R, MetricType metric)
-        : IndexNSG(new IndexFlat(d, metric), R) {
+IndexNSGFlat::IndexNSGFlat(int d_in, int R, MetricType metric)
+        : IndexNSG(new IndexFlat(d_in, metric), R) {
     own_fields = true;
     is_trained = true;
 }
@@ -285,8 +302,8 @@ IndexNSGFlat::IndexNSGFlat(int d, int R, MetricType metric)
 
 IndexNSGPQ::IndexNSGPQ() = default;
 
-IndexNSGPQ::IndexNSGPQ(int d, int pq_m, int M, int pq_nbits)
-        : IndexNSG(new IndexPQ(d, pq_m, pq_nbits), M) {
+IndexNSGPQ::IndexNSGPQ(int d_in, int pq_m, int M, int pq_nbits)
+        : IndexNSG(new IndexPQ(d_in, pq_m, pq_nbits), M) {
     own_fields = true;
     is_trained = false;
 }
@@ -301,11 +318,11 @@ void IndexNSGPQ::train(idx_t n, const float* x) {
  **************************************************************/
 
 IndexNSGSQ::IndexNSGSQ(
-        int d,
+        int d_in,
         ScalarQuantizer::QuantizerType qtype,
         int M,
         MetricType metric)
-        : IndexNSG(new IndexScalarQuantizer(d, qtype, metric), M) {
+        : IndexNSG(new IndexScalarQuantizer(d_in, qtype, metric), M) {
     is_trained = this->storage->is_trained;
     own_fields = true;
 }

data/vendor/faiss/faiss/IndexNSG.h

@@ -9,8 +9,6 @@
 
 #pragma once
 
-#include <vector>
-
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexNNDescent.h>
 #include <faiss/IndexPQ.h>

data/vendor/faiss/faiss/IndexNeuralNetCodec.cpp

@@ -16,29 +16,43 @@ namespace faiss {
  *********************************************************/
 
 IndexNeuralNetCodec::IndexNeuralNetCodec(
-        int d,
-        int M,
-        int nbits,
+        int d_in,
+        int M_in,
+        int nbits_in,
         MetricType metric)
-        : IndexFlatCodes((M * nbits + 7) / 8, d, metric), M(M), nbits(nbits) {
+        : IndexFlatCodes((M_in * nbits_in + 7) / 8, d_in, metric),
+          M(M_in),
+          nbits(nbits_in) {
     is_trained = false;
 }
 
-void IndexNeuralNetCodec::train(idx_t n, const float* x) {
+void IndexNeuralNetCodec::train(idx_t /*n*/, const float* /*x*/) {
     FAISS_THROW_MSG("Training not implemented in C++, use Pytorch");
 }
 
-void IndexNeuralNetCodec::sa_encode(idx_t n, const float* x, uint8_t* codes)
+void IndexNeuralNetCodec::sa_encode(idx_t n, const float* x, uint8_t* bytes)
         const {
     nn::Tensor2D x_tensor(n, d, x);
     nn::Int32Tensor2D codes_tensor = net->encode(x_tensor);
-    pack_bitstrings(n, M, nbits, codes_tensor.data(), codes, code_size);
+    pack_bitstrings(
+            n,
+            M,
+            static_cast<int>(nbits),
+            codes_tensor.data(),
+            bytes,
+            code_size);
 }
 
-void IndexNeuralNetCodec::sa_decode(idx_t n, const uint8_t* codes, float* x)
+void IndexNeuralNetCodec::sa_decode(idx_t n, const uint8_t* bytes, float* x)
         const {
     nn::Int32Tensor2D codes_tensor(n, M);
-    unpack_bitstrings(n, M, nbits, codes, code_size, codes_tensor.data());
+    unpack_bitstrings(
+            n,
+            M,
+            static_cast<int>(nbits),
+            bytes,
+            code_size,
+            codes_tensor.data());
     nn::Tensor2D x_tensor = net->decode(codes_tensor);
     memcpy(x, x_tensor.data(), d * n * sizeof(float));
 }
@@ -47,9 +61,15 @@ void IndexNeuralNetCodec::sa_decode(idx_t n, const uint8_t* codes, float* x)
  * IndexQINeuralNetCodec implementation
  *********************************************************/
 
-IndexQINCo::IndexQINCo(int d, int M, int nbits, int L, int h, MetricType metric)
-        : IndexNeuralNetCodec(d, M, nbits, metric),
-          qinco(d, 1 << nbits, L, M, h) {
+IndexQINCo::IndexQINCo(
+        int d_in,
+        int M_in,
+        int nbits_in,
+        int L,
+        int h,
+        MetricType metric)
+        : IndexNeuralNetCodec(d_in, M_in, nbits_in, metric),
+          qinco(d_in, 1 << nbits_in, L, M_in, h) {
     net = &qinco;
 }