faiss 0.6.0 → 0.6.2

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

@@ -18,10 +18,201 @@
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 #include <vector>
 
+#if defined(COMPILE_SIMD_AVX2) && defined(__AVX2__) && defined(__BMI2__)
+#include <immintrin.h>
+#endif
+
 namespace faiss {
 
+#ifndef SWIG
+
+/// Compute dot products between query_level and active vectors.
+///
+/// @tparam AllActive  If true, vectors are at sequential positions 0..N-1
+///                    (first level, full batch). If false, positions come
+///                    from active_indices (subsequent levels after pruning).
+/// @tparam LevelWidth Compile-time level width in floats (0 = use runtime
+///                    level_width_dims). Enables full loop unrolling.
+// Skip pragmas under nvcc: its EDG frontend warns on `#pragma GCC optimize`
+// (#1675-D) for every `.cu` that transitively includes this header. These
+// templates are CPU-only, so the hint is irrelevant during nvcc parse.
+#if !defined(__NVCC__)
+FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
+#endif
+template <bool AllActive = false, size_t LevelWidth = 0>
+static inline void compute_level_dot_kernel(
+        const float* FAISS_RESTRICT query_level,
+        const float* FAISS_RESTRICT level_storage,
+        const uint32_t* active_indices,
+        const size_t num_active,
+        const size_t level_width_dims,
+        float* FAISS_RESTRICT dot_products) {
+    const size_t width = LevelWidth > 0 ? LevelWidth : level_width_dims;
+    size_t i = 0;
+    for (; i + 4 <= num_active; i += 4) {
+        const float* y0 = level_storage +
+                (AllActive ? (i + 0) : active_indices[i + 0]) * width;
+        const float* y1 = level_storage +
+                (AllActive ? (i + 1) : active_indices[i + 1]) * width;
+        const float* y2 = level_storage +
+                (AllActive ? (i + 2) : active_indices[i + 2]) * width;
+        const float* y3 = level_storage +
+                (AllActive ? (i + 3) : active_indices[i + 3]) * width;
+
+        float dp0 = 0, dp1 = 0, dp2 = 0, dp3 = 0;
+        FAISS_PRAGMA_IMPRECISE_LOOP
+        for (size_t j = 0; j < width; j++) {
+            float q = query_level[j];
+            dp0 += q * y0[j];
+            dp1 += q * y1[j];
+            dp2 += q * y2[j];
+            dp3 += q * y3[j];
+        }
+
+        dot_products[i + 0] = dp0;
+        dot_products[i + 1] = dp1;
+        dot_products[i + 2] = dp2;
+        dot_products[i + 3] = dp3;
+    }
+    for (; i < num_active; i++) {
+        const float* yj =
+                level_storage + (AllActive ? i : active_indices[i]) * width;
+        float dp = 0;
+        FAISS_PRAGMA_IMPRECISE_LOOP
+        for (size_t j = 0; j < width; j++) {
+            dp += query_level[j] * yj[j];
+        }
+        dot_products[i] = dp;
+    }
+}
+#if !defined(__NVCC__)
+FAISS_PRAGMA_IMPRECISE_FUNCTION_END
+#endif
+
+/// Update exact distances with the current level's dot products, then apply
+/// Panorama pruning: for each active vector, compute a lower bound on
+/// the final distance and mark it for removal if it cannot beat the current
+/// threshold. Writes 0/1 into active_byteset for subsequent compaction.
+///
+/// Uses `if constexpr` on C::is_max rather than C::cmp() to ensure the
+/// comparison autovectorizes (C::cmp generates scalar function calls).
+#if !defined(__NVCC__)
+FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
+#endif
+template <bool AllActive, typename C, MetricType M>
+static inline void prune_kernel(
+        float* FAISS_RESTRICT exact_distances,
+        const float* FAISS_RESTRICT dot_buffer,
+        const float* FAISS_RESTRICT level_cum_sums,
+        uint8_t* FAISS_RESTRICT active_byteset,
+        const uint32_t* FAISS_RESTRICT active_indices,
+        const uint32_t num_active,
+        const float query_cum_norm,
+        const float threshold) {
+    FAISS_PRAGMA_IMPRECISE_LOOP
+    for (uint32_t i = 0; i < num_active; i++) {
+        uint32_t idx = AllActive ? i : active_indices[i];
+        if constexpr (M == METRIC_INNER_PRODUCT) {
+            exact_distances[idx] += dot_buffer[i];
+        } else {
+            exact_distances[idx] -= 2.0f * dot_buffer[i];
+        }
+
+        float cum_sum = level_cum_sums[idx];
+        float cauchy_schwarz_bound;
+        if constexpr (M == METRIC_INNER_PRODUCT) {
+            cauchy_schwarz_bound = -cum_sum * query_cum_norm;
+        } else {
+            cauchy_schwarz_bound = 2.0f * cum_sum * query_cum_norm;
+        }
+
+        float lower_bound = exact_distances[idx] - cauchy_schwarz_bound;
+        if constexpr (C::is_max) {
+            active_byteset[i] = (threshold > lower_bound) ? 1 : 0;
+        } else {
+            active_byteset[i] = (threshold < lower_bound) ? 1 : 0;
+        }
+    }
+}
+#if !defined(__NVCC__)
+FAISS_PRAGMA_IMPRECISE_FUNCTION_END
+#endif
+
+/// Compact active_indices in-place, removing entries where active_byteset[i]
+/// is zero. Returns the new count of active elements. Uses a branchless BMI2 +
+/// AVX2 fast path (8 elements/iteration via _pext_u64 permutation) with a
+/// scalar fallback for the tail and non-x86 platforms.
+inline size_t compact_active_kernel(
+        uint32_t* active_indices,
+        const uint8_t* FAISS_RESTRICT active_byteset,
+        const size_t num_active) {
+    size_t next_active = 0;
+    size_t i = 0;
+
+#if defined(COMPILE_SIMD_AVX2) && defined(__AVX2__) && defined(__BMI2__)
+    for (; i + 8 <= num_active; i += 8) {
+        uint64_t bytes;
+        memcpy(&bytes, &active_byteset[i], 8);
+
+        uint64_t expanded = bytes * 0xFFULL;
+        uint64_t packed = _pext_u64(0x0706050403020100ULL, expanded);
+
+        __m256i perm = _mm256_cvtepu8_epi32(_mm_cvtsi64_si128((int64_t)packed));
+        __m256i data = _mm256_loadu_si256((const __m256i*)&active_indices[i]);
+        __m256i compacted = _mm256_permutevar8x32_epi32(data, perm);
+        _mm256_storeu_si256((__m256i*)&active_indices[next_active], compacted);
+
+        next_active += __builtin_popcountll(bytes);
+    }
+#endif
+
+    for (; i < num_active; i++) {
+        active_indices[next_active] = active_indices[i];
+        next_active += active_byteset[i] ? 1 : 0;
+    }
+
+    return next_active;
+}
+
+/// Compile-time dispatch: converts a runtime `width` value into a template
+/// parameter by generating an if-else chain over [Lo, Hi] in steps of Step.
+/// Falls through to LevelWidth=0 (runtime path) if no specialization matches.
+/// Allows for specialization of common level widths.
+namespace detail {
+template <size_t Lo, size_t Hi, size_t Step, typename Lambda>
+inline auto dispatch_width(size_t width, Lambda&& fn) {
+    if constexpr (Lo > Hi) {
+        return fn.template operator()<0>();
+    } else {
+        if (width == Lo) {
+            return fn.template operator()<Lo>();
+        }
+        return dispatch_width<Lo + Step, Hi, Step>(
+                width, std::forward<Lambda>(fn));
+    }
+}
+} // namespace detail
+
+/// Specialize for common float level widths (multiples of 8 up to 128).
+template <typename LambdaType>
+inline auto with_level_width(size_t width, LambdaType&& action) {
+    return detail::dispatch_width<8, 128, 8>(
+            width, std::forward<LambdaType>(action));
+}
+
+template <typename Lambda>
+inline auto with_bool(bool value, Lambda&& fn) {
+    if (value) {
+        return fn.template operator()<true>();
+    } else {
+        return fn.template operator()<false>();
+    }
+}
+#endif // SWIG
+
 /**
  * Implements the core logic of Panorama-based refinement.
  * arXiv: https://arxiv.org/abs/2510.00566
@@ -42,6 +233,8 @@ namespace faiss {
  * accelerating the refinement stage.
  */
 struct Panorama {
+    static constexpr size_t kDefaultBatchSize = 128;
+
     size_t d = 0;
     size_t code_size = 0;
     size_t n_levels = 0;
@@ -98,6 +291,7 @@ struct Panorama {
     /// 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.
+#ifndef SWIG
     template <typename C, MetricType M>
     size_t progressive_filter_batch(
             const uint8_t* codes_base,
@@ -110,111 +304,99 @@ struct Panorama {
             const idx_t* ids,
             bool use_sel,
             std::vector<uint32_t>& active_indices,
+            std::vector<uint8_t>& active_byteset,
             std::vector<float>& exact_distances,
+            std::vector<float>& dot_buffer,
             float threshold,
-            PanoramaStats& local_stats) const;
-
-    void reconstruct(idx_t key, float* recons, const uint8_t* codes_base) const;
-};
+            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);
 
-template <typename C, MetricType M>
-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];
+        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];
 
-        if constexpr (M == METRIC_INNER_PRODUCT) {
-            exact_distances[i] = 0.0f;
-        } else {
-            exact_distances[i] = cum_sum * cum_sum + q_norm;
-        }
+        size_t batch_offset = batch_no * batch_size * code_size;
+        const uint8_t* storage_base = codes_base + batch_offset;
 
-        num_active += include;
-    }
+        // 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);
 
-    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;
+            active_indices[num_active] = i;
+            float cum_sum = batch_cum_sums[i];
 
-        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);
+            if constexpr (M == METRIC_INNER_PRODUCT) {
+                exact_distances[i] = 0.0f;
+            } else {
+                exact_distances[i] = cum_sum * cum_sum + q_norm;
+            }
 
-        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);
+            num_active += include;
+        }
 
-            const float* yj = level_storage + idx * actual_level_width;
-            const float* query_level = query + level * level_width_floats;
+        size_t total_active = num_active;
+        const bool first_level_full = (num_active == curr_batch_size);
 
-            float dot_product =
-                    fvec_inner_product(query_level, yj, actual_level_width);
+        local_stats.total_dims += total_active * n_levels;
 
-            if constexpr (M == METRIC_INNER_PRODUCT) {
-                exact_distances[idx] += dot_product;
-            } else {
-                exact_distances[idx] -= 2.0f * dot_product;
-            }
+        for (size_t level = 0; (level < n_levels) && (num_active > 0);
+             level++) {
+            local_stats.total_dims_scanned += num_active;
 
-            float cum_sum = level_cum_sums[idx];
-            float cauchy_schwarz_bound;
-            if constexpr (M == METRIC_INNER_PRODUCT) {
-                cauchy_schwarz_bound = -cum_sum * query_cum_norm;
-            } else {
-                cauchy_schwarz_bound = 2.0f * cum_sum * query_cum_norm;
-            }
+            float query_cum_norm = query_cum_sums[level + 1];
 
-            float lower_bound = exact_distances[idx] - cauchy_schwarz_bound;
+            size_t level_offset = level * level_width * batch_size;
+            const float* level_storage =
+                    (const float*)(storage_base + level_offset);
+            const float* query_level = query + level * level_width_floats;
+            size_t actual_level_width = std::min(
+                    level_width_floats, d - level * level_width_floats);
 
-            active_indices[next_active] = idx;
-            next_active += C::cmp(threshold, lower_bound) ? 1 : 0;
+            num_active = with_bool(
+                    level == 0 && first_level_full, [&]<bool AllActive>() {
+                        with_level_width(
+                                actual_level_width, [&]<size_t LevelWidth>() {
+                                    compute_level_dot_kernel<
+                                            AllActive,
+                                            LevelWidth>(
+                                            query_level,
+                                            level_storage,
+                                            active_indices.data(),
+                                            num_active,
+                                            actual_level_width,
+                                            dot_buffer.data());
+                                });
+
+                        prune_kernel<AllActive, C, M>(
+                                exact_distances.data(),
+                                dot_buffer.data(),
+                                level_cum_sums,
+                                active_byteset.data(),
+                                active_indices.data(),
+                                (uint32_t)num_active,
+                                query_cum_norm,
+                                threshold);
+
+                        return compact_active_kernel(
+                                active_indices.data(),
+                                active_byteset.data(),
+                                num_active);
+                    });
+
+            level_cum_sums += batch_size;
         }
 
-        num_active = next_active;
-        level_cum_sums += batch_size;
+        return num_active;
     }
+#endif // SWIG
 
-    return num_active;
-}
+    void reconstruct(idx_t key, float* recons, const uint8_t* codes_base) const;
+};
 } // namespace faiss
 
 #endif

data/vendor/faiss/faiss/impl/PdxLayout.cpp

@@ -0,0 +1,93 @@
+/*
+ * 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/PdxLayout.h>
+
+#include <cstddef>
+#include <cstring>
+
+namespace faiss {
+namespace detail {
+
+void pdxify(
+        const float* Y,
+        int k,
+        int d_trail,
+        int pdx_block_size,
+        float* Y_pdx) {
+    const int n_full_blocks = d_trail / pdx_block_size;
+    const int tail = d_trail % pdx_block_size;
+    size_t offset = 0;
+    for (int b = 0; b < n_full_blocks; ++b) {
+        const size_t src_start = static_cast<size_t>(b) * pdx_block_size;
+        for (int j = 0; j < k; ++j) {
+            std::memcpy(
+                    Y_pdx + offset,
+                    Y + static_cast<size_t>(j) * d_trail + src_start,
+                    pdx_block_size * sizeof(float));
+            offset += pdx_block_size;
+        }
+    }
+    if (tail > 0) {
+        const size_t src_start =
+                static_cast<size_t>(n_full_blocks) * pdx_block_size;
+        for (int j = 0; j < k; ++j) {
+            std::memcpy(
+                    Y_pdx + offset,
+                    Y + static_cast<size_t>(j) * d_trail + src_start,
+                    tail * sizeof(float));
+            offset += tail;
+        }
+    }
+}
+
+void de_pdxify(
+        const float* Y_pdx,
+        int k,
+        int d_trail,
+        int pdx_block_size,
+        float* Y) {
+    const int n_full_blocks = d_trail / pdx_block_size;
+    const int tail = d_trail % pdx_block_size;
+    size_t offset = 0;
+    for (int b = 0; b < n_full_blocks; ++b) {
+        const size_t dst_start = static_cast<size_t>(b) * pdx_block_size;
+        for (int j = 0; j < k; ++j) {
+            std::memcpy(
+                    Y + static_cast<size_t>(j) * d_trail + dst_start,
+                    Y_pdx + offset,
+                    pdx_block_size * sizeof(float));
+            offset += pdx_block_size;
+        }
+    }
+    if (tail > 0) {
+        const size_t dst_start =
+                static_cast<size_t>(n_full_blocks) * pdx_block_size;
+        for (int j = 0; j < k; ++j) {
+            std::memcpy(
+                    Y + static_cast<size_t>(j) * d_trail + dst_start,
+                    Y_pdx + offset,
+                    tail * sizeof(float));
+            offset += tail;
+        }
+    }
+}
+
+void compute_partial_norms(const float* X, int n, int d, int p, float* norms) {
+#pragma omp parallel for
+    for (int i = 0; i < n; ++i) {
+        float s = 0.0f;
+        const float* row = X + static_cast<size_t>(i) * d;
+        for (int m = 0; m < p; ++m) {
+            s += row[m] * row[m];
+        }
+        norms[i] = s;
+    }
+}
+
+} // namespace detail
+} // namespace faiss

data/vendor/faiss/faiss/impl/PdxLayout.h

@@ -0,0 +1,41 @@
+/*
+ * 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
+
+namespace faiss {
+namespace detail {
+
+/** Reorder a row-major (k, d_trail) matrix into PDX block-column-major
+ * layout. Inside each block of `pdx_block_size` dims the layout is
+ * column-major across centroids, so all k centroids' values for the same
+ * dim are contiguous — the access pattern that makes progressive pruning
+ * cache-friendly. Trailing block (size `d_trail % pdx_block_size`) uses
+ * the same convention. `Y_pdx` must already be sized to `k * d_trail`. */
+void pdxify(
+        const float* Y,
+        int k,
+        int d_trail,
+        int pdx_block_size,
+        float* Y_pdx);
+
+/** Inverse of pdxify (used in tests for the bit-identical round-trip
+ * check). */
+void de_pdxify(
+        const float* Y_pdx,
+        int k,
+        int d_trail,
+        int pdx_block_size,
+        float* Y);
+
+/** norms[i] = sum_{m<p} X[i, m]^2  for row-major X of shape (n, d).
+ * Parallel over rows. Used by SuperKMeans to keep partial-norm caches
+ * in sync with the current d_prime. */
+void compute_partial_norms(const float* X, int n, int d, int p, float* norms);
+
+} // namespace detail
+} // namespace faiss