faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/IndexRaBitQFastScan.cpp

@@ -7,10 +7,10 @@
 
 #include <faiss/IndexRaBitQFastScan.h>
 #include <faiss/impl/CodePackerRaBitQ.h>
-#include <faiss/impl/FastScanDistancePostProcessing.h>
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/RaBitQuantizerMultiBit.h>
-#include <faiss/impl/pq4_fast_scan.h>
+#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
+#include <faiss/impl/fast_scan/fast_scan.h>
 #include <faiss/utils/utils.h>
 #include <algorithm>
 #include <cmath>
@@ -28,13 +28,13 @@ size_t IndexRaBitQFastScan::compute_per_vector_storage_size() const {
 IndexRaBitQFastScan::IndexRaBitQFastScan() = default;
 
 IndexRaBitQFastScan::IndexRaBitQFastScan(
-        idx_t d,
+        idx_t d_in,
         MetricType metric,
-        int bbs,
+        int bbs_in,
         uint8_t nb_bits)
-        : rabitq(d, metric, nb_bits) {
+        : rabitq(d_in, metric, nb_bits) {
     // RaBitQ-specific validation
-    FAISS_THROW_IF_NOT_MSG(d > 0, "Dimension must be positive");
+    FAISS_THROW_IF_NOT_MSG(d_in > 0, "Dimension must be positive");
     FAISS_THROW_IF_NOT_MSG(
             metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT,
             "RaBitQ FastScan only supports L2 and Inner Product metrics");
@@ -43,18 +43,19 @@ IndexRaBitQFastScan::IndexRaBitQFastScan(
 
     // RaBitQ uses 1 bit per dimension packed into 4-bit FastScan sub-quantizers
     // Each FastScan sub-quantizer handles 4 RaBitQ dimensions
-    const size_t M_fastscan = (d + 3) / 4;
+    const size_t M_fastscan = (d_in + 3) / 4;
     constexpr size_t nbits_fastscan = 4;
 
     // init_fastscan will validate bbs % 32 == 0 and nbits_fastscan == 4
-    init_fastscan(static_cast<int>(d), M_fastscan, nbits_fastscan, metric, bbs);
+    init_fastscan(
+            static_cast<int>(d_in), M_fastscan, nbits_fastscan, metric, bbs_in);
 
     // Compute code_size directly using RaBitQuantizer
-    code_size = rabitq.compute_code_size(d, nb_bits);
+    code_size = rabitq.compute_code_size(d_in, nb_bits);
 
     // Set RaBitQ-specific parameters
     qb = 8;
-    center.resize(d, 0.0f);
+    center.resize(d_in, 0.0f);
 }
 
 CodePacker* IndexRaBitQFastScan::get_CodePacker() const {
@@ -102,7 +103,7 @@ size_t IndexRaBitQFastScan::remove_ids(const IDSelector& sel) {
     return nremove;
 }
 
-IndexRaBitQFastScan::IndexRaBitQFastScan(const IndexRaBitQ& orig, int bbs)
+IndexRaBitQFastScan::IndexRaBitQFastScan(const IndexRaBitQ& orig, int bbs_in)
         : rabitq(orig.rabitq) {
     // RaBitQ-specific validation
     FAISS_THROW_IF_NOT_MSG(orig.d > 0, "Dimension must be positive");
@@ -122,7 +123,7 @@ IndexRaBitQFastScan::IndexRaBitQFastScan(const IndexRaBitQ& orig, int bbs)
             M_fastscan,
             nbits_fastscan,
             orig.metric_type,
-            bbs);
+            bbs_in);
 
     code_size = rabitq.compute_code_size(d, rabitq.nb_bits);
 
@@ -148,7 +149,7 @@ IndexRaBitQFastScan::IndexRaBitQFastScan(const IndexRaBitQ& orig, int bbs)
             const uint8_t* orig_code = orig.codes.data() + i * orig.code_size;
             uint8_t* fs_code = fastscan_codes.get() + i * code_size;
 
-            for (size_t j = 0; j < orig.d; j++) {
+            for (size_t j = 0; j < static_cast<size_t>(orig.d); j++) {
                 const size_t orig_byte_idx = j / 8;
                 const size_t orig_bit_offset = j % 8;
                 const bool bit_value =
@@ -197,13 +198,13 @@ void IndexRaBitQFastScan::train(idx_t n, const float* x) {
     // compute a centroid
     std::vector<float> centroid(d, 0);
     for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
-        for (size_t j = 0; j < d; j++) {
+        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 < d; j++) {
+        for (size_t j = 0; j < static_cast<size_t>(d); j++) {
             centroid[j] /= (float)n;
         }
     }
@@ -279,9 +280,11 @@ void IndexRaBitQFastScan::add(idx_t n, const float* x) {
     ntotal += n;
 }
 
-void IndexRaBitQFastScan::compute_codes(uint8_t* codes, idx_t n, const float* x)
-        const {
-    FAISS_ASSERT(codes != nullptr);
+void IndexRaBitQFastScan::compute_codes(
+        uint8_t* out_codes,
+        idx_t n,
+        const float* x) const {
+    FAISS_ASSERT(out_codes != nullptr);
     FAISS_ASSERT(x != nullptr);
     FAISS_ASSERT(
             (metric_type == MetricType::METRIC_L2 ||
@@ -296,23 +299,23 @@ void IndexRaBitQFastScan::compute_codes(uint8_t* codes, idx_t n, const float* x)
     const size_t ex_bits = rabitq.nb_bits - 1;
     const size_t ex_code_size = (d * ex_bits + 7) / 8;
 
-    memset(codes, 0, n * code_size);
+    memset(out_codes, 0, n * code_size);
 
 #pragma omp parallel for if (n > 1000)
     for (int64_t i = 0; i < n; i++) {
-        uint8_t* const code = codes + i * code_size;
+        uint8_t* const code = out_codes + i * code_size;
         const float* const x_row = x + i * d;
 
         // Compute residual once, reuse for both sign bits and ex-bits
         std::vector<float> residual(d);
-        for (size_t j = 0; j < d; j++) {
+        for (size_t j = 0; j < static_cast<size_t>(d); j++) {
             const float centroid_val = centroid_data ? centroid_data[j] : 0.0f;
             residual[j] = x_row[j] - centroid_val;
         }
 
         // Pack sign bits directly into FastScan format using precomputed
         // residual
-        for (size_t j = 0; j < d; j++) {
+        for (size_t j = 0; j < static_cast<size_t>(d); j++) {
             if (residual[j] > 0.0f) {
                 rabitq_utils::set_bit_fastscan(code, j);
             }
@@ -412,7 +415,7 @@ void IndexRaBitQFastScan::compute_float_LUT(
                     for (size_t dim_offset = 0; dim_offset < 4; dim_offset++) {
                         const size_t dim_idx = dim_start + dim_offset;
 
-                        if (dim_idx < d) {
+                        if (dim_idx < static_cast<size_t>(d)) {
                             const bool db_bit = (code_val >> dim_offset) & 1;
                             const float query_value = rotated_qq[dim_idx];
 
@@ -447,7 +450,8 @@ void IndexRaBitQFastScan::compute_float_LUT(
                     for (size_t dim_offset = 0; dim_offset < 4; dim_offset++) {
                         const size_t dim_idx = dim_start + dim_offset;
 
-                        if (dim_idx < d && ((code_val >> dim_offset) & 1)) {
+                        if (dim_idx < static_cast<size_t>(d) &&
+                            ((code_val >> dim_offset) & 1)) {
                             inner_product += rotated_qq[dim_idx];
                             popcount++;
                         }
@@ -463,12 +467,16 @@ void IndexRaBitQFastScan::compute_float_LUT(
     }
 }
 
+size_t IndexRaBitQFastScan::fast_scan_code_size() const {
+    return (d + 7) / 8;
+}
+
 void IndexRaBitQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x)
         const {
     const float* centroid_in =
             (center.data() == nullptr) ? nullptr : center.data();
-    const uint8_t* codes = bytes;
-    FAISS_ASSERT(codes != nullptr);
+    const uint8_t* input_codes = bytes;
+    FAISS_ASSERT(input_codes != nullptr);
     FAISS_ASSERT(x != nullptr);
 
     const float inv_d_sqrt = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d));
@@ -477,7 +485,7 @@ void IndexRaBitQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x)
 #pragma omp parallel for if (n > 1000)
     for (int64_t i = 0; i < n; i++) {
         // Access code using correct FastScan format
-        const uint8_t* code = codes + i * code_size;
+        const uint8_t* code = input_codes + i * code_size;
 
         // Extract factors directly from embedded codes
         const uint8_t* factors_ptr = code + bit_pattern_size;
@@ -485,7 +493,7 @@ void IndexRaBitQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x)
                 reinterpret_cast<const rabitq_utils::SignBitFactors*>(
                         factors_ptr);
 
-        for (size_t j = 0; j < d; j++) {
+        for (size_t j = 0; j < static_cast<size_t>(d); j++) {
             // Use RaBitQUtils for consistent bit extraction
             bool bit_value = rabitq_utils::extract_bit_fastscan(code, j);
             float bit = bit_value ? 1.0f : 0.0f;
@@ -522,264 +530,20 @@ void IndexRaBitQFastScan::search(
     }
 }
 
-// Template implementations for RaBitQHeapHandler
-template <class C, bool with_id_map>
-RaBitQHeapHandler<C, with_id_map>::RaBitQHeapHandler(
-        const IndexRaBitQFastScan* index,
-        size_t nq_val,
-        size_t k_val,
-        float* distances,
-        int64_t* labels,
-        const IDSelector* sel_in,
-        const FastScanDistancePostProcessing& ctx,
-        bool multi_bit)
-        : RHC(nq_val, index->ntotal, sel_in),
-          rabitq_index(index),
-          heap_distances(distances),
-          heap_labels(labels),
-          nq(nq_val),
-          k(k_val),
-          context(ctx),
-          is_multi_bit(multi_bit),
-          storage_size(index->compute_per_vector_storage_size()),
-          packed_block_size(((index->M2 + 1) / 2) * index->bbs),
-          full_block_size(index->get_block_stride()),
-          packer(index->get_CodePacker()) {
-    // Initialize heaps for all queries in constructor
-    // This allows us to support direct normalizer assignment
-#pragma omp parallel for if (nq > 100)
-    for (int64_t q = 0; q < static_cast<int64_t>(nq); q++) {
-        float* heap_dis = heap_distances + q * k;
-        int64_t* heap_ids = heap_labels + q * k;
-        heap_heapify<Cfloat>(k, heap_dis, heap_ids);
-    }
-}
-
-template <class C, bool with_id_map>
-void RaBitQHeapHandler<C, with_id_map>::handle(
-        size_t q,
-        size_t b,
-        simd16uint16 d0,
-        simd16uint16 d1) {
-    ALIGNED(32) uint16_t d32tab[32];
-    d0.store(d32tab);
-    d1.store(d32tab + 16);
-
-    // Get heap pointers and query factors (computed once per batch)
-    float* const heap_dis = heap_distances + q * k;
-    int64_t* const heap_ids = heap_labels + q * k;
-
-    // Access query factors from query_factors pointer
-    rabitq_utils::QueryFactorsData query_factors_data = {};
-    if (context.query_factors != nullptr) {
-        query_factors_data = context.query_factors[q];
-    }
-
-    // Compute normalizers once per batch
-    const float one_a = normalizers ? (1.0f / normalizers[2 * q]) : 1.0f;
-    const float bias = normalizers ? normalizers[2 * q + 1] : 0.0f;
-
-    // Compute loop bounds to avoid redundant bounds checking
-    const size_t base_db_idx = this->j0 + b * 32;
-    const size_t max_vectors = (base_db_idx < rabitq_index->ntotal)
-            ? std::min<size_t>(32, rabitq_index->ntotal - base_db_idx)
-            : 0;
-
-    // Compute block auxiliary region base pointer once per batch.
-    // Since bbs=32, each batch of 32 vectors aligns to one block.
-    const size_t block_idx = base_db_idx / rabitq_index->bbs;
-    const uint8_t* aux_base = rabitq_index->codes.get() +
-            block_idx * full_block_size + packed_block_size;
-
-    // Stats tracking for multi-bit two-stage search only
-    // n_1bit_evaluations: candidates evaluated using 1-bit lower bound
-    // n_multibit_evaluations: candidates requiring full multi-bit distance
-    size_t local_1bit_evaluations = 0;
-    size_t local_multibit_evaluations = 0;
-
-    // Process distances in batch
-    for (size_t i = 0; i < max_vectors; i++) {
-        const size_t db_idx = base_db_idx + i;
-
-        // Normalize distance from LUT lookup
-        const float normalized_distance = d32tab[i] * one_a + bias;
-
-        // Access factors from block auxiliary region
-        const uint8_t* base_ptr = aux_base + i * storage_size;
-
-        if (is_multi_bit) {
-            // Track candidates actually considered for two-stage filtering
-            local_1bit_evaluations++;
-
-            const SignBitFactorsWithError& full_factors =
-                    *reinterpret_cast<const SignBitFactorsWithError*>(base_ptr);
-
-            float dist_1bit = rabitq_utils::compute_1bit_adjusted_distance(
-                    normalized_distance,
-                    full_factors,
-                    query_factors_data,
-                    rabitq_index->centered,
-                    rabitq_index->qb,
-                    rabitq_index->d);
-
-            // Adaptive filtering: decide whether to compute full distance
-            const bool is_similarity = rabitq_index->metric_type ==
-                    MetricType::METRIC_INNER_PRODUCT;
-            bool should_refine = rabitq_utils::should_refine_candidate(
-                    dist_1bit,
-                    full_factors.f_error,
-                    context.query_factors ? context.query_factors[q].g_error
-                                          : 0.0f,
-                    heap_dis[0],
-                    is_similarity);
-
-            if (should_refine) {
-                local_multibit_evaluations++;
-                float dist_full = compute_full_multibit_distance(db_idx, q);
-
-                if (Cfloat::cmp(heap_dis[0], dist_full)) {
-                    heap_replace_top<Cfloat>(
-                            k, heap_dis, heap_ids, dist_full, db_idx);
-                }
-            }
-        } else {
-            const rabitq_utils::SignBitFactors& db_factors =
-                    *reinterpret_cast<const rabitq_utils::SignBitFactors*>(
-                            base_ptr);
-
-            float adjusted_distance =
-                    rabitq_utils::compute_1bit_adjusted_distance(
-                            normalized_distance,
-                            db_factors,
-                            query_factors_data,
-                            rabitq_index->centered,
-                            rabitq_index->qb,
-                            rabitq_index->d);
-
-            // Add to heap if better than current worst
-            if (Cfloat::cmp(heap_dis[0], adjusted_distance)) {
-                heap_replace_top<Cfloat>(
-                        k, heap_dis, heap_ids, adjusted_distance, db_idx);
-            }
-        }
-    }
-
-    // Update global stats atomically
-#pragma omp atomic
-    rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
-#pragma omp atomic
-    rabitq_stats.n_multibit_evaluations += local_multibit_evaluations;
-}
-
-template <class C, bool with_id_map>
-void RaBitQHeapHandler<C, with_id_map>::begin(const float* norms) {
-    normalizers = norms;
-    // Heap initialization is now done in constructor
-}
-
-template <class C, bool with_id_map>
-void RaBitQHeapHandler<C, with_id_map>::end() {
-// Reorder final results
-#pragma omp parallel for if (nq > 100)
-    for (int64_t q = 0; q < static_cast<int64_t>(nq); q++) {
-        float* heap_dis = heap_distances + q * k;
-        int64_t* heap_ids = heap_labels + q * k;
-        heap_reorder<Cfloat>(k, heap_dis, heap_ids);
-    }
-}
-
-template <class C, bool with_id_map>
-float RaBitQHeapHandler<C, with_id_map>::compute_lower_bound(
-        float dist_1bit,
-        size_t db_idx,
-        size_t q) const {
-    // Access f_error from block auxiliary region
-    const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
-            rabitq_index->codes.get(),
-            db_idx,
-            rabitq_index->bbs,
-            packed_block_size,
-            full_block_size,
-            storage_size);
-    const SignBitFactorsWithError& db_factors =
-            *reinterpret_cast<const SignBitFactorsWithError*>(base_ptr);
-    float f_error = db_factors.f_error;
-
-    // Get g_error from query factors (query-dependent error term)
-    float g_error = 0.0f;
-    if (context.query_factors != nullptr) {
-        g_error = context.query_factors[q].g_error;
-    }
-
-    // Compute error adjustment: f_error * g_error
-    float error_adjustment = f_error * g_error;
-
-    return dist_1bit - error_adjustment;
-}
-
-template <class C, bool with_id_map>
-float RaBitQHeapHandler<C, with_id_map>::compute_full_multibit_distance(
-        size_t db_idx,
-        size_t q) const {
-    const size_t ex_bits = rabitq_index->rabitq.nb_bits - 1;
-    const size_t dim = rabitq_index->d;
-
-    const uint8_t* base_ptr = rabitq_utils::get_block_aux_ptr(
-            rabitq_index->codes.get(),
-            db_idx,
-            rabitq_index->bbs,
-            packed_block_size,
-            full_block_size,
-            storage_size);
-
-    const size_t ex_code_size = (dim * ex_bits + 7) / 8;
-    const uint8_t* ex_code = base_ptr + sizeof(SignBitFactorsWithError);
-    const ExtraBitsFactors& ex_fac = *reinterpret_cast<const ExtraBitsFactors*>(
-            base_ptr + sizeof(SignBitFactorsWithError) + ex_code_size);
-
-    // Get query factors reference (avoid copying)
-    const rabitq_utils::QueryFactorsData& query_factors =
-            context.query_factors[q];
-
-    // Get sign bits from FastScan packed format
-    std::vector<uint8_t> unpacked_code(rabitq_index->code_size);
-    packer->unpack_1(rabitq_index->codes.get(), db_idx, unpacked_code.data());
-    const uint8_t* sign_bits = unpacked_code.data();
-
-    return rabitq_utils::compute_full_multibit_distance(
-            sign_bits,
-            ex_code,
-            ex_fac,
-            query_factors.rotated_q.data(),
-            (rabitq_index->metric_type == MetricType::METRIC_INNER_PRODUCT)
-                    ? query_factors.q_dot_c
-                    : query_factors.qr_to_c_L2sqr,
-            dim,
-            ex_bits,
-            rabitq_index->metric_type);
-}
+std::unique_ptr<FastScanCodeScanner> IndexRaBitQFastScan::make_knn_scanner(
 
-// Implementation of virtual make_knn_handler method
-SIMDResultHandlerToFloat* IndexRaBitQFastScan::make_knn_handler(
         bool is_max,
-        int /*impl*/,
         idx_t n,
         idx_t k,
         size_t /*ntotal*/,
         float* distances,
         idx_t* labels,
         const IDSelector* sel,
+        int /*impl*/,
         const FastScanDistancePostProcessing& context) const {
-    // Use runtime boolean for multi-bit mode
-    const bool multi_bit = rabitq.nb_bits > 1;
-
-    if (is_max) {
-        return new RaBitQHeapHandler<CMax<uint16_t, int>, false>(
-                this, n, k, distances, labels, sel, context, multi_bit);
-    } else {
-        return new RaBitQHeapHandler<CMin<uint16_t, int>, false>(
-                this, n, k, distances, labels, sel, context, multi_bit);
-    }
+    const bool is_multi_bit = rabitq.nb_bits > 1;
+    return rabitq_make_knn_scanner(
+            this, is_max, n, k, distances, labels, sel, context, is_multi_bit);
 }
 
 } // namespace faiss