faiss 0.5.0 → 0.5.1

data/vendor/faiss/faiss/impl/Panorama.cpp

@@ -0,0 +1,193 @@
+/*
+ * 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/Panorama.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+#include <vector>
+
+namespace faiss {
+
+/**************************************************************
+ * Panorama structure implementation
+ **************************************************************/
+
+Panorama::Panorama(size_t code_size, size_t n_levels, size_t batch_size)
+        : code_size(code_size), n_levels(n_levels), batch_size(batch_size) {
+    set_derived_values();
+}
+
+void Panorama::set_derived_values() {
+    this->d = code_size / sizeof(float);
+    this->level_width_floats = ((d + n_levels - 1) / n_levels);
+    this->level_width = this->level_width_floats * sizeof(float);
+}
+
+/**
+ * @brief Copy codes to level-oriented layout
+ * @param codes The base pointer to codes
+ * @param offset Where to start writing new data (in number of vectors)
+ * @param n_entry The number of new vectors to write
+ * @param code The new vector data
+ */
+void Panorama::copy_codes_to_level_layout(
+        uint8_t* codes,
+        size_t offset,
+        size_t n_entry,
+        const uint8_t* code) {
+    for (size_t entry_idx = 0; entry_idx < n_entry; entry_idx++) {
+        size_t current_pos = offset + entry_idx;
+
+        // Determine which batch we're in and position within that batch.
+        size_t batch_no = current_pos / batch_size;
+        size_t pos_in_batch = current_pos % batch_size;
+
+        // Copy entry into level-oriented layout for this batch.
+        size_t batch_offset = batch_no * batch_size * code_size;
+        for (size_t level = 0; level < n_levels; level++) {
+            size_t level_offset = level * level_width * batch_size;
+            size_t start_byte = level * level_width;
+            size_t actual_level_width =
+                    std::min(level_width, code_size - level * level_width);
+
+            const uint8_t* src = code + entry_idx * code_size + start_byte;
+            uint8_t* dest = codes + batch_offset + level_offset +
+                    pos_in_batch * actual_level_width;
+
+            memcpy(dest, src, actual_level_width);
+        }
+    }
+}
+
+void Panorama::compute_cumulative_sums(
+        float* cumsum_base,
+        size_t offset,
+        size_t n_entry,
+        const float* vectors) {
+    std::vector<float> suffix_sums(d + 1);
+
+    for (size_t entry_idx = 0; entry_idx < n_entry; entry_idx++) {
+        size_t current_pos = offset + entry_idx;
+        size_t batch_no = current_pos / batch_size;
+        size_t pos_in_batch = current_pos % batch_size;
+
+        const float* vector = vectors + entry_idx * d;
+
+        // Compute suffix sums of squared values.
+        suffix_sums[d] = 0.0f;
+        for (int j = d - 1; j >= 0; j--) {
+            float squared_val = vector[j] * vector[j];
+            suffix_sums[j] = suffix_sums[j + 1] + squared_val;
+        }
+
+        // Store cumulative sums in batch-oriented layout.
+        size_t cumsum_batch_offset = batch_no * batch_size * (n_levels + 1);
+
+        for (size_t level = 0; level < n_levels; level++) {
+            size_t start_idx = level * level_width_floats;
+            size_t cumsum_offset =
+                    cumsum_batch_offset + level * batch_size + pos_in_batch;
+            if (start_idx < d) {
+                cumsum_base[cumsum_offset] = std::sqrt(suffix_sums[start_idx]);
+            } else {
+                cumsum_base[cumsum_offset] = 0.0f;
+            }
+        }
+
+        // Last level sum is always 0.
+        size_t cumsum_offset =
+                cumsum_batch_offset + n_levels * batch_size + pos_in_batch;
+        cumsum_base[cumsum_offset] = 0.0f;
+    }
+}
+
+void Panorama::compute_query_cum_sums(const float* query, float* query_cum_sums)
+        const {
+    std::vector<float> suffix_sums(d + 1);
+    suffix_sums[d] = 0.0f;
+
+    for (int j = d - 1; j >= 0; j--) {
+        float squared_val = query[j] * query[j];
+        suffix_sums[j] = suffix_sums[j + 1] + squared_val;
+    }
+
+    for (size_t level = 0; level < n_levels; level++) {
+        size_t start_idx = level * level_width_floats;
+        if (start_idx < d) {
+            query_cum_sums[level] = std::sqrt(suffix_sums[start_idx]);
+        } else {
+            query_cum_sums[level] = 0.0f;
+        }
+    }
+
+    query_cum_sums[n_levels] = 0.0f;
+}
+
+void Panorama::reconstruct(idx_t key, float* recons, const uint8_t* codes_base)
+        const {
+    uint8_t* recons_buffer = reinterpret_cast<uint8_t*>(recons);
+
+    size_t batch_no = key / batch_size;
+    size_t pos_in_batch = key % batch_size;
+    size_t batch_offset = batch_no * batch_size * code_size;
+
+    for (size_t level = 0; level < n_levels; level++) {
+        size_t level_offset = level * level_width * batch_size;
+        const uint8_t* src = codes_base + batch_offset + level_offset +
+                pos_in_batch * level_width;
+        uint8_t* dest = recons_buffer + level * level_width;
+        size_t copy_size =
+                std::min(level_width, code_size - level * level_width);
+        memcpy(dest, src, copy_size);
+    }
+}
+
+void Panorama::copy_entry(
+        uint8_t* dest_codes,
+        uint8_t* src_codes,
+        float* dest_cum_sums,
+        float* src_cum_sums,
+        size_t dest_idx,
+        size_t src_idx) const {
+    // Calculate positions
+    size_t src_batch_no = src_idx / batch_size;
+    size_t src_pos_in_batch = src_idx % batch_size;
+    size_t dest_batch_no = dest_idx / batch_size;
+    size_t dest_pos_in_batch = dest_idx % batch_size;
+
+    // Calculate offsets
+    size_t src_batch_offset = src_batch_no * batch_size * code_size;
+    size_t dest_batch_offset = dest_batch_no * batch_size * code_size;
+    size_t src_cumsum_batch_offset = src_batch_no * batch_size * (n_levels + 1);
+    size_t dest_cumsum_batch_offset =
+            dest_batch_no * batch_size * (n_levels + 1);
+
+    for (size_t level = 0; level < n_levels; level++) {
+        // Copy code
+        size_t level_offset = level * level_width * batch_size;
+        size_t actual_level_width =
+                std::min(level_width, code_size - level * level_width);
+
+        const uint8_t* src = src_codes + src_batch_offset + level_offset +
+                src_pos_in_batch * actual_level_width;
+        uint8_t* dest = dest_codes + dest_batch_offset + level_offset +
+                dest_pos_in_batch * actual_level_width;
+        memcpy(dest, src, actual_level_width);
+
+        // Copy cum_sums
+        size_t cumsum_level_offset = level * batch_size;
+
+        const size_t src_offset = src_cumsum_batch_offset +
+                cumsum_level_offset + src_pos_in_batch;
+        size_t dest_offset = dest_cumsum_batch_offset + cumsum_level_offset +
+                dest_pos_in_batch;
+        dest_cum_sums[dest_offset] = src_cum_sums[src_offset];
+    }
+}
+} // namespace faiss

data/vendor/faiss/faiss/impl/Panorama.h

@@ -0,0 +1,204 @@
+/*
+ * 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.
+ */
+
+// -*- c++ -*-
+
+#ifndef FAISS_PANORAMA_H
+#define FAISS_PANORAMA_H
+
+#include <faiss/impl/IDSelector.h>
+#include <faiss/impl/PanoramaStats.h>
+#include <faiss/utils/distances.h>
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+namespace faiss {
+
+/**
+ * Implements the core logic of Panorama-based refinement.
+ * arXiv: https://arxiv.org/abs/2510.00566
+ *
+ * Panorama partitions the dimensions of all vectors into L contiguous levels.
+ * During the refinement stage of ANNS, it computes distances between the query
+ * and its candidates level-by-level. After processing each level, it prunes the
+ * candidates whose lower bound exceeds the k-th best distance.
+ *
+ * In order to enable speedups, the dimensions (or codes) of each vector are
+ * stored in a batched, level-major manner. Within each batch of b vectors, the
+ * dimensions corresponding to level 1 will be stored first (for all elements in
+ * that batch), followed by level 2, and so on. This allows for efficient memory
+ * access patterns.
+ *
+ * Coupled with the appropriate orthogonal PreTransform (e.g. PCA, Cayley,
+ * etc.), Panorama can prune the vast majority of dimensions, greatly
+ * accelerating the refinement stage.
+ */
+struct Panorama {
+    size_t d = 0;
+    size_t code_size = 0;
+    size_t n_levels = 0;
+    size_t level_width = 0;
+    size_t level_width_floats = 0;
+    size_t batch_size = 0;
+
+    explicit Panorama(size_t code_size, size_t n_levels, size_t batch_size);
+
+    void set_derived_values();
+
+    /// Helper method to copy codes into level-oriented batch layout at a given
+    /// offset in the list.
+    void copy_codes_to_level_layout(
+            uint8_t* codes,
+            size_t offset,
+            size_t n_entry,
+            const uint8_t* code);
+
+    /// Helper method to compute the cumulative sums of the codes.
+    /// The cumsums also follow the level-oriented batch layout to minimize the
+    /// number of random memory accesses.
+    void compute_cumulative_sums(
+            float* cumsum_base,
+            size_t offset,
+            size_t n_entry,
+            const float* vectors);
+
+    /// Compute the cumulative sums of the query vector.
+    void compute_query_cum_sums(const float* query, float* query_cum_sums)
+            const;
+
+    /// Copy single entry (code and cum_sum) from one location to another.
+    void copy_entry(
+            uint8_t* dest_codes,
+            uint8_t* src_codes,
+            float* dest_cum_sums,
+            float* src_cum_sums,
+            size_t dest_idx,
+            size_t src_idx) const;
+
+    /// Panorama's core progressive filtering algorithm:
+    /// Process vectors in batches for cache efficiency. For each batch:
+    /// 1. Apply ID selection filter and initialize distances
+    /// (||y||^2 + ||x||^2).
+    /// 2. Maintain an "active set" of candidate indices that haven't been
+    /// pruned yet.
+    /// 3. For each level, refine distances incrementally and compact the active
+    /// set:
+    ///    - Compute dot product for current level: exact_dist -= 2*<x,y>.
+    ///    - Use Cauchy-Schwarz bound on remaining levels to get lower bound
+    ///    - Prune candidates whose lower bound exceeds k-th best distance.
+    ///    - Compact active_indices to remove pruned candidates (branchless)
+    /// 4. After all levels, survivors are exact distances; update heap.
+    /// This achieves early termination while maintaining SIMD-friendly
+    /// sequential access patterns in the level-oriented storage layout.
+    template <typename C>
+    size_t progressive_filter_batch(
+            const uint8_t* codes_base,
+            const float* cum_sums,
+            const float* query,
+            const float* query_cum_sums,
+            size_t batch_no,
+            size_t list_size,
+            const IDSelector* sel,
+            const idx_t* ids,
+            bool use_sel,
+            std::vector<uint32_t>& active_indices,
+            std::vector<float>& exact_distances,
+            float threshold,
+            PanoramaStats& local_stats) const;
+
+    void reconstruct(idx_t key, float* recons, const uint8_t* codes_base) const;
+};
+
+template <typename C>
+size_t Panorama::progressive_filter_batch(
+        const uint8_t* codes_base,
+        const float* cum_sums,
+        const float* query,
+        const float* query_cum_sums,
+        size_t batch_no,
+        size_t list_size,
+        const IDSelector* sel,
+        const idx_t* ids,
+        bool use_sel,
+        std::vector<uint32_t>& active_indices,
+        std::vector<float>& exact_distances,
+        float threshold,
+        PanoramaStats& local_stats) const {
+    size_t batch_start = batch_no * batch_size;
+    size_t curr_batch_size = std::min(list_size - batch_start, batch_size);
+
+    size_t cumsum_batch_offset = batch_no * batch_size * (n_levels + 1);
+    const float* batch_cum_sums = cum_sums + cumsum_batch_offset;
+    const float* level_cum_sums = batch_cum_sums + batch_size;
+    float q_norm = query_cum_sums[0] * query_cum_sums[0];
+
+    size_t batch_offset = batch_no * batch_size * code_size;
+    const uint8_t* storage_base = codes_base + batch_offset;
+
+    // Initialize active set with ID-filtered vectors.
+    size_t num_active = 0;
+    for (size_t i = 0; i < curr_batch_size; i++) {
+        size_t global_idx = batch_start + i;
+        idx_t id = (ids == nullptr) ? global_idx : ids[global_idx];
+        bool include = !use_sel || sel->is_member(id);
+
+        active_indices[num_active] = i;
+        float cum_sum = batch_cum_sums[i];
+        exact_distances[i] = cum_sum * cum_sum + q_norm;
+
+        num_active += include;
+    }
+
+    if (num_active == 0) {
+        return 0;
+    }
+
+    size_t total_active = num_active;
+    for (size_t level = 0; level < n_levels; level++) {
+        local_stats.total_dims_scanned += num_active;
+        local_stats.total_dims += total_active;
+
+        float query_cum_norm = query_cum_sums[level + 1];
+
+        size_t level_offset = level * level_width * batch_size;
+        const float* level_storage =
+                (const float*)(storage_base + level_offset);
+
+        size_t next_active = 0;
+        for (size_t i = 0; i < num_active; i++) {
+            uint32_t idx = active_indices[i];
+            size_t actual_level_width = std::min(
+                    level_width_floats, d - level * level_width_floats);
+
+            const float* yj = level_storage + idx * actual_level_width;
+            const float* query_level = query + level * level_width_floats;
+
+            float dot_product =
+                    fvec_inner_product(query_level, yj, actual_level_width);
+
+            exact_distances[idx] -= 2.0f * dot_product;
+
+            float cum_sum = level_cum_sums[idx];
+            float cauchy_schwarz_bound = 2.0f * cum_sum * query_cum_norm;
+            float lower_bound = exact_distances[idx] - cauchy_schwarz_bound;
+
+            active_indices[next_active] = idx;
+            next_active += C::cmp(threshold, lower_bound) ? 1 : 0;
+        }
+
+        num_active = next_active;
+        level_cum_sums += batch_size;
+    }
+
+    return num_active;
+}
+} // namespace faiss
+
+#endif

data/vendor/faiss/faiss/impl/RaBitQStats.cpp

@@ -0,0 +1,29 @@
+/*
+ * 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/RaBitQStats.h>
+
+namespace faiss {
+
+// NOLINTNEXTLINE(facebook-avoid-non-const-global-variables)
+RaBitQStats rabitq_stats;
+
+void RaBitQStats::reset() {
+    n_1bit_evaluations = 0;
+    n_multibit_evaluations = 0;
+}
+
+double RaBitQStats::skip_percentage() const {
+    const size_t copy_n_1bit_evaluations = n_1bit_evaluations;
+    const size_t copy_n_multibit_evaluations = n_multibit_evaluations;
+    return copy_n_1bit_evaluations > 0
+            ? 100.0 * (copy_n_1bit_evaluations - copy_n_multibit_evaluations) /
+                    copy_n_1bit_evaluations
+            : 0.0;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/RaBitQStats.h

@@ -0,0 +1,56 @@
+/*
+ * 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/impl/platform_macros.h>
+#include <cstddef>
+
+namespace faiss {
+
+/// Statistics for RaBitQ multi-bit two-stage search.
+///
+/// These stats are ONLY collected for multi-bit mode (nb_bits > 1).
+/// In 1-bit mode, there is no two-stage filtering - all candidates are
+/// evaluated with a single distance computation, so there is nothing
+/// meaningful to track. For 1-bit mode, both counters remain 0.
+///
+/// Multi-bit mode uses a two-stage search:
+///   Stage 1: Compute 1-bit lower bound distance for all candidates
+///   Stage 2: Compute full multi-bit distance only for promising candidates
+///
+/// The skip_percentage() metric measures filtering effectiveness:
+/// how many candidates were filtered out by the 1-bit lower bound
+/// without needing the more expensive multi-bit distance computation.
+///
+/// WARNING: Statistics are not robust to internal threading nor to
+/// concurrent RaBitQ searches. Use these values in a single-threaded
+/// context to accurately gauge RaBitQ's filtering effectiveness.
+/// Call reset() before search, then read stats after search completes.
+struct RaBitQStats {
+    /// Number of candidates evaluated using 1-bit (lower bound) distance.
+    /// This is the first stage of two-stage search in multi-bit mode.
+    /// Always 0 in 1-bit mode (stats not tracked).
+    size_t n_1bit_evaluations = 0;
+
+    /// Number of candidates that passed 1-bit filtering and required
+    /// full multi-bit distance computation (second stage).
+    /// Always 0 in 1-bit mode (stats not tracked).
+    size_t n_multibit_evaluations = 0;
+
+    void reset();
+
+    /// Compute percentage of candidates skipped (filtered out by 1-bit stage).
+    /// Returns 0 if no candidates were evaluated (including 1-bit mode).
+    double skip_percentage() const;
+};
+
+/// Global stats for RaBitQ indexes
+// NOLINTNEXTLINE(facebook-avoid-non-const-global-variables)
+FAISS_API extern RaBitQStats rabitq_stats;
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/RaBitQUtils.cpp

@@ -16,6 +16,18 @@
 namespace faiss {
 namespace rabitq_utils {
 
+// Verify no unexpected padding in structures used for per-vector storage.
+// These checks ensure compute_per_vector_storage_size() remains accurate.
+static_assert(
+        sizeof(SignBitFactors) == 8,
+        "SignBitFactors has unexpected padding");
+static_assert(
+        sizeof(SignBitFactorsWithError) == 12,
+        "SignBitFactorsWithError has unexpected padding");
+static_assert(
+        sizeof(ExtraBitsFactors) == 8,
+        "ExtraBitsFactors has unexpected padding");
+
 // Ideal quantizer radii for quantizers of 1..8 bits, optimized to minimize
 // L2 reconstruction error.
 const float Z_MAX_BY_QB[8] = {
@@ -54,13 +66,16 @@ void compute_vector_intermediate_values(
     }
 }
 
-FactorsData compute_factors_from_intermediates(
+SignBitFactorsWithError compute_factors_from_intermediates(
         float norm_L2sqr,
         float or_L2sqr,
         float dp_oO,
         size_t d,
-        MetricType metric_type) {
+        MetricType metric_type,
+        bool compute_error) {
     constexpr float epsilon = std::numeric_limits<float>::epsilon();
+    constexpr float kConstEpsilon =
+            1.9f; // Error bound constant from RaBitQ paper
     const float inv_d_sqrt =
             (d == 0) ? 1.0f : (1.0f / std::sqrt(static_cast<float>(d)));
 
@@ -72,25 +87,57 @@ FactorsData compute_factors_from_intermediates(
     const float inv_dp_oO =
             (std::abs(normalized_dp) < epsilon) ? 1.0f : (1.0f / normalized_dp);
 
-    FactorsData factors;
+    SignBitFactorsWithError factors;
     factors.or_minus_c_l2sqr = (metric_type == MetricType::METRIC_INNER_PRODUCT)
             ? (norm_L2sqr - or_L2sqr)
             : norm_L2sqr;
     factors.dp_multiplier = inv_dp_oO * sqrt_norm_L2;
 
+    // Compute error bound only if needed (skip for 1-bit mode)
+    if (compute_error) {
+        const float xu_cb_norm_sqr = static_cast<float>(d) * 0.25f;
+        const float ip_resi_xucb = 0.5f * dp_oO;
+
+        float tmp_error = 0.0f;
+        if (std::abs(ip_resi_xucb) > epsilon) {
+            const float ratio_sq = (norm_L2sqr * xu_cb_norm_sqr) /
+                    (ip_resi_xucb * ip_resi_xucb);
+            if (ratio_sq > 1.0f) {
+                if (d == 1) {
+                    tmp_error = sqrt_norm_L2 * kConstEpsilon *
+                            std::sqrt(ratio_sq - 1.0f);
+                } else {
+                    tmp_error = sqrt_norm_L2 * kConstEpsilon *
+                            std::sqrt((ratio_sq - 1.0f) /
+                                      static_cast<float>(d - 1));
+                }
+            }
+        }
+
+        // Apply metric-specific multiplier
+        if (metric_type == MetricType::METRIC_L2) {
+            factors.f_error = 2.0f * tmp_error;
+        } else if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
+            factors.f_error = 1.0f * tmp_error;
+        } else {
+            factors.f_error = 0.0f;
+        }
+    }
+
     return factors;
 }
 
-FactorsData compute_vector_factors(
+SignBitFactorsWithError compute_vector_factors(
         const float* x,
         size_t d,
         const float* centroid,
-        MetricType metric_type) {
+        MetricType metric_type,
+        bool compute_error) {
     float norm_L2sqr, or_L2sqr, dp_oO;
     compute_vector_intermediate_values(
             x, d, centroid, norm_L2sqr, or_L2sqr, dp_oO);
     return compute_factors_from_intermediates(
-            norm_L2sqr, or_L2sqr, dp_oO, d, metric_type);
+            norm_L2sqr, or_L2sqr, dp_oO, d, metric_type, compute_error);
 }
 
 QueryFactorsData compute_query_factors(
@@ -113,6 +160,7 @@ QueryFactorsData compute_query_factors(
     } else {
         query_factors.qr_to_c_L2sqr = fvec_norm_L2sqr(query, d);
     }
+    query_factors.g_error = std::sqrt(query_factors.qr_to_c_L2sqr);
 
     // Rotate the query (subtract centroid)
     rotated_q.resize(d);