faiss 0.6.2 → 0.6.3

data/vendor/faiss/faiss/impl/PolysemousTraining.cpp

@@ -26,6 +26,8 @@
 
 #include <faiss/impl/FaissAssert.h>
 
+#include <faiss/impl/polysemous_training/dispatch.h>
+
 /*****************************************
  * Mixed PQ / Hamming
  ******************************************/
@@ -168,6 +170,94 @@ static inline int hamming_dis(uint64_t a, uint64_t b) {
     return popcount64(a ^ b);
 }
 
+static inline double sqr(double x) {
+    return x * x;
+}
+
+// Scalar (NONE) kernels for the two objectives. The dispatch boundary lives in
+// the objective methods below, which route to these or to the AVX-512
+// specializations (in polysemous_training/avx512.cpp) via
+// with_selected_simd_levels.
+namespace polysemous_training {
+
+template <>
+double hamming_compute_cost<SIMDLevel::NONE>(
+        int n,
+        const int* perm,
+        const double* target_dis,
+        const double* weights) {
+    double cost = 0;
+    for (int i = 0; i < n; i++) {
+        for (int j = 0; j < n; j++) {
+            double wanted = target_dis[i * n + j];
+            double w = weights[i * n + j];
+            double actual = hamming_dis(perm[i], perm[j]);
+            cost += w * sqr(wanted - actual);
+        }
+    }
+    return cost;
+}
+
+template <>
+double hamming_cost_update<SIMDLevel::NONE>(
+        int n,
+        const int* perm,
+        int iw,
+        int jw,
+        const double* target_dis,
+        const double* weights) {
+    double delta_cost = 0;
+
+    for (int i = 0; i < n; i++) {
+        if (i == iw) {
+            for (int j = 0; j < n; j++) {
+                double wanted = target_dis[i * n + j], w = weights[i * n + j];
+                double actual = hamming_dis(perm[i], perm[j]);
+                delta_cost -= w * sqr(wanted - actual);
+                double new_actual = hamming_dis(
+                        perm[jw],
+                        perm[j == iw           ? jw
+                                     : j == jw ? iw
+                                               : j]);
+                delta_cost += w * sqr(wanted - new_actual);
+            }
+        } else if (i == jw) {
+            for (int j = 0; j < n; j++) {
+                double wanted = target_dis[i * n + j], w = weights[i * n + j];
+                double actual = hamming_dis(perm[i], perm[j]);
+                delta_cost -= w * sqr(wanted - actual);
+                double new_actual = hamming_dis(
+                        perm[iw],
+                        perm[j == iw           ? jw
+                                     : j == jw ? iw
+                                               : j]);
+                delta_cost += w * sqr(wanted - new_actual);
+            }
+        } else {
+            int j = iw;
+            {
+                double wanted = target_dis[i * n + j], w = weights[i * n + j];
+                double actual = hamming_dis(perm[i], perm[j]);
+                delta_cost -= w * sqr(wanted - actual);
+                double new_actual = hamming_dis(perm[i], perm[jw]);
+                delta_cost += w * sqr(wanted - new_actual);
+            }
+            j = jw;
+            {
+                double wanted = target_dis[i * n + j], w = weights[i * n + j];
+                double actual = hamming_dis(perm[i], perm[j]);
+                delta_cost -= w * sqr(wanted - actual);
+                double new_actual = hamming_dis(perm[i], perm[iw]);
+                delta_cost += w * sqr(wanted - new_actual);
+            }
+        }
+    }
+
+    return delta_cost;
+}
+
+} // namespace polysemous_training
+
 namespace {
 
 /// optimize permutation to reproduce a distance table with Hamming distances
@@ -175,10 +265,6 @@ struct ReproduceWithHammingObjective : PermutationObjective {
     int nbits;
     double dis_weight_factor;
 
-    static double sqr(double x) {
-        return x * x;
-    }
-
     // weighting of distances: it is more important to reproduce small
     // distances well
     double dis_weight(double x) const {
@@ -190,73 +276,21 @@ struct ReproduceWithHammingObjective : PermutationObjective {
 
     // cost = quadratic difference between actual distance and Hamming distance
     double compute_cost(const int* perm) const override {
-        double cost = 0;
-        for (int i = 0; i < n; i++) {
-            for (int j = 0; j < n; j++) {
-                double wanted = target_dis[i * n + j];
-                double w = weights[i * n + j];
-                double actual = hamming_dis(perm[i], perm[j]);
-                cost += w * sqr(wanted - actual);
-            }
-        }
-        return cost;
+        return with_selected_simd_levels<polysemous_training::SIMD_LEVELS>(
+                [&]<SIMDLevel SL>() {
+                    return polysemous_training::hamming_compute_cost<SL>(
+                            n, perm, target_dis.data(), weights.data());
+                });
     }
 
     // what would the cost update be if iw and jw were swapped?
     // computed in O(n) instead of O(n^2) for the full re-computation
     double cost_update(const int* perm, int iw, int jw) const override {
-        double delta_cost = 0;
-
-        for (int i = 0; i < n; i++) {
-            if (i == iw) {
-                for (int j = 0; j < n; j++) {
-                    double wanted = target_dis[i * n + j],
-                           w = weights[i * n + j];
-                    double actual = hamming_dis(perm[i], perm[j]);
-                    delta_cost -= w * sqr(wanted - actual);
-                    double new_actual = hamming_dis(
-                            perm[jw],
-                            perm[j == iw           ? jw
-                                         : j == jw ? iw
-                                                   : j]);
-                    delta_cost += w * sqr(wanted - new_actual);
-                }
-            } else if (i == jw) {
-                for (int j = 0; j < n; j++) {
-                    double wanted = target_dis[i * n + j],
-                           w = weights[i * n + j];
-                    double actual = hamming_dis(perm[i], perm[j]);
-                    delta_cost -= w * sqr(wanted - actual);
-                    double new_actual = hamming_dis(
-                            perm[iw],
-                            perm[j == iw           ? jw
-                                         : j == jw ? iw
-                                                   : j]);
-                    delta_cost += w * sqr(wanted - new_actual);
-                }
-            } else {
-                int j = iw;
-                {
-                    double wanted = target_dis[i * n + j],
-                           w = weights[i * n + j];
-                    double actual = hamming_dis(perm[i], perm[j]);
-                    delta_cost -= w * sqr(wanted - actual);
-                    double new_actual = hamming_dis(perm[i], perm[jw]);
-                    delta_cost += w * sqr(wanted - new_actual);
-                }
-                j = jw;
-                {
-                    double wanted = target_dis[i * n + j],
-                           w = weights[i * n + j];
-                    double actual = hamming_dis(perm[i], perm[j]);
-                    delta_cost -= w * sqr(wanted - actual);
-                    double new_actual = hamming_dis(perm[i], perm[iw]);
-                    delta_cost += w * sqr(wanted - new_actual);
-                }
-            }
-        }
-
-        return delta_cost;
+        return with_selected_simd_levels<polysemous_training::SIMD_LEVELS>(
+                [&]<SIMDLevel SL>() {
+                    return polysemous_training::hamming_cost_update<SL>(
+                            n, perm, iw, jw, target_dis.data(), weights.data());
+                });
     }
 
     ReproduceWithHammingObjective(
@@ -306,14 +340,20 @@ double ReproduceDistancesObjective::get_source_dis(int i, int j) const {
     return source_dis[i * n + j];
 }
 
+namespace polysemous_training {
+
 // cost = quadratic difference between actual distance and Hamming distance
-double ReproduceDistancesObjective::compute_cost(const int* perm) const {
+template <>
+double distances_compute_cost<SIMDLevel::NONE>(
+        const ReproduceDistancesObjective& obj,
+        const int* perm) {
+    const int n = obj.n;
     double cost = 0;
     for (int i = 0; i < n; i++) {
         for (int j = 0; j < n; j++) {
-            double wanted = target_dis[i * n + j];
-            double w = weights[i * n + j];
-            double actual = get_source_dis(perm[i], perm[j]);
+            double wanted = obj.target_dis[i * n + j];
+            double w = obj.weights[i * n + j];
+            double actual = obj.get_source_dis(perm[i], perm[j]);
             cost += w * sqr(wanted - actual);
         }
     }
@@ -322,16 +362,22 @@ double ReproduceDistancesObjective::compute_cost(const int* perm) const {
 
 // what would the cost update be if iw and jw were swapped?
 // computed in O(n) instead of O(n^2) for the full re-computation
-double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
-        const {
+template <>
+double distances_cost_update<SIMDLevel::NONE>(
+        const ReproduceDistancesObjective& obj,
+        const int* perm,
+        int iw,
+        int jw) {
+    const int n = obj.n;
     double delta_cost = 0;
     for (int i = 0; i < n; i++) {
         if (i == iw) {
             for (int j = 0; j < n; j++) {
-                double wanted = target_dis[i * n + j], w = weights[i * n + j];
-                double actual = get_source_dis(perm[i], perm[j]);
+                double wanted = obj.target_dis[i * n + j],
+                       w = obj.weights[i * n + j];
+                double actual = obj.get_source_dis(perm[i], perm[j]);
                 delta_cost -= w * sqr(wanted - actual);
-                double new_actual = get_source_dis(
+                double new_actual = obj.get_source_dis(
                         perm[jw],
                         perm[j == iw           ? jw
                                      : j == jw ? iw
@@ -340,10 +386,11 @@ double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
             }
         } else if (i == jw) {
             for (int j = 0; j < n; j++) {
-                double wanted = target_dis[i * n + j], w = weights[i * n + j];
-                double actual = get_source_dis(perm[i], perm[j]);
+                double wanted = obj.target_dis[i * n + j],
+                       w = obj.weights[i * n + j];
+                double actual = obj.get_source_dis(perm[i], perm[j]);
                 delta_cost -= w * sqr(wanted - actual);
-                double new_actual = get_source_dis(
+                double new_actual = obj.get_source_dis(
                         perm[iw],
                         perm[j == iw           ? jw
                                      : j == jw ? iw
@@ -353,18 +400,20 @@ double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
         } else {
             int j = iw;
             {
-                double wanted = target_dis[i * n + j], w = weights[i * n + j];
-                double actual = get_source_dis(perm[i], perm[j]);
+                double wanted = obj.target_dis[i * n + j],
+                       w = obj.weights[i * n + j];
+                double actual = obj.get_source_dis(perm[i], perm[j]);
                 delta_cost -= w * sqr(wanted - actual);
-                double new_actual = get_source_dis(perm[i], perm[jw]);
+                double new_actual = obj.get_source_dis(perm[i], perm[jw]);
                 delta_cost += w * sqr(wanted - new_actual);
             }
             j = jw;
             {
-                double wanted = target_dis[i * n + j], w = weights[i * n + j];
-                double actual = get_source_dis(perm[i], perm[j]);
+                double wanted = obj.target_dis[i * n + j],
+                       w = obj.weights[i * n + j];
+                double actual = obj.get_source_dis(perm[i], perm[j]);
                 delta_cost -= w * sqr(wanted - actual);
-                double new_actual = get_source_dis(perm[i], perm[iw]);
+                double new_actual = obj.get_source_dis(perm[i], perm[iw]);
                 delta_cost += w * sqr(wanted - new_actual);
             }
         }
@@ -372,6 +421,25 @@ double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
     return delta_cost;
 }
 
+} // namespace polysemous_training
+
+double ReproduceDistancesObjective::compute_cost(const int* perm) const {
+    return with_selected_simd_levels<polysemous_training::SIMD_LEVELS>(
+            [&]<SIMDLevel SL>() {
+                return polysemous_training::distances_compute_cost<SL>(
+                        *this, perm);
+            });
+}
+
+double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
+        const {
+    return with_selected_simd_levels<polysemous_training::SIMD_LEVELS>(
+            [&]<SIMDLevel SL>() {
+                return polysemous_training::distances_cost_update<SL>(
+                        *this, perm, iw, jw);
+            });
+}
+
 ReproduceDistancesObjective::ReproduceDistancesObjective(
         int n_in,
         const double* source_dis_in,

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

@@ -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;

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.