faiss 0.2.7 → 0.3.1

data/vendor/faiss/faiss/IndexFastScan.cpp

@@ -7,8 +7,8 @@
 
 #include <faiss/IndexFastScan.h>
 
-#include <limits.h>
 #include <cassert>
+#include <climits>
 #include <memory>
 
 #include <omp.h>
@@ -37,22 +37,22 @@ inline size_t roundup(size_t a, size_t b) {
 
 void IndexFastScan::init_fastscan(
         int d,
-        size_t M,
-        size_t nbits,
+        size_t M_2,
+        size_t nbits_2,
         MetricType metric,
         int bbs) {
-    FAISS_THROW_IF_NOT(nbits == 4);
+    FAISS_THROW_IF_NOT(nbits_2 == 4);
     FAISS_THROW_IF_NOT(bbs % 32 == 0);
     this->d = d;
-    this->M = M;
-    this->nbits = nbits;
+    this->M = M_2;
+    this->nbits = nbits_2;
     this->metric_type = metric;
     this->bbs = bbs;
-    ksub = (1 << nbits);
+    ksub = (1 << nbits_2);
 
-    code_size = (M * nbits + 7) / 8;
+    code_size = (M_2 * nbits_2 + 7) / 8;
     ntotal = ntotal2 = 0;
-    M2 = roundup(M, 2);
+    M2 = roundup(M_2, 2);
     is_trained = false;
 }
 
@@ -158,7 +158,7 @@ void IndexFastScan::merge_from(Index& otherIndex, idx_t add_id) {
 
 namespace {
 
-template <class C, typename dis_t, class Scaler>
+template <class C, typename dis_t>
 void estimators_from_tables_generic(
         const IndexFastScan& index,
         const uint8_t* codes,
@@ -167,25 +167,28 @@ void estimators_from_tables_generic(
         size_t k,
         typename C::T* heap_dis,
         int64_t* heap_ids,
-        const Scaler& scaler) {
+        const NormTableScaler* scaler) {
     using accu_t = typename C::T;
 
     for (size_t j = 0; j < ncodes; ++j) {
         BitstringReader bsr(codes + j * index.code_size, index.code_size);
         accu_t dis = 0;
         const dis_t* dt = dis_table;
-        for (size_t m = 0; m < index.M - scaler.nscale; m++) {
+        int nscale = scaler ? scaler->nscale : 0;
+
+        for (size_t m = 0; m < index.M - nscale; m++) {
             uint64_t c = bsr.read(index.nbits);
             dis += dt[c];
             dt += index.ksub;
         }
 
-        for (size_t m = 0; m < scaler.nscale; m++) {
-            uint64_t c = bsr.read(index.nbits);
-            dis += scaler.scale_one(dt[c]);
-            dt += index.ksub;
+        if (nscale) {
+            for (size_t m = 0; m < nscale; m++) {
+                uint64_t c = bsr.read(index.nbits);
+                dis += scaler->scale_one(dt[c]);
+                dt += index.ksub;
+            }
         }
-
         if (C::cmp(heap_dis[0], dis)) {
             heap_pop<C>(k, heap_dis, heap_ids);
             heap_push<C>(k, heap_dis, heap_ids, dis, j);
@@ -193,6 +196,27 @@ void estimators_from_tables_generic(
     }
 }
 
+template <class C>
+ResultHandlerCompare<C, false>* make_knn_handler(
+        int impl,
+        idx_t n,
+        idx_t k,
+        size_t ntotal,
+        float* distances,
+        idx_t* labels) {
+    using HeapHC = HeapHandler<C, false>;
+    using ReservoirHC = ReservoirHandler<C, false>;
+    using SingleResultHC = SingleResultHandler<C, false>;
+
+    if (k == 1) {
+        return new SingleResultHC(n, ntotal, distances, labels);
+    } else if (impl % 2 == 0) {
+        return new HeapHC(n, ntotal, k, distances, labels);
+    } else /* if (impl % 2 == 1) */ {
+        return new ReservoirHC(n, ntotal, k, 2 * k, distances, labels);
+    }
+}
+
 } // anonymous namespace
 
 using namespace quantize_lut;
@@ -241,22 +265,21 @@ void IndexFastScan::search(
             !params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
-    DummyScaler scaler;
     if (metric_type == METRIC_L2) {
-        search_dispatch_implem<true>(n, x, k, distances, labels, scaler);
+        search_dispatch_implem<true>(n, x, k, distances, labels, nullptr);
     } else {
-        search_dispatch_implem<false>(n, x, k, distances, labels, scaler);
+        search_dispatch_implem<false>(n, x, k, distances, labels, nullptr);
     }
 }
 
-template <bool is_max, class Scaler>
+template <bool is_max>
 void IndexFastScan::search_dispatch_implem(
         idx_t n,
         const float* x,
         idx_t k,
         float* distances,
         idx_t* labels,
-        const Scaler& scaler) const {
+        const NormTableScaler* scaler) const {
     using Cfloat = typename std::conditional<
             is_max,
             CMax<float, int64_t>,
@@ -319,14 +342,14 @@ void IndexFastScan::search_dispatch_implem(
     }
 }
 
-template <class Cfloat, class Scaler>
+template <class Cfloat>
 void IndexFastScan::search_implem_234(
         idx_t n,
         const float* x,
         idx_t k,
         float* distances,
         idx_t* labels,
-        const Scaler& scaler) const {
+        const NormTableScaler* scaler) const {
     FAISS_THROW_IF_NOT(implem == 2 || implem == 3 || implem == 4);
 
     const size_t dim12 = ksub * M;
@@ -378,7 +401,7 @@ void IndexFastScan::search_implem_234(
     }
 }
 
-template <class C, class Scaler>
+template <class C>
 void IndexFastScan::search_implem_12(
         idx_t n,
         const float* x,
@@ -386,7 +409,8 @@ void IndexFastScan::search_implem_12(
         float* distances,
         idx_t* labels,
         int impl,
-        const Scaler& scaler) const {
+        const NormTableScaler* scaler) const {
+    using RH = ResultHandlerCompare<C, false>;
     FAISS_THROW_IF_NOT(bbs == 32);
 
     // handle qbs2 blocking by recursive call
@@ -432,63 +456,31 @@ void IndexFastScan::search_implem_12(
             pq4_pack_LUT_qbs(qbs, M2, quantized_dis_tables.get(), LUT.get());
     FAISS_THROW_IF_NOT(LUT_nq == n);
 
-    if (k == 1) {
-        SingleResultHandler<C> handler(n, ntotal);
-        if (skip & 4) {
-            // pass
-        } else {
-            handler.disable = bool(skip & 2);
-            pq4_accumulate_loop_qbs(
-                    qbs, ntotal2, M2, codes.get(), LUT.get(), handler, scaler);
-        }
+    std::unique_ptr<RH> handler(
+            make_knn_handler<C>(impl, n, k, ntotal, distances, labels));
+    handler->disable = bool(skip & 2);
+    handler->normalizers = normalizers.get();
 
-        handler.to_flat_arrays(distances, labels, normalizers.get());
-
-    } else if (impl == 12) {
-        std::vector<uint16_t> tmp_dis(n * k);
-        std::vector<int32_t> tmp_ids(n * k);
-
-        if (skip & 4) {
-            // skip
-        } else {
-            HeapHandler<C> handler(
-                    n, tmp_dis.data(), tmp_ids.data(), k, ntotal);
-            handler.disable = bool(skip & 2);
-
-            pq4_accumulate_loop_qbs(
-                    qbs, ntotal2, M2, codes.get(), LUT.get(), handler, scaler);
-
-            if (!(skip & 8)) {
-                handler.to_flat_arrays(distances, labels, normalizers.get());
-            }
-        }
-
-    } else { // impl == 13
-
-        ReservoirHandler<C> handler(n, ntotal, k, 2 * k);
-        handler.disable = bool(skip & 2);
-
-        if (skip & 4) {
-            // skip
-        } else {
-            pq4_accumulate_loop_qbs(
-                    qbs, ntotal2, M2, codes.get(), LUT.get(), handler, scaler);
-        }
-
-        if (!(skip & 8)) {
-            handler.to_flat_arrays(distances, labels, normalizers.get());
-        }
-
-        FastScan_stats.t0 += handler.times[0];
-        FastScan_stats.t1 += handler.times[1];
-        FastScan_stats.t2 += handler.times[2];
-        FastScan_stats.t3 += handler.times[3];
+    if (skip & 4) {
+        // pass
+    } else {
+        pq4_accumulate_loop_qbs(
+                qbs,
+                ntotal2,
+                M2,
+                codes.get(),
+                LUT.get(),
+                *handler.get(),
+                scaler);
+    }
+    if (!(skip & 8)) {
+        handler->end();
     }
 }
 
 FastScanStats FastScan_stats;
 
-template <class C, class Scaler>
+template <class C>
 void IndexFastScan::search_implem_14(
         idx_t n,
         const float* x,
@@ -496,7 +488,8 @@ void IndexFastScan::search_implem_14(
         float* distances,
         idx_t* labels,
         int impl,
-        const Scaler& scaler) const {
+        const NormTableScaler* scaler) const {
+    using RH = ResultHandlerCompare<C, false>;
     FAISS_THROW_IF_NOT(bbs % 32 == 0);
 
     int qbs2 = qbs == 0 ? 4 : qbs;
@@ -531,91 +524,29 @@ void IndexFastScan::search_implem_14(
     AlignedTable<uint8_t> LUT(n * dim12);
     pq4_pack_LUT(n, M2, quantized_dis_tables.get(), LUT.get());
 
-    if (k == 1) {
-        SingleResultHandler<C> handler(n, ntotal);
-        if (skip & 4) {
-            // pass
-        } else {
-            handler.disable = bool(skip & 2);
-            pq4_accumulate_loop(
-                    n,
-                    ntotal2,
-                    bbs,
-                    M2,
-                    codes.get(),
-                    LUT.get(),
-                    handler,
-                    scaler);
-        }
-        handler.to_flat_arrays(distances, labels, normalizers.get());
-
-    } else if (impl == 14) {
-        std::vector<uint16_t> tmp_dis(n * k);
-        std::vector<int32_t> tmp_ids(n * k);
-
-        if (skip & 4) {
-            // skip
-        } else if (k > 1) {
-            HeapHandler<C> handler(
-                    n, tmp_dis.data(), tmp_ids.data(), k, ntotal);
-            handler.disable = bool(skip & 2);
-
-            pq4_accumulate_loop(
-                    n,
-                    ntotal2,
-                    bbs,
-                    M2,
-                    codes.get(),
-                    LUT.get(),
-                    handler,
-                    scaler);
-
-            if (!(skip & 8)) {
-                handler.to_flat_arrays(distances, labels, normalizers.get());
-            }
-        }
-
-    } else { // impl == 15
-
-        ReservoirHandler<C> handler(n, ntotal, k, 2 * k);
-        handler.disable = bool(skip & 2);
-
-        if (skip & 4) {
-            // skip
-        } else {
-            pq4_accumulate_loop(
-                    n,
-                    ntotal2,
-                    bbs,
-                    M2,
-                    codes.get(),
-                    LUT.get(),
-                    handler,
-                    scaler);
-        }
+    std::unique_ptr<RH> handler(
+            make_knn_handler<C>(impl, n, k, ntotal, distances, labels));
+    handler->disable = bool(skip & 2);
+    handler->normalizers = normalizers.get();
 
-        if (!(skip & 8)) {
-            handler.to_flat_arrays(distances, labels, normalizers.get());
-        }
+    if (skip & 4) {
+        // pass
+    } else {
+        pq4_accumulate_loop(
+                n,
+                ntotal2,
+                bbs,
+                M2,
+                codes.get(),
+                LUT.get(),
+                *handler.get(),
+                scaler);
+    }
+    if (!(skip & 8)) {
+        handler->end();
     }
 }
 
-template void IndexFastScan::search_dispatch_implem<true, NormTableScaler>(
-        idx_t n,
-        const float* x,
-        idx_t k,
-        float* distances,
-        idx_t* labels,
-        const NormTableScaler& scaler) const;
-
-template void IndexFastScan::search_dispatch_implem<false, NormTableScaler>(
-        idx_t n,
-        const float* x,
-        idx_t k,
-        float* distances,
-        idx_t* labels,
-        const NormTableScaler& scaler) const;
-
 void IndexFastScan::reconstruct(idx_t key, float* recons) const {
     std::vector<uint8_t> code(code_size, 0);
     BitstringWriter bsw(code.data(), code_size);

data/vendor/faiss/faiss/IndexFastScan.h

@@ -13,6 +13,7 @@
 namespace faiss {
 
 struct CodePacker;
+struct NormTableScaler;
 
 /** Fast scan version of IndexPQ and IndexAQ. Works for 4-bit PQ and AQ for now.
  *
@@ -87,25 +88,25 @@ struct IndexFastScan : Index {
             uint8_t* lut,
             float* normalizers) const;
 
-    template <bool is_max, class Scaler>
+    template <bool is_max>
     void search_dispatch_implem(
             idx_t n,
             const float* x,
             idx_t k,
             float* distances,
             idx_t* labels,
-            const Scaler& scaler) const;
+            const NormTableScaler* scaler) const;
 
-    template <class Cfloat, class Scaler>
+    template <class Cfloat>
     void search_implem_234(
             idx_t n,
             const float* x,
             idx_t k,
             float* distances,
             idx_t* labels,
-            const Scaler& scaler) const;
+            const NormTableScaler* scaler) const;
 
-    template <class C, class Scaler>
+    template <class C>
     void search_implem_12(
             idx_t n,
             const float* x,
@@ -113,9 +114,9 @@ struct IndexFastScan : Index {
             float* distances,
             idx_t* labels,
             int impl,
-            const Scaler& scaler) const;
+            const NormTableScaler* scaler) const;
 
-    template <class C, class Scaler>
+    template <class C>
     void search_implem_14(
             idx_t n,
             const float* x,
@@ -123,7 +124,7 @@ struct IndexFastScan : Index {
             float* distances,
             idx_t* labels,
             int impl,
-            const Scaler& scaler) const;
+            const NormTableScaler* scaler) const;
 
     void reconstruct(idx_t key, float* recons) const override;
     size_t remove_ids(const IDSelector& sel) override;

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -14,6 +14,7 @@
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/extra_distances.h>
+#include <faiss/utils/prefetch.h>
 #include <faiss/utils/sorting.h>
 #include <faiss/utils/utils.h>
 #include <cstring>
@@ -122,6 +123,39 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
     void set_query(const float* x) override {
         q = x;
     }
+
+    // compute four distances
+    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) final override {
+        ndis += 4;
+
+        // compute first, assign next
+        const float* __restrict y0 =
+                reinterpret_cast<const float*>(codes + idx0 * code_size);
+        const float* __restrict y1 =
+                reinterpret_cast<const float*>(codes + idx1 * code_size);
+        const float* __restrict y2 =
+                reinterpret_cast<const float*>(codes + idx2 * code_size);
+        const float* __restrict y3 =
+                reinterpret_cast<const float*>(codes + idx3 * code_size);
+
+        float dp0 = 0;
+        float dp1 = 0;
+        float dp2 = 0;
+        float dp3 = 0;
+        fvec_L2sqr_batch_4(q, y0, y1, y2, y3, d, dp0, dp1, dp2, dp3);
+        dis0 = dp0;
+        dis1 = dp1;
+        dis2 = dp2;
+        dis3 = dp3;
+    }
 };
 
 struct FlatIPDis : FlatCodesDistanceComputer {
@@ -131,13 +165,13 @@ struct FlatIPDis : FlatCodesDistanceComputer {
     const float* b;
     size_t ndis;
 
-    float symmetric_dis(idx_t i, idx_t j) override {
+    float symmetric_dis(idx_t i, idx_t j) final override {
         return fvec_inner_product(b + j * d, b + i * d, d);
     }
 
-    float distance_to_code(const uint8_t* code) final {
+    float distance_to_code(const uint8_t* code) final override {
         ndis++;
-        return fvec_inner_product(q, (float*)code, d);
+        return fvec_inner_product(q, (const float*)code, d);
     }
 
     explicit FlatIPDis(const IndexFlat& storage, const float* q = nullptr)
@@ -153,6 +187,39 @@ struct FlatIPDis : FlatCodesDistanceComputer {
     void set_query(const float* x) override {
         q = x;
     }
+
+    // compute four distances
+    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) final override {
+        ndis += 4;
+
+        // compute first, assign next
+        const float* __restrict y0 =
+                reinterpret_cast<const float*>(codes + idx0 * code_size);
+        const float* __restrict y1 =
+                reinterpret_cast<const float*>(codes + idx1 * code_size);
+        const float* __restrict y2 =
+                reinterpret_cast<const float*>(codes + idx2 * code_size);
+        const float* __restrict y3 =
+                reinterpret_cast<const float*>(codes + idx3 * code_size);
+
+        float dp0 = 0;
+        float dp1 = 0;
+        float dp2 = 0;
+        float dp3 = 0;
+        fvec_inner_product_batch_4(q, y0, y1, y2, y3, d, dp0, dp1, dp2, dp3);
+        dis0 = dp0;
+        dis1 = dp1;
+        dis2 = dp2;
+        dis3 = dp3;
+    }
 };
 
 } // namespace
@@ -184,6 +251,131 @@ void IndexFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
     }
 }
 
+/***************************************************
+ * IndexFlatL2
+ ***************************************************/
+
+namespace {
+struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
+    size_t d;
+    idx_t nb;
+    const float* q;
+    const float* b;
+    size_t ndis;
+
+    const float* l2norms;
+    float query_l2norm;
+
+    float distance_to_code(const uint8_t* code) final override {
+        ndis++;
+        return fvec_L2sqr(q, (float*)code, d);
+    }
+
+    float operator()(const idx_t i) final override {
+        const float* __restrict y =
+                reinterpret_cast<const float*>(codes + i * code_size);
+
+        prefetch_L2(l2norms + i);
+        const float dp0 = fvec_inner_product(q, y, d);
+        return query_l2norm + l2norms[i] - 2 * dp0;
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) final override {
+        const float* __restrict yi =
+                reinterpret_cast<const float*>(codes + i * code_size);
+        const float* __restrict yj =
+                reinterpret_cast<const float*>(codes + j * code_size);
+
+        prefetch_L2(l2norms + i);
+        prefetch_L2(l2norms + j);
+        const float dp0 = fvec_inner_product(yi, yj, d);
+        return l2norms[i] + l2norms[j] - 2 * dp0;
+    }
+
+    explicit FlatL2WithNormsDis(
+            const IndexFlatL2& storage,
+            const float* q = nullptr)
+            : FlatCodesDistanceComputer(
+                      storage.codes.data(),
+                      storage.code_size),
+              d(storage.d),
+              nb(storage.ntotal),
+              q(q),
+              b(storage.get_xb()),
+              ndis(0),
+              l2norms(storage.cached_l2norms.data()),
+              query_l2norm(0) {}
+
+    void set_query(const float* x) override {
+        q = x;
+        query_l2norm = fvec_norm_L2sqr(q, d);
+    }
+
+    // compute four distances
+    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) final override {
+        ndis += 4;
+
+        // compute first, assign next
+        const float* __restrict y0 =
+                reinterpret_cast<const float*>(codes + idx0 * code_size);
+        const float* __restrict y1 =
+                reinterpret_cast<const float*>(codes + idx1 * code_size);
+        const float* __restrict y2 =
+                reinterpret_cast<const float*>(codes + idx2 * code_size);
+        const float* __restrict y3 =
+                reinterpret_cast<const float*>(codes + idx3 * code_size);
+
+        prefetch_L2(l2norms + idx0);
+        prefetch_L2(l2norms + idx1);
+        prefetch_L2(l2norms + idx2);
+        prefetch_L2(l2norms + idx3);
+
+        float dp0 = 0;
+        float dp1 = 0;
+        float dp2 = 0;
+        float dp3 = 0;
+        fvec_inner_product_batch_4(q, y0, y1, y2, y3, d, dp0, dp1, dp2, dp3);
+        dis0 = query_l2norm + l2norms[idx0] - 2 * dp0;
+        dis1 = query_l2norm + l2norms[idx1] - 2 * dp1;
+        dis2 = query_l2norm + l2norms[idx2] - 2 * dp2;
+        dis3 = query_l2norm + l2norms[idx3] - 2 * dp3;
+    }
+};
+
+} // namespace
+
+void IndexFlatL2::sync_l2norms() {
+    cached_l2norms.resize(ntotal);
+    fvec_norms_L2sqr(
+            cached_l2norms.data(),
+            reinterpret_cast<const float*>(codes.data()),
+            d,
+            ntotal);
+}
+
+void IndexFlatL2::clear_l2norms() {
+    cached_l2norms.clear();
+    cached_l2norms.shrink_to_fit();
+}
+
+FlatCodesDistanceComputer* IndexFlatL2::get_FlatCodesDistanceComputer() const {
+    if (metric_type == METRIC_L2) {
+        if (!cached_l2norms.empty()) {
+            return new FlatL2WithNormsDis(*this);
+        }
+    }
+
+    return IndexFlat::get_FlatCodesDistanceComputer();
+}
+
 /***************************************************
  * IndexFlat1D
  ***************************************************/