faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/impl/ClusteringHelpers.cpp

@@ -0,0 +1,244 @@
+/*
+ * 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/impl/ClusteringHelpers.h>
+
+#include <cassert>
+#include <chrono>
+#include <cinttypes>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <limits>
+#include <vector>
+
+#include <omp.h>
+
+#include <faiss/Index.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/random.h>
+
+namespace faiss {
+namespace detail {
+
+uint64_t get_actual_rng_seed(const int seed) {
+    return (seed >= 0)
+            ? seed
+            : static_cast<uint64_t>(std::chrono::high_resolution_clock::now()
+                                            .time_since_epoch()
+                                            .count());
+}
+
+idx_t subsample_training_set(
+        const Clustering& clus,
+        idx_t nx,
+        const uint8_t* x,
+        size_t line_size,
+        const float* weights,
+        uint8_t** x_out,
+        float** weights_out) {
+    FAISS_THROW_IF_NOT(clus.k > 0 && clus.max_points_per_centroid > 0);
+    if (clus.verbose) {
+        printf("Sampling a subset of %zd / %" PRId64 " for training\n",
+               clus.k * clus.max_points_per_centroid,
+               nx);
+    }
+
+    const uint64_t actual_seed = get_actual_rng_seed(clus.seed);
+
+    std::vector<idx_t> perm;
+    if (clus.use_faster_subsampling) {
+        SplitMix64RandomGenerator rng(actual_seed);
+
+        const idx_t new_nx = clus.k * clus.max_points_per_centroid;
+        perm.resize(new_nx);
+        assert(!perm.empty());
+        for (idx_t i = 0; i < new_nx; i++) {
+            perm[i] = rng.rand_int64() % nx;
+        }
+    } else {
+        FAISS_THROW_IF_NOT_FMT(
+                nx <= static_cast<idx_t>(std::numeric_limits<int>::max()),
+                "Dataset too large (%" PRId64
+                ") for standard subsampling; "
+                "set use_faster_subsampling=true",
+                nx);
+        std::vector<int> int_perm(nx);
+        rand_perm(int_perm.data(), nx, actual_seed);
+        perm.assign(int_perm.begin(), int_perm.end());
+    }
+
+    nx = clus.k * clus.max_points_per_centroid;
+    FAISS_THROW_IF_NOT_FMT(
+            perm.size() >= static_cast<size_t>(nx),
+            "subsample_training_set: perm size %zu < required nx %" PRId64,
+            perm.size(),
+            nx);
+    assert(!perm.empty());
+
+    uint8_t* x_new = new uint8_t[nx * line_size];
+    *x_out = x_new;
+
+    for (idx_t i = 0; i < nx; i++) {
+        memcpy(x_new + i * line_size, x + perm[i] * line_size, line_size);
+    }
+    if (weights) {
+        float* weights_new = new float[nx];
+        for (idx_t i = 0; i < nx; i++) {
+            weights_new[i] = weights[perm[i]];
+        }
+        *weights_out = weights_new;
+    } else {
+        *weights_out = nullptr;
+    }
+    return nx;
+}
+
+void compute_centroids(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        const uint8_t* x,
+        const Index* codec,
+        const int64_t* assign,
+        const float* weights,
+        float* hassign,
+        float* centroids) {
+    k -= k_frozen;
+    centroids += k_frozen * d;
+
+    memset(centroids, 0, sizeof(*centroids) * d * k);
+
+    size_t line_size = codec ? codec->sa_code_size() : d * sizeof(float);
+
+#pragma omp parallel
+    {
+        int nt = omp_get_num_threads();
+        int rank = omp_get_thread_num();
+
+        // this thread is taking care of centroids c0:c1
+        size_t c0 = (k * rank) / nt;
+        size_t c1 = (k * (rank + 1)) / nt;
+        std::vector<float> decode_buffer(d);
+
+        for (size_t i = 0; i < n; i++) {
+            int64_t ci = assign[i];
+            FAISS_THROW_IF_NOT_MSG(
+                    ci >= 0 && ci < k + k_frozen, "invalid cluster assignment");
+            ci -= k_frozen;
+            if (ci >= static_cast<int64_t>(c0) &&
+                ci < static_cast<int64_t>(c1)) {
+                float* c = centroids + ci * d;
+                const float* xi;
+                if (!codec) {
+                    xi = reinterpret_cast<const float*>(x + i * line_size);
+                } else {
+                    float* xif = decode_buffer.data();
+                    codec->sa_decode(1, x + i * line_size, xif);
+                    xi = xif;
+                }
+                if (weights) {
+                    float w = weights[i];
+                    hassign[ci] += w;
+                    for (size_t j = 0; j < d; j++) {
+                        c[j] += xi[j] * w;
+                    }
+                } else {
+                    hassign[ci] += 1.0;
+                    for (size_t j = 0; j < d; j++) {
+                        c[j] += xi[j];
+                    }
+                }
+            }
+        }
+    }
+
+#pragma omp parallel for
+    for (idx_t ci = 0; ci < static_cast<idx_t>(k); ci++) {
+        if (hassign[ci] == 0) {
+            continue;
+        }
+        float norm = 1 / hassign[ci];
+        float* c = centroids + ci * d;
+        for (size_t j = 0; j < d; j++) {
+            c[j] *= norm;
+        }
+    }
+}
+
+// a bit above machine epsilon for float16
+static constexpr float EPS = 1.f / 1024.f;
+
+int split_clusters(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        float* hassign,
+        float* centroids) {
+    k -= k_frozen;
+    centroids += k_frozen * d;
+    FAISS_THROW_IF_NOT_MSG(
+            n > k,
+            "split_clusters: n must exceed k to find a non-empty donor centroid");
+
+    size_t nsplit = 0;
+    RandomGenerator rng(1234);
+    for (size_t ci = 0; ci < k; ci++) {
+        if (hassign[ci] == 0) {
+            // Probabilistic donor pick weighted by hassign; deterministic
+            // fallback to the largest cluster if too many iterations pass.
+            size_t cj;
+            size_t max_tries = 10 * k;
+            size_t n_tries = 0;
+            bool found = false;
+            for (cj = 0; n_tries < max_tries; cj = (cj + 1) % k) {
+                float p = (hassign[cj] - 1.0) / (float)(n - k);
+                float r = rng.rand_float();
+                if (r < p) {
+                    found = true;
+                    break;
+                }
+                n_tries++;
+            }
+            if (!found) {
+                // Deterministic fallback: split the largest cluster.
+                cj = 0;
+                for (size_t j = 1; j < k; j++) {
+                    if (hassign[j] > hassign[cj]) {
+                        cj = j;
+                    }
+                }
+            }
+            memcpy(centroids + ci * d,
+                   centroids + cj * d,
+                   sizeof(*centroids) * d);
+
+            /* small symmetric perturbation */
+            for (size_t j = 0; j < d; j++) {
+                if (j % 2 == 0) {
+                    centroids[ci * d + j] *= 1 + EPS;
+                    centroids[cj * d + j] *= 1 - EPS;
+                } else {
+                    centroids[ci * d + j] *= 1 - EPS;
+                    centroids[cj * d + j] *= 1 + EPS;
+                }
+            }
+
+            /* assume even split of the cluster */
+            hassign[ci] = hassign[cj] / 2;
+            hassign[cj] -= hassign[ci];
+            nsplit++;
+        }
+    }
+
+    return static_cast<int>(nsplit);
+}
+
+} // namespace detail
+} // namespace faiss

data/vendor/faiss/faiss/impl/ClusteringHelpers.h

@@ -0,0 +1,94 @@
+/*
+ * 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 <cstddef>
+#include <cstdint>
+
+#include <faiss/Clustering.h>
+#include <faiss/Index.h>
+
+namespace faiss {
+namespace detail {
+
+/** Resolve the actual RNG seed for clustering helpers.
+ *
+ * If `seed >= 0`, returns `seed`. Otherwise returns a high-resolution
+ * timestamp so that callers get a non-deterministic seed.
+ *
+ * @param seed  user-provided seed; negative values request a time-based seed
+ * @return      the resolved seed
+ */
+uint64_t get_actual_rng_seed(const int seed);
+
+/** Subsample a training set down to `clus.k * clus.max_points_per_centroid`
+ * rows.
+ *
+ * Allocates `*x_out` (and `*weights_out` when `weights` is non-null) with
+ * `new[]`; ownership is transferred to the caller.
+ *
+ * @param clus        clustering parameters (reads `k`,
+ * `max_points_per_centroid`, `use_faster_subsampling`, `seed`, `verbose`)
+ * @param nx          number of input training rows
+ * @param x           input training data, row-major, `nx * line_size` bytes
+ * @param line_size   bytes per training row
+ * @param weights     optional per-row weights (length `nx`), or null
+ * @param x_out       output: newly allocated subsampled rows
+ * @param weights_out output: newly allocated subsampled weights, or null
+ * @return            number of rows in the subsampled set
+ */
+idx_t subsample_training_set(
+        const Clustering& clus,
+        idx_t nx,
+        const uint8_t* x,
+        size_t line_size,
+        const float* weights,
+        uint8_t** x_out,
+        float** weights_out);
+
+/** compute centroids as (weighted) sum of training points
+ *
+ * @param x            training vectors, size n * code_size (from codec)
+ * @param codec        how to decode the vectors (if NULL then cast to float*)
+ * @param weights      per-training vector weight, size n (or NULL)
+ * @param assign       nearest centroid for each training vector, size n
+ * @param k_frozen     do not update the k_frozen first centroids
+ * @param centroids    centroid vectors (output only), size k * d
+ * @param hassign      histogram of assignments per centroid (size k),
+ *                     should be 0 on input
+ *
+ */
+void compute_centroids(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        const uint8_t* x,
+        const Index* codec,
+        const int64_t* assign,
+        const float* weights,
+        float* hassign,
+        float* centroids);
+
+/** Handle empty clusters by splitting larger ones.
+ *
+ * It works by slightly changing the centroids to make 2 clusters from
+ * a single one. Takes the same arguments as compute_centroids.
+ *
+ * @return           nb of splitting operations (larger is worse)
+ */
+int split_clusters(
+        size_t d,
+        size_t k,
+        size_t n,
+        size_t k_frozen,
+        float* hassign,
+        float* centroids);
+
+} // namespace detail
+} // namespace faiss

data/vendor/faiss/faiss/impl/ClusteringInitialization.cpp

@@ -142,8 +142,8 @@ size_t sample_from_cumsum(
 
 } // namespace
 
-ClusteringInitialization::ClusteringInitialization(size_t d, size_t k)
-        : d(d), k(k) {}
+ClusteringInitialization::ClusteringInitialization(size_t d_in, size_t k_in)
+        : d(d_in), k(k_in) {}
 
 void ClusteringInitialization::init_centroids(
         size_t n,
@@ -221,27 +221,27 @@ void ClusteringInitialization::init_kmeans_plus_plus(
     std::vector<double> cumsum(n);
 
     // Select remaining centroids using D² sampling
-    for (size_t c = result.first_new_centroid_idx; c < k; c++) {
-        // Compute cumulative sum
-        cumsum[0] = min_distances[0];
-        for (size_t i = 1; i < n; i++) {
-            cumsum[i] = cumsum[i - 1] + min_distances[i];
-        }
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (size_t c = result.first_new_centroid_idx; c < k; c++) {
+            // Compute cumulative sum
+            cumsum[0] = min_distances[0];
+            for (size_t i = 1; i < n; i++) {
+                cumsum[i] = cumsum[i - 1] + min_distances[i];
+            }
 
-        // Sample using precomputed cumsum
-        size_t next_idx = sample_from_cumsum(cumsum, rng);
+            // Sample using precomputed cumsum
+            size_t next_idx = sample_from_cumsum(cumsum, rng);
 
-        float* new_centroid = centroids + c * d;
-        std::memcpy(new_centroid, x + next_idx * d, d * sizeof(float));
+            float* new_centroid = centroids + c * d;
+            std::memcpy(new_centroid, x + next_idx * d, d * sizeof(float));
 
-        // Update min distances incrementally
-        with_simd_level([&]<SIMDLevel SL>() {
+            // Update min distances incrementally
             for (size_t i = 0; i < n; i++) {
                 double dist = fvec_L2sqr<SL>(x + i * d, new_centroid, d);
                 min_distances[i] = std::min(min_distances[i], dist);
             }
-        });
-    }
+        }
+    });
 }
 
 void ClusteringInitialization::init_afkmc2(

data/vendor/faiss/faiss/impl/CodePacker.cpp

@@ -33,10 +33,10 @@ void CodePacker::unpack_all(const uint8_t* block, uint8_t* flat_codes) const {
  * CodePackerFlat
  */
 
-CodePackerFlat::CodePackerFlat(size_t code_size) {
-    this->code_size = code_size;
+CodePackerFlat::CodePackerFlat(size_t code_size_in) {
+    this->code_size = code_size_in;
     nvec = 1;
-    block_size = code_size;
+    block_size = code_size_in;
 }
 
 void CodePackerFlat::pack_all(const uint8_t* flat_codes, uint8_t* block) const {

data/vendor/faiss/faiss/impl/CodePackerRaBitQ.cpp

@@ -6,7 +6,7 @@
  */
 
 #include <faiss/impl/CodePackerRaBitQ.h>
-#include <faiss/impl/pq4_fast_scan.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
 
 #include <cstring>
 

data/vendor/faiss/faiss/impl/DistanceComputer.h

@@ -66,8 +66,8 @@ struct NegativeDistanceComputer : DistanceComputer {
     /// owned by this
     DistanceComputer* basedis;
 
-    explicit NegativeDistanceComputer(DistanceComputer* basedis)
-            : basedis(basedis) {}
+    explicit NegativeDistanceComputer(DistanceComputer* basedis_)
+            : basedis(basedis_) {}
 
     void set_query(const float* x) override {
         basedis->set_query(x);
@@ -116,13 +116,13 @@ struct FlatCodesDistanceComputer : DistanceComputer {
     const float* q = nullptr; // not used in all distance computers
 
     FlatCodesDistanceComputer(
-            const uint8_t* codes,
-            size_t code_size,
-            const float* q = nullptr)
-            : codes(codes), code_size(code_size), q(q) {}
+            const uint8_t* codes_,
+            size_t code_size_,
+            const float* q_ = nullptr)
+            : codes(codes_), code_size(code_size_), q(q_) {}
 
-    explicit FlatCodesDistanceComputer(const float* q)
-            : codes(nullptr), code_size(0), q(q) {}
+    explicit FlatCodesDistanceComputer(const float* q_)
+            : codes(nullptr), code_size(0), q(q_) {}
 
     FlatCodesDistanceComputer() : codes(nullptr), code_size(0), q(nullptr) {}
 

data/vendor/faiss/faiss/impl/FaissAssert.h

@@ -78,9 +78,12 @@
     do {                                                       \
         std::string __s;                                       \
         int __size = snprintf(nullptr, 0, FMT, __VA_ARGS__);   \
-        __s.resize(__size + 1);                                \
-        snprintf(&__s[0], __s.size(), FMT, __VA_ARGS__);       \
-        throw ::faiss::FaissException(                         \
+        if (__size > 0) {                                      \
+            __s.resize(static_cast<size_t>(__size) + 1);       \
+            snprintf(&__s[0], __s.size(), FMT, __VA_ARGS__);   \
+            __s.resize(static_cast<size_t>(__size));           \
+        }                                                      \
+        throw faiss::FaissException(                           \
                 __s, __PRETTY_FUNCTION__, __FILE__, __LINE__); \
     } while (false)
 

data/vendor/faiss/faiss/impl/FaissException.h

@@ -46,9 +46,8 @@ void handleExceptions(
 struct TransformedVectors {
     const float* x;
     bool own_x;
-    TransformedVectors(const float* x_orig, const float* x) : x(x) {
-        own_x = x_orig != x;
-    }
+    TransformedVectors(const float* x_orig, const float* x_in)
+            : x(x_in), own_x(x_orig != x_in) {}
 
     ~TransformedVectors() {
         if (own_x) {
@@ -60,6 +59,54 @@ struct TransformedVectors {
 /// make typeids more readable
 std::string demangle_cpp_symbol(const char* name);
 
+/// Capture the current exception into `ex` if no prior exception has been
+/// recorded.  Call from a catch block inside an OpenMP parallel region.
+/// Uses `#pragma omp critical` to serialize access to `ex`.
+///
+/// The optional `cleanup` callable runs inside the critical section
+/// alongside the exception capture, so that side-effects visible to
+/// other threads (e.g. setting an interrupt flag) are serialized with
+/// the exception_ptr write.
+///
+/// Usage:
+///   std::exception_ptr ex;
+///   bool interrupt = false;
+///   #pragma omp parallel
+///   {
+///       try { ... } catch (...) {
+///           omp_capture_exception(ex, [&] { interrupt = true; });
+///       }
+///   }
+///   omp_rethrow_if_exception(ex);
+inline void omp_capture_exception(std::exception_ptr& ex) {
+#pragma omp critical(faiss_omp_exception)
+    {
+        if (!ex) {
+            ex = std::current_exception();
+        }
+    }
+}
+
+/// Overload with cleanup that runs inside the critical section.
+template <typename Cleanup>
+inline void omp_capture_exception(std::exception_ptr& ex, Cleanup&& cleanup) {
+#pragma omp critical(faiss_omp_exception)
+    {
+        cleanup();
+        if (!ex) {
+            ex = std::current_exception();
+        }
+    }
+}
+
+/// Rethrow the captured exception, if any.  Call on the main thread
+/// after the parallel region completes.
+inline void omp_rethrow_if_exception(std::exception_ptr& ex) {
+    if (ex) {
+        std::rethrow_exception(ex);
+    }
+}
+
 } // namespace faiss
 
 #endif