faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/RaBitQUtils.h

@@ -9,8 +9,13 @@
 
 #include <faiss/MetricType.h>
 #include <faiss/impl/platform_macros.h>
+#include <faiss/utils/AlignedTable.h>
+#include <faiss/utils/rabitq_simd.h>
+#include <faiss/utils/simd_levels.h>
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 #include <vector>
 
 namespace faiss {
@@ -68,6 +73,7 @@ struct QueryFactorsData {
 
     float qr_to_c_L2sqr = 0;
     float qr_norm_L2sqr = 0;
+    float q_dot_c = 0; // <query, centroid> for IP metric; 0 for L2
 
     float int_dot_scale = 1;
 
@@ -239,6 +245,41 @@ inline float compute_1bit_adjusted_distance(
     return adjusted_distance;
 }
 
+/** Determine whether a candidate should be refined in two-stage search.
+ * Consolidates the filtering logic for both L2 and IP metrics.
+ *
+ * For L2 (min-heap): uses lower_bound = est_distance - error_adjustment
+ *   - Skip if lower_bound >= threshold (can't beat current worst)
+ * For IP (max-heap): uses upper_bound = est_distance + error_adjustment
+ *   - Skip if upper_bound <= threshold (can't beat current best)
+ *
+ * @param est_distance     Estimated 1-bit distance
+ * @param f_error          Database vector error factor
+ * @param g_error          Query vector error factor
+ * @param threshold        Current heap threshold (worst result in heap)
+ * @param is_similarity    True for IP metric (max-heap), false for L2
+ * (min-heap)
+ * @return                 True if candidate should be refined with full
+ * multi-bit distance
+ */
+inline bool should_refine_candidate(
+        float est_distance,
+        float f_error,
+        float g_error,
+        float threshold,
+        bool is_similarity) {
+    float error_adjustment = f_error * g_error;
+    if (is_similarity) {
+        // IP (max-heap): use upper bound for filtering
+        float upper_bound = est_distance + error_adjustment;
+        return upper_bound > threshold;
+    } else {
+        // L2 (min-heap): use lower bound for filtering
+        float lower_bound = std::max(0.0f, est_distance - error_adjustment);
+        return lower_bound < threshold;
+    }
+}
+
 /** Extract multi-bit code on-the-fly from packed ex-bit codes.
  * This inline function extracts a single code value without unpacking the
  * entire array, enabling efficient on-the-fly decoding during distance
@@ -276,55 +317,141 @@ inline int extract_code_inline(
  *
  * The multi-bit distance combines the sign bit (1-bit) with additional
  * magnitude bits (ex_bits) to compute a more accurate distance estimate.
+ * Uses SIMD-optimized bit-plane decomposition (AVX2+BMI2) for ex_bits 1-7,
+ * with scalar fallback for non-x86 or non-BMI2 platforms.
  *
  * @param sign_bits       unpacked sign bits (1-bit codes in standard format)
  * @param ex_code         packed ex-bit codes
  * @param ex_fac          ex-bit factors (f_add_ex, f_rescale_ex)
  * @param rotated_q       rotated query vector
- * @param qr_to_c_L2sqr   precomputed ||query_rotated - centroid||^2
- * @param qr_norm_L2sqr   precomputed ||query_rotated||^2 (0 for L2 metric)
+ * @param qr_base         precomputed base term: ||q-c||^2 for L2, <q,c> for IP
  * @param d               dimensionality
  * @param ex_bits         number of extra bits (nb_bits - 1)
  * @param metric_type     distance metric (L2 or Inner Product)
  * @return                computed full multi-bit distance
  */
+float compute_full_multibit_distance(
+        const uint8_t* sign_bits,
+        const uint8_t* ex_code,
+        const ExtraBitsFactors& ex_fac,
+        const float* rotated_q,
+        float qr_base,
+        size_t d,
+        size_t ex_bits,
+        MetricType metric_type);
+
+// SIMDLevel-templatized version — avoids per-call dynamic dispatch.
+// Inline so it can be used from templatized distance computers without
+// needing explicit instantiations in per-SIMD TUs.
+template <SIMDLevel SL>
 inline float compute_full_multibit_distance(
         const uint8_t* sign_bits,
         const uint8_t* ex_code,
         const ExtraBitsFactors& ex_fac,
         const float* rotated_q,
-        float qr_to_c_L2sqr,
-        float qr_norm_L2sqr,
+        float qr_base,
         size_t d,
         size_t ex_bits,
         MetricType metric_type) {
-    float ex_ip = 0.0f;
     const float cb = -(static_cast<float>(1 << ex_bits) - 0.5f);
 
-    for (size_t i = 0; i < d; i++) {
-        const size_t byte_idx = i / 8;
-        const size_t bit_offset = i % 8;
-        const bool sign_bit = (sign_bits[byte_idx] >> bit_offset) & 1;
+    float ex_ip = rabitq::multibit::compute_inner_product<SL>(
+            sign_bits, ex_code, rotated_q, d, ex_bits, cb);
 
-        int ex_code_val = extract_code_inline(ex_code, i, ex_bits);
+    float dist = qr_base + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;
 
-        int total_code = (sign_bit ? 1 : 0) << ex_bits;
-        total_code += ex_code_val;
-        float reconstructed = static_cast<float>(total_code) + cb;
-
-        ex_ip += rotated_q[i] * reconstructed;
+    if (metric_type == MetricType::METRIC_L2) {
+        dist = std::max(0.0f, dist);
     }
 
-    float dist = qr_to_c_L2sqr + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;
+    return dist;
+}
 
-    if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
-        dist = -0.5f * (dist - qr_norm_L2sqr);
-    } else {
-        dist = std::max(0.0f, dist);
+/** Compute pointer to a vector's auxiliary data within block layout. */
+template <typename T>
+inline T* get_block_aux_ptr(
+        T* block_data,
+        size_t vec_pos,
+        size_t bbs,
+        size_t packed_block_size,
+        size_t full_block_size,
+        size_t storage_size) {
+    return block_data + (vec_pos / bbs) * full_block_size + packed_block_size +
+            (vec_pos % bbs) * storage_size;
+}
+
+/// Extract sign bits from PQ4-interleaved block into flat byte packing.
+/// Like CodePackerRaBitQ::unpack_1 but sign-bits-only and with the
+/// vector's in-block address hoisted out of the per-SQ loop.
+inline void unpack_sign_bits_from_packed(
+        const uint8_t* block,
+        size_t bbs,
+        size_t nsq,
+        size_t offset,
+        size_t block_stride,
+        uint8_t* sign_bits_out) {
+    block += (offset / bbs) * block_stride;
+    offset = offset % bbs;
+
+    const bool nibble_high = offset > 15;
+    const size_t vid = offset & 15;
+    const size_t in_group_addr =
+            (vid < 8) ? (vid << 1) : (((vid - 8) << 1) + 1);
+
+    const size_t num_pairs = nsq / 2;
+    for (size_t k = 0; k < num_pairs; k++) {
+        const size_t base = k * bbs;
+        const uint8_t raw_even = block[base + in_group_addr];
+        const uint8_t raw_odd = block[base + in_group_addr + 16];
+
+        const uint8_t nib0 = nibble_high ? (raw_even >> 4) : (raw_even & 0xF);
+        const uint8_t nib1 = nibble_high ? (raw_odd >> 4) : (raw_odd & 0xF);
+        sign_bits_out[k] = nib0 | (nib1 << 4);
     }
 
-    return dist;
+    if (nsq & 1) {
+        const uint8_t raw = block[num_pairs * bbs + in_group_addr];
+        sign_bits_out[num_pairs] = nibble_high ? (raw >> 4) : (raw & 0xF);
+    }
 }
 
+/** Compute per-vector auxiliary storage size.
+ *
+ * @param nb_bits  number of quantization bits (1 = sign-bit only)
+ * @param d        dimensionality
+ * @return         storage size in bytes
+ */
+size_t compute_per_vector_storage_size(size_t nb_bits, size_t d);
+
+/** [LEGACY FORMAT SUPPORT] Migrate block data from old I/O format to new
+ * format.
+ *
+ * This function is used only when reading indexes saved with the legacy format
+ * (fourcc "Irfs"/"Iwrf") to convert them to the new embedded auxiliary data
+ * format. Not needed for indexes saved with the new format ("Irfn"/"Iwrn").
+ *
+ * Re-layouts blocks in-place and copies aux data from flat_storage.
+ *
+ * @param flat_storage       legacy per-vector aux data indexed by global ID
+ * @param codes              block data (will be resized and re-laid out)
+ * @param num_vectors        number of vectors in this segment
+ * @param bbs                block batch size (vectors per block)
+ * @param M2                 rounded sub-quantizer count
+ * @param old_block_stride   old block size (packed codes only, or current)
+ * @param new_block_stride   new block size (packed codes + aux region)
+ * @param storage_size       per-vector aux storage size in bytes
+ * @param id_map             maps local offset to global ID; null = sequential
+ */
+void populate_block_aux_from_flat_storage(
+        const std::vector<uint8_t>& flat_storage,
+        AlignedTable<uint8_t>& codes,
+        size_t num_vectors,
+        size_t bbs,
+        size_t M2,
+        size_t old_block_stride,
+        size_t new_block_stride,
+        size_t storage_size,
+        const int64_t* id_map = nullptr);
+
 } // namespace rabitq_utils
 } // namespace faiss