faiss 0.6.1 → 0.6.3

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);
@@ -296,8 +274,7 @@ FlatCodesDistanceComputer* IndexFlat::get_FlatCodesDistanceComputer() const {
     } else if (metric_type == METRIC_INNER_PRODUCT) {
         with_simd_level([&]<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 +304,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 +333,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 +356,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);
@@ -487,8 +451,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),