faiss 0.6.2 → 0.6.4

data/vendor/faiss/faiss/IndexEDEN.cpp

@@ -0,0 +1,273 @@
+/*
+ * 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.
+ */
+
+#include <faiss/IndexEDEN.h>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/ResultHandler.h>
+#include <memory>
+
+namespace faiss {
+
+IndexEDEN::IndexEDEN() = default;
+
+IndexEDEN::IndexEDEN(
+        idx_t d_in,
+        MetricType metric,
+        uint8_t nb_bits_in,
+        EDENScaleType scale_type_in)
+        : IndexFlatCodes(0, d_in, metric),
+          sq(d_in, eden_utils::quantizer_type_for_bits(nb_bits_in)),
+          scale_type(scale_type_in) {
+    FAISS_THROW_IF_NOT_MSG(
+            scale_type == EDENScaleType_UNBIASED ||
+                    scale_type == EDENScaleType_BIASED,
+            "invalid EDEN scale type");
+    sq.train(0, nullptr);
+    code_size = eden_utils::code_size(d, sq.bits);
+    is_trained = false;
+}
+
+void IndexEDEN::train(idx_t n, const float* x) {
+    std::vector<float> centroid(d, 0.0f);
+    for (idx_t i = 0; i < n; i++) {
+        for (size_t j = 0; j < static_cast<size_t>(d); j++) {
+            centroid[j] += x[i * d + j];
+        }
+    }
+
+    if (n != 0) {
+        for (size_t j = 0; j < static_cast<size_t>(d); j++) {
+            centroid[j] /= static_cast<float>(n);
+        }
+    }
+
+    center = std::move(centroid);
+    sq.train(n, x);
+    is_trained = true;
+}
+
+void IndexEDEN::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
+    FAISS_THROW_IF_NOT(is_trained);
+    eden_utils::compute_codes(
+            sq, metric_type, scale_type, x, bytes, n, center.data());
+}
+
+void IndexEDEN::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
+    FAISS_THROW_IF_NOT(is_trained);
+    eden_utils::decode(sq, bytes, x, n, center.data());
+}
+
+FlatCodesDistanceComputer* IndexEDEN::get_FlatCodesDistanceComputer() const {
+    FlatCodesDistanceComputer* dc =
+            eden_utils::get_distance_computer(sq, metric_type, center.data());
+    dc->code_size = code_size;
+    dc->codes = codes.data();
+    return dc;
+}
+
+namespace {
+
+bool use_eden_batch_scan(size_t d, size_t nb_bits) {
+    const size_t values_per_byte = 8 / nb_bits;
+    const size_t num_bytes = (d + values_per_byte - 1) / values_per_byte;
+    if (nb_bits == 1) {
+        return num_bytes >= 32;
+    }
+    if (nb_bits == 2) {
+        return num_bytes >= 16;
+    }
+    if (nb_bits == 4) {
+        return num_bytes >= 128;
+    }
+    return false;
+}
+
+struct Run_search_with_eden_dc {
+    using T = void;
+
+    template <class BlockResultHandler>
+    void f(BlockResultHandler& res, const IndexEDEN* index, const float* xq) {
+        const size_t ntotal = index->ntotal;
+        using SingleResultHandler =
+                typename BlockResultHandler::SingleResultHandler;
+        const int d = index->d;
+        const bool use_batch_scan = use_eden_batch_scan(d, index->sq.bits);
+
+#pragma omp parallel
+        {
+            std::unique_ptr<EDENFlatCodesDistanceComputer> dc(
+                    eden_utils::get_distance_computer(
+                            index->sq,
+                            index->metric_type,
+                            index->center.data()));
+            dc->code_size = index->code_size;
+            dc->codes = index->codes.data();
+            SingleResultHandler resi(res);
+
+#pragma omp for
+            for (int64_t q = 0; q < static_cast<int64_t>(res.nq); q++) {
+                resi.begin(q);
+                dc->set_query(xq + d * q);
+
+                if (use_batch_scan) {
+                    size_t i = 0;
+                    for (; i + 16 <= ntotal; i += 16) {
+                        bool keep[16];
+                        bool keep_all = true;
+                        for (size_t j = 0; j < 16; j++) {
+                            keep[j] = res.is_in_selection(i + j);
+                            keep_all &= keep[j];
+                        }
+
+                        if (keep_all) {
+                            float dis[16];
+                            dc->consecutive_distances_batch_16(i, dis);
+                            for (size_t j = 0; j < 16; j++) {
+                                resi.add_result(dis[j], i + j);
+                            }
+                        } else {
+                            for (size_t j = 0; j < 16; j++) {
+                                if (keep[j]) {
+                                    resi.add_result((*dc)(i + j), i + j);
+                                }
+                            }
+                        }
+                    }
+                    for (; i + 8 <= ntotal; i += 8) {
+                        const bool keep0 = res.is_in_selection(i);
+                        const bool keep1 = res.is_in_selection(i + 1);
+                        const bool keep2 = res.is_in_selection(i + 2);
+                        const bool keep3 = res.is_in_selection(i + 3);
+                        const bool keep4 = res.is_in_selection(i + 4);
+                        const bool keep5 = res.is_in_selection(i + 5);
+                        const bool keep6 = res.is_in_selection(i + 6);
+                        const bool keep7 = res.is_in_selection(i + 7);
+                        if (keep0 && keep1 && keep2 && keep3 && keep4 &&
+                            keep5 && keep6 && keep7) {
+                            float dis[8];
+                            dc->consecutive_distances_batch_8(i, dis);
+                            resi.add_result(dis[0], i);
+                            resi.add_result(dis[1], i + 1);
+                            resi.add_result(dis[2], i + 2);
+                            resi.add_result(dis[3], i + 3);
+                            resi.add_result(dis[4], i + 4);
+                            resi.add_result(dis[5], i + 5);
+                            resi.add_result(dis[6], i + 6);
+                            resi.add_result(dis[7], i + 7);
+                        } else {
+                            if (keep0) {
+                                resi.add_result((*dc)(i), i);
+                            }
+                            if (keep1) {
+                                resi.add_result((*dc)(i + 1), i + 1);
+                            }
+                            if (keep2) {
+                                resi.add_result((*dc)(i + 2), i + 2);
+                            }
+                            if (keep3) {
+                                resi.add_result((*dc)(i + 3), i + 3);
+                            }
+                            if (keep4) {
+                                resi.add_result((*dc)(i + 4), i + 4);
+                            }
+                            if (keep5) {
+                                resi.add_result((*dc)(i + 5), i + 5);
+                            }
+                            if (keep6) {
+                                resi.add_result((*dc)(i + 6), i + 6);
+                            }
+                            if (keep7) {
+                                resi.add_result((*dc)(i + 7), i + 7);
+                            }
+                        }
+                    }
+                    for (; i + 4 <= ntotal; i += 4) {
+                        const bool keep0 = res.is_in_selection(i);
+                        const bool keep1 = res.is_in_selection(i + 1);
+                        const bool keep2 = res.is_in_selection(i + 2);
+                        const bool keep3 = res.is_in_selection(i + 3);
+                        if (keep0 && keep1 && keep2 && keep3) {
+                            float dis0;
+                            float dis1;
+                            float dis2;
+                            float dis3;
+                            dc->distances_batch_4(
+                                    i,
+                                    i + 1,
+                                    i + 2,
+                                    i + 3,
+                                    dis0,
+                                    dis1,
+                                    dis2,
+                                    dis3);
+                            resi.add_result(dis0, i);
+                            resi.add_result(dis1, i + 1);
+                            resi.add_result(dis2, i + 2);
+                            resi.add_result(dis3, i + 3);
+                        } else {
+                            if (keep0) {
+                                resi.add_result((*dc)(i), i);
+                            }
+                            if (keep1) {
+                                resi.add_result((*dc)(i + 1), i + 1);
+                            }
+                            if (keep2) {
+                                resi.add_result((*dc)(i + 2), i + 2);
+                            }
+                            if (keep3) {
+                                resi.add_result((*dc)(i + 3), i + 3);
+                            }
+                        }
+                    }
+                    for (; i < ntotal; i++) {
+                        if (res.is_in_selection(i)) {
+                            resi.add_result((*dc)(i), i);
+                        }
+                    }
+                } else {
+                    for (size_t i = 0; i < ntotal; i++) {
+                        if (res.is_in_selection(i)) {
+                            resi.add_result((*dc)(i), i);
+                        }
+                    }
+                }
+
+                resi.end();
+            }
+        }
+    }
+};
+
+} // namespace
+
+void IndexEDEN::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    FAISS_THROW_IF_NOT(is_trained);
+    const IDSelector* sel = params != nullptr ? params->sel : nullptr;
+    Run_search_with_eden_dc r;
+    dispatch_knn_ResultHandler(
+            n, distances, labels, k, metric_type, sel, r, this, x);
+}
+
+void IndexEDEN::range_search(
+        idx_t /*n*/,
+        const float* x,
+        float radius,
+        RangeSearchResult* result,
+        const SearchParameters* params) const {
+    const IDSelector* sel = params != nullptr ? params->sel : nullptr;
+    Run_search_with_eden_dc r;
+    dispatch_range_ResultHandler(result, radius, metric_type, sel, r, this, x);
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexEDEN.h

@@ -0,0 +1,57 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <faiss/IndexFlatCodes.h>
+#include <faiss/impl/EDENQuantizer.h>
+
+namespace faiss {
+
+struct IndexEDEN : IndexFlatCodes {
+    ScalarQuantizer sq;
+    EDENScaleType scale_type = EDENScaleType_UNBIASED;
+
+    // Center of all training points. Random rotation, if desired, is expected
+    // to be supplied externally via IndexPreTransform.
+    std::vector<float> center;
+
+    IndexEDEN();
+
+    // Factory strings: EDEN, EDEN<n>, EDENBIASED, EDEN<n>BIASED.
+    // <n> is the Lloyd-Max bit budget, 1..8. The default is EDEN's original
+    // unbiased scale; the BIASED suffix selects the MSE-minimizing scale.
+    explicit IndexEDEN(
+            idx_t d,
+            MetricType metric = METRIC_L2,
+            uint8_t nb_bits = 1,
+            EDENScaleType scale_type = EDENScaleType_UNBIASED);
+
+    void train(idx_t n, const float* x) override;
+
+    void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
+    void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+
+    FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result,
+            const SearchParameters* params = nullptr) const override;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexFastScan.cpp

@@ -8,6 +8,7 @@
 #include <faiss/IndexFastScan.h>
 
 #include <omp.h>
+#include <algorithm>
 #include <cstring>
 #include <memory>
 
@@ -266,8 +267,7 @@ void IndexFastScan::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
     FastScanDistancePostProcessing empty_context{};
@@ -316,7 +316,7 @@ void IndexFastScan::search_dispatch_implem(
     if (implem == 1) {
         FAISS_THROW_MSG("not implemented");
     } else if (implem == 2 || implem == 3 || implem == 4) {
-        FAISS_THROW_IF_NOT(orig_codes != nullptr);
+        FAISS_THROW_IF_NOT(orig_codes);
         search_implem_234<Cfloat>(n, x, k, distances, labels, context);
     } else if (impl >= 12 && impl <= 15) {
         FAISS_THROW_IF_NOT(ntotal < INT_MAX);
@@ -617,8 +617,19 @@ template void IndexFastScan::search_dispatch_implem<false>(
         const FastScanDistancePostProcessing& context) const;
 
 void IndexFastScan::reconstruct(idx_t key, float* recons) const {
-    std::vector<uint8_t> code(code_size, 0);
+    FAISS_THROW_IF_NOT_FMT(
+            key >= 0 && key < ntotal,
+            "IndexFastScan::reconstruct: key %zd out of range (ntotal=%zd)",
+            (size_t)key,
+            (size_t)ntotal);
     std::unique_ptr<CodePacker> packer(get_CodePacker());
+    size_t block_no = (size_t)key / packer->nvec;
+    FAISS_THROW_IF_NOT_MSG(
+            mul_no_overflow(
+                    block_no + 1, packer->block_size, "IndexFastScan codes") <=
+                    codes.size(),
+            "IndexFastScan::reconstruct: packed codes buffer too small");
+    std::vector<uint8_t> code(std::max(code_size, packer->code_size), 0);
     packer->unpack_1(codes.data(), key, code.data());
     sa_decode(1, code.data(), recons);
 }

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -77,7 +77,8 @@ void IndexFlat::range_search(
             range_search_L2sqr(x, get_xb(), d, n, ntotal, radius, result, sel);
             break;
         default:
-            FAISS_THROW_MSG("metric type not supported");
+            IndexFlatCodes::range_search(n, x, radius, result, params);
+            break;
     }
 }
 
@@ -105,13 +106,8 @@ namespace {
 template <SIMDLevel SL>
 struct FlatL2Dis : FlatCodesDistanceComputer {
     size_t d;
-    idx_t nb;
-    const float* b;
-    size_t ndis;
-    size_t npartial_dot_products;
 
     float distance_to_code(const uint8_t* code) final {
-        ndis++;
         return fvec_L2sqr<SL>(q, (float*)code, d);
     }
 
@@ -119,25 +115,21 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
             const idx_t i,
             const uint32_t offset,
             const uint32_t num_components) final override {
-        npartial_dot_products++;
+        const float* b = (const float*)this->codes;
         return fvec_inner_product<SL>(
                 q + offset, b + i * d + offset, num_components);
     }
 
     float symmetric_dis(idx_t i, idx_t j) override {
+        const float* b = (const float*)this->codes;
         return fvec_L2sqr<SL>(b + j * d, b + i * d, d);
     }
 
-    explicit FlatL2Dis(const IndexFlat& storage, const float* q_ = nullptr)
+    explicit FlatL2Dis(const IndexFlat& storage)
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
-                      storage.code_size,
-                      q_),
-              d(storage.d),
-              nb(storage.ntotal),
-              b(storage.get_xb()),
-              ndis(0),
-              npartial_dot_products(0) {}
+                      storage.code_size),
+              d(storage.d) {}
 
     void set_query(const float* x) override {
         q = x;
@@ -153,8 +145,6 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
             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);
@@ -187,8 +177,6 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
             float& dp3,
             const uint32_t offset,
             const uint32_t num_components) final override {
-        npartial_dot_products += 4;
-
         // compute first, assign next
         const float* __restrict y0 =
                 reinterpret_cast<const float*>(codes + idx0 * code_size);
@@ -224,29 +212,21 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
 template <SIMDLevel SL>
 struct FlatIPDis : FlatCodesDistanceComputer {
     size_t d;
-    idx_t nb;
-    const float* q;
-    const float* b;
-    size_t ndis;
 
     float symmetric_dis(idx_t i, idx_t j) final override {
+        const float* b = (const float*)this->codes;
         return fvec_inner_product<SL>(b + j * d, b + i * d, d);
     }
 
     float distance_to_code(const uint8_t* code) final override {
-        ndis++;
         return fvec_inner_product<SL>(q, (const float*)code, d);
     }
 
-    explicit FlatIPDis(const IndexFlat& storage, const float* q_in = nullptr)
+    explicit FlatIPDis(const IndexFlat& storage)
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
                       storage.code_size),
-              d(storage.d),
-              nb(storage.ntotal),
-              q(q_in),
-              b(storage.get_xb()),
-              ndis(0) {}
+              d(storage.d) {}
 
     void set_query(const float* x) override {
         q = x;
@@ -262,8 +242,6 @@ struct FlatIPDis : FlatCodesDistanceComputer {
             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);
@@ -292,12 +270,13 @@ struct FlatIPDis : FlatCodesDistanceComputer {
 FlatCodesDistanceComputer* IndexFlat::get_FlatCodesDistanceComputer() const {
     FlatCodesDistanceComputer* dc = nullptr;
     if (metric_type == METRIC_L2) {
-        with_simd_level([&]<SIMDLevel SL>() { dc = new FlatL2Dis<SL>(*this); });
+        with_simd_level_with_sve(
+                [&]<SIMDLevel SL>() { dc = new FlatL2Dis<SL>(*this); });
     } else if (metric_type == METRIC_INNER_PRODUCT) {
-        with_simd_level([&]<SIMDLevel SL>() { dc = new FlatIPDis<SL>(*this); });
+        with_simd_level_with_sve(
+                [&]<SIMDLevel SL>() { dc = new FlatIPDis<SL>(*this); });
     } else {
-        dc = get_extra_distance_computer(
-                d, metric_type, metric_arg, ntotal, get_xb());
+        dc = get_extra_distance_computer(d, metric_type, metric_arg, get_xb());
     }
     return dc;
 }
@@ -327,16 +306,11 @@ namespace {
 template <SIMDLevel SL>
 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<SL>(q, (float*)code, d);
     }
 
@@ -361,17 +335,11 @@ struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
         return l2norms[i] + l2norms[j] - 2 * dp0;
     }
 
-    explicit FlatL2WithNormsDis(
-            const IndexFlatL2& storage,
-            const float* q_in = nullptr)
+    explicit FlatL2WithNormsDis(const IndexFlatL2& storage)
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
                       storage.code_size),
               d(storage.d),
-              nb(storage.ntotal),
-              q(q_in),
-              b(storage.get_xb()),
-              ndis(0),
               l2norms(storage.cached_l2norms.data()),
               query_l2norm(0) {}
 
@@ -390,8 +358,6 @@ struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
             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);
@@ -440,7 +406,7 @@ FlatCodesDistanceComputer* IndexFlatL2::get_FlatCodesDistanceComputer() const {
     if (metric_type == METRIC_L2) {
         if (!cached_l2norms.empty()) {
             FlatCodesDistanceComputer* dc = nullptr;
-            with_simd_level([&]<SIMDLevel SL>() {
+            with_simd_level_with_sve([&]<SIMDLevel SL>() {
                 dc = new FlatL2WithNormsDis<SL>(*this);
             });
             return dc;
@@ -487,8 +453,7 @@ void IndexFlat1D::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
     FAISS_THROW_IF_NOT_MSG(
             perm.size() == static_cast<size_t>(ntotal),
@@ -741,10 +706,12 @@ void IndexFlatPanorama::reset() {
 }
 
 void IndexFlatPanorama::reconstruct(idx_t key, float* recons) const {
+    FAISS_THROW_IF_NOT(key >= 0 && key < ntotal);
     pano.reconstruct(key, recons, codes.data());
 }
 
 void IndexFlatPanorama::reconstruct_n(idx_t i, idx_t n, float* recons) const {
+    FAISS_THROW_IF_NOT(i >= 0 && i <= ntotal && n >= 0 && n <= ntotal - i);
     Index::reconstruct_n(i, n, recons);
 }
 
@@ -821,7 +788,7 @@ void IndexFlatPanorama::search_subset(
         idx_t k,
         float* distances,
         idx_t* labels) const {
-    with_simd_level([&]<SIMDLevel SL>() {
+    with_simd_level_with_sve([&]<SIMDLevel SL>() {
         with_metric_type(metric_type, [&]<MetricType M>() {
             constexpr bool is_sim = is_similarity_metric(M);
             using C = std::conditional_t<

data/vendor/faiss/faiss/IndexFlat.h

@@ -176,7 +176,7 @@ struct IndexFlatL2Panorama : IndexFlatPanorama {
     explicit IndexFlatL2Panorama(
             idx_t d_in,
             size_t n_levels_in,
-            size_t batch_size_in = 512)
+            size_t batch_size_in = Panorama::kDefaultBatchSize)
             : IndexFlatPanorama(d_in, METRIC_L2, n_levels_in, batch_size_in) {}
 };
 
@@ -189,7 +189,7 @@ struct IndexFlatIPPanorama : IndexFlatPanorama {
     explicit IndexFlatIPPanorama(
             idx_t d_in,
             size_t n_levels_in,
-            size_t batch_size_in = 512)
+            size_t batch_size_in = Panorama::kDefaultBatchSize)
             : IndexFlatPanorama(
                       d_in,
                       METRIC_INNER_PRODUCT,