faiss 0.6.2 → 0.6.4

data/vendor/faiss/faiss/impl/ProductQuantizer.cpp

@@ -280,7 +280,7 @@ void compute_1_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
 } // namespace
 
 void ProductQuantizer::compute_code(const float* x, uint8_t* code) const {
-    with_simd_level([&]<SIMDLevel SL>() {
+    with_simd_level_with_sve([&]<SIMDLevel SL>() {
         switch (nbits) {
             case 8:
                 compute_1_code<PQEncoder8, SL>(*this, x, code);
@@ -294,7 +294,7 @@ void ProductQuantizer::compute_code(const float* x, uint8_t* code) const {
                 compute_1_code<PQEncoderGeneric, SL>(*this, x, code);
                 break;
         }
-    }); // with_simd_level
+    }); // with_simd_level_with_sve
 }
 
 template <class PQDecoder>
@@ -442,7 +442,7 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
 
 void ProductQuantizer::compute_distance_table(const float* x, float* dis_table)
         const {
-    with_simd_level([&]<SIMDLevel SL>() {
+    with_simd_level_with_sve([&]<SIMDLevel SL>() {
         if (transposed_centroids.empty()) {
             // use regular version
             for (size_t m = 0; m < M; m++) {
@@ -824,7 +824,7 @@ void ProductQuantizer::compute_sdc_table() {
     sdc_table.resize(M * ksub * ksub);
 
     if (dsub < 4) {
-        with_simd_level([&]<SIMDLevel SL>() {
+        with_simd_level_with_sve([&]<SIMDLevel SL>() {
 #pragma omp parallel for
             for (int64_t mk = 0; mk < static_cast<int64_t>(M * ksub); mk++) {
                 // allow omp to schedule in a more fine-grained way
@@ -862,32 +862,44 @@ void ProductQuantizer::search_sdc(
     size_t k = res->k;
     int64_t nq_signed = nq;
 
-#pragma omp parallel for
-    for (int64_t i = 0; i < nq_signed; i++) {
-        /* Compute distances and keep smallest values */
-        idx_t* heap_ids = res->ids + i * k;
-        float* heap_dis = res->val + i * k;
-        const uint8_t* qcode = qcodes + i * code_size;
-
-        if (init_finalize_heap)
-            maxheap_heapify(k, heap_dis, heap_ids);
-
-        const uint8_t* bcode = bcodes;
-        for (size_t j = 0; j < nb; j++) {
-            float dis = 0;
-            const float* tab = sdc_table.data();
+#pragma omp parallel
+    {
+        // One allocation per OMP thread instead of one per query.
+        std::vector<const float*> q_row(M);
+#pragma omp for
+        for (int64_t i = 0; i < nq_signed; i++) {
+            idx_t* heap_ids = res->ids + i * k;
+            float* heap_dis = res->val + i * k;
+            const uint8_t* qcode = qcodes + i * code_size;
+
+            if (init_finalize_heap)
+                maxheap_heapify(k, heap_dis, heap_ids);
+
+            // Precompute per-subquantizer row pointers: q_row[m] points to
+            // sdc_table[m*ksub^2 + qcode[m]*ksub], eliminating M
+            // multiplications and M pointer advances per database vector in the
+            // j-loop.
+            const float* sdc = sdc_table.data();
             for (size_t m = 0; m < M; m++) {
-                dis += tab[bcode[m] + qcode[m] * ksub];
-                tab += ksub * ksub;
+                q_row[m] = sdc + m * (size_t)(ksub * ksub) +
+                        (size_t)qcode[m] * ksub;
             }
-            if (dis < heap_dis[0]) {
-                maxheap_replace_top(k, heap_dis, heap_ids, dis, j);
+
+            const uint8_t* bcode = bcodes;
+            for (size_t j = 0; j < nb; j++) {
+                float dis = 0;
+                for (size_t m = 0; m < M; m++) {
+                    dis += q_row[m][bcode[m]];
+                }
+                if (dis < heap_dis[0]) {
+                    maxheap_replace_top(k, heap_dis, heap_ids, dis, j);
+                }
+                bcode += code_size;
             }
-            bcode += code_size;
-        }
 
-        if (init_finalize_heap)
-            maxheap_reorder(k, heap_dis, heap_ids);
+            if (init_finalize_heap)
+                maxheap_reorder(k, heap_dis, heap_ids);
+        }
     }
 }
 

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

@@ -188,46 +188,54 @@ QueryFactorsData compute_query_factors(
 
     const float inv_d_sqrt = 1.0f / std::sqrt(static_cast<float>(d));
 
-    // Compute quantization range
-    float v_min = std::numeric_limits<float>::max();
-    float v_max = std::numeric_limits<float>::lowest();
-
     const float* rq = rotated_q.data();
-    if (centered) {
-        float z_max = Z_MAX_BY_QB[qb - 1];
-        float v_radius = z_max * std::sqrt(query_factors.qr_to_c_L2sqr / d);
-        v_min = -v_radius;
-        v_max = v_radius;
-    } else {
-        for (size_t i = 0; i < d; i++) {
-            const float v_q = rq[i];
-            v_min = std::min(v_min, v_q);
-            v_max = std::max(v_max, v_q);
-        }
-    }
-
-    // Quantize the query
-    const uint8_t max_code = (1 << qb) - 1;
-    const float delta = (v_max - v_min) / max_code;
-    const float inv_delta = 1.0f / delta;
-
-    rotated_qq.resize(d);
+    float v_min;
+    float v_max;
+    float delta;
     size_t sum_qq = 0;
     int64_t sum2_signed_odd_int = 0;
-
+    const uint8_t max_code = (1 << qb) - 1;
+    rotated_qq.resize(d);
     uint8_t* rqq = rotated_qq.data();
-    for (size_t i = 0; i < d; i++) {
-        const float v_q = rq[i];
-        const uint8_t v_qq = std::clamp<float>(
-                std::round((v_q - v_min) * inv_delta), 0, max_code);
-        rqq[i] = v_qq;
-        sum_qq += v_qq;
-
-        if (centered) {
-            int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
-            sum2_signed_odd_int += signed_odd_int * signed_odd_int;
-        }
-    }
+
+    // Select the SIMD implementation once for both range computation and
+    // quantization. This function runs once per query/probe pair.
+    with_selected_simd_levels<rabitq::RABITQ_QUANTIZATION_SIMD_LEVELS>(
+            [&]<SIMDLevel SL>() {
+                if (centered) {
+                    const float z_max = Z_MAX_BY_QB[qb - 1];
+                    const float v_radius =
+                            z_max * std::sqrt(query_factors.qr_to_c_L2sqr / d);
+                    v_min = -v_radius;
+                    v_max = v_radius;
+                } else {
+                    v_min = std::numeric_limits<float>::max();
+                    v_max = std::numeric_limits<float>::lowest();
+                    rabitq::minmax_values<SL>(rq, d, v_min, v_max);
+                }
+
+                delta = (v_max - v_min) / max_code;
+                // A constant (or zero-norm) query has delta == 0. Preserve the
+                // scalar path's centered correction terms while avoiding
+                // 0 * inf during quantization.
+                if (delta <= 0.0f) {
+                    memset(rqq, 0, d * sizeof(uint8_t));
+                    if (centered) {
+                        sum2_signed_odd_int = int64_t(d) * max_code * max_code;
+                    }
+                } else {
+                    rabitq::quantize_query_values<SL>(
+                            rq,
+                            d,
+                            v_min,
+                            1.0f / delta,
+                            max_code,
+                            centered,
+                            rqq,
+                            sum_qq,
+                            sum2_signed_odd_int);
+                }
+            });
 
     // Compute query factors
     query_factors.c1 = 2 * delta * inv_d_sqrt;
@@ -321,7 +329,7 @@ float compute_full_multibit_distance(
         size_t d,
         size_t ex_bits,
         MetricType metric_type) {
-    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0_SPR>(
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_BASE_WITH_VPOPCNT>(
             [&]<SIMDLevel SL>() {
                 return compute_full_multibit_distance<SL>(
                         sign_bits,

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

@@ -13,6 +13,7 @@
 #include <faiss/utils/rabitq_simd.h>
 #include <faiss/utils/simd_levels.h>
 #include <algorithm>
+#include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
@@ -86,6 +87,40 @@ struct QueryFactorsData {
  */
 FAISS_API extern const float Z_MAX_BY_QB[8];
 
+/** Fast half-away rounding for non-negative RaBitQ quantization values.
+ *
+ * Contract: x must be finite, non-NaN, and in [0, 255.5). This is not a
+ * bit-exact replacement for roundf within a few ulps below k + 0.5, where the
+ * addition may round across the tie. It is intended only for RaBitQ LUT
+ * quantization paths where a +/- 1 code difference at those boundaries is
+ * acceptable.
+ */
+inline uint8_t round_nonnegative_to_uint8(float x) {
+    assert(x == x);
+    assert(x >= 0.0f);
+    assert(x < 255.5f);
+    return rabitq::round_nonnegative_byte_scalar(x);
+}
+
+/** Same as round_nonnegative_to_uint8 for uint16 RaBitQ bias values. */
+inline uint16_t round_nonnegative_to_uint16(float x) {
+    assert(x == x);
+    assert(x >= 0.0f);
+    assert(x < 65535.5f);
+    return static_cast<uint16_t>(static_cast<int>(x + 0.5f));
+}
+
+/** Fast clamped rounding for query-byte quantization.
+ *
+ * Contract: x must be non-NaN. Values outside [0, max_code] are clamped before
+ * applying the same non-bit-exact rounding used by
+ * round_nonnegative_to_uint8().
+ */
+inline uint8_t round_clamped_to_uint8(float x, uint8_t max_code) {
+    assert(x == x);
+    return rabitq::round_clamped_byte_scalar(x, max_code);
+}
+
 /** Compute factors for a single database vector using RaBitQ algorithm.
  * This function consolidates the mathematical logic that was duplicated
  * between IndexRaBitQ and IndexRaBitQFastScan.