faiss 0.6.4 → 0.6.5

data/vendor/faiss/faiss/impl/sq_fastscan_utils.h

@@ -0,0 +1,232 @@
+/*
+ * 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
+
+#include <algorithm>
+#include <cstring>
+
+#include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/utils/hamming.h>
+
+/** Helpers shared by IndexSQFastScan and IndexIVFSQFastScan.
+ *
+ * Both re-quantise a higher-precision scalar quantizer down to 4-bit nibbles
+ * for the SIMD scan and keep the original codes to rerank with, so the
+ * classification of quantizer types and the nibble packing have to agree
+ * between them.
+ */
+
+namespace faiss {
+namespace sq_fastscan {
+
+inline size_t roundup(size_t a, size_t b) {
+    return (a + b - 1) / b * b;
+}
+
+inline bool is_native_4bit(ScalarQuantizer::QuantizerType qtype) {
+    return qtype == ScalarQuantizer::QT_4bit ||
+            qtype == ScalarQuantizer::QT_4bit_uniform;
+}
+
+inline bool needs_rerank(ScalarQuantizer::QuantizerType qtype) {
+    return qtype == ScalarQuantizer::QT_6bit ||
+            qtype == ScalarQuantizer::QT_8bit ||
+            qtype == ScalarQuantizer::QT_8bit_uniform ||
+            qtype == ScalarQuantizer::QT_8bit_direct ||
+            qtype == ScalarQuantizer::QT_8bit_direct_signed;
+}
+
+inline bool is_fallback(ScalarQuantizer::QuantizerType qtype) {
+    return !is_native_4bit(qtype) && !needs_rerank(qtype);
+}
+
+inline bool is_uniform_range(ScalarQuantizer::QuantizerType qtype) {
+    return qtype == ScalarQuantizer::QT_8bit_uniform ||
+            qtype == ScalarQuantizer::QT_8bit_direct ||
+            qtype == ScalarQuantizer::QT_8bit_direct_signed;
+}
+
+inline void get_uniform_range(
+        const ScalarQuantizer& sq,
+        float& vmin,
+        float& vdiff) {
+    if (sq.qtype == ScalarQuantizer::QT_8bit_direct) {
+        vmin = 0;
+        vdiff = 255;
+    } else if (sq.qtype == ScalarQuantizer::QT_8bit_direct_signed) {
+        vmin = -128;
+        vdiff = 255;
+    } else {
+        vmin = sq.trained[0];
+        vdiff = sq.trained[1];
+    }
+}
+
+/// Bits per dimension the quantizer stores.
+inline int sq_bits(ScalarQuantizer::QuantizerType qtype) {
+    if (qtype == ScalarQuantizer::QT_6bit) {
+        return 6;
+    }
+    return is_native_4bit(qtype) ? 4 : 8;
+}
+
+/// Bits left over once the scan has taken its 4. Zero for a native 4-bit
+/// code, which needs no rerank and therefore stores nothing extra.
+inline int sq_rerank_bits(ScalarQuantizer::QuantizerType qtype) {
+    return sq_bits(qtype) - 4;
+}
+
+/// Bytes per vector needed to hold the leftover bits of every dimension.
+inline size_t sq_rerank_size(int d, ScalarQuantizer::QuantizerType qtype) {
+    const int bits = sq_rerank_bits(qtype);
+    return bits > 0 ? (size_t(d) * bits + 7) / 8 : 0;
+}
+
+/// The code for one value at the quantizer's own width.
+inline int sq_code_of(float x, float vmin, float vdiff, int bits) {
+    const int maxv = (1 << bits) - 1;
+    const float f = (vdiff > 0) ? (x - vmin) / vdiff : 0.0f;
+    return std::min(maxv, std::max(0, int(f * float(maxv))));
+}
+
+/// Per-dimension range of the trained quantizer, resolved once for a batch.
+struct SQRanges {
+    SQRanges(const ScalarQuantizer& sq, int d)
+            : uniform(is_uniform_range(sq.qtype) ||
+                      sq.qtype == ScalarQuantizer::QT_4bit_uniform) {
+        if (uniform) {
+            get_uniform_range(sq, vmin_s, vdiff_s);
+        } else {
+            vmin_arr = sq.trained.data();
+            vdiff_arr = sq.trained.data() + d;
+        }
+    }
+
+    float vmin(int m) const {
+        return uniform ? vmin_s : vmin_arr[m];
+    }
+    float vdiff(int m) const {
+        return uniform ? vdiff_s : vdiff_arr[m];
+    }
+
+    const bool uniform;
+    const float* vmin_arr = nullptr;
+    const float* vdiff_arr = nullptr;
+    float vmin_s = 0, vdiff_s = 0;
+};
+
+/** Split each code into the top 4 bits the SIMD kernel scans and the rest,
+ *  which rerank keeps. `hi` is packed two dimensions per byte as
+ *  pq4_pack_codes expects. `lo` is bit-packed at `bits - 4` per dimension.
+ *  Total width is the quantizer's own, not that width plus 4.
+ */
+inline void float_to_split_codes(
+        const float* x,
+        uint8_t* hi,
+        uint8_t* lo,
+        idx_t n,
+        int d,
+        size_t M2,
+        const ScalarQuantizer& sq) {
+    const int bits = sq_bits(sq.qtype);
+    const int lo_bits = bits - 4;
+    const size_t hi_stride = M2 / 2;
+    const size_t lo_stride = sq_rerank_size(d, sq.qtype);
+    const SQRanges ranges(sq, d);
+    for (idx_t i = 0; i < n; i++) {
+        const float* xi = x + i * d;
+        uint8_t* hdst = hi + i * hi_stride;
+        memset(hdst, 0, hi_stride);
+        uint8_t* ldst = lo + i * lo_stride;
+        memset(ldst, 0, lo_stride);
+        BitstringWriter bsw(ldst, lo_stride);
+        for (int m = 0; m < d; m++) {
+            const int code =
+                    sq_code_of(xi[m], ranges.vmin(m), ranges.vdiff(m), bits);
+            const int h = code >> lo_bits;
+            hdst[m / 2] |= (m & 1) ? uint8_t(h << 4) : uint8_t(h);
+            bsw.write(code & ((1 << lo_bits) - 1), lo_bits);
+        }
+    }
+}
+
+/// Rebuild the code values of one vector from its two halves.
+inline void split_to_codes(
+        const uint8_t* hi,
+        const uint8_t* lo,
+        int d,
+        ScalarQuantizer::QuantizerType qtype,
+        int* codes_out) {
+    const int lo_bits = sq_rerank_bits(qtype);
+    BitstringReader bsr(lo, sq_rerank_size(d, qtype));
+    for (int m = 0; m < d; m++) {
+        const int h = (m & 1) ? (hi[m / 2] >> 4) : (hi[m / 2] & 0xF);
+        codes_out[m] = (h << lo_bits) | int(bsr.read(lo_bits));
+    }
+}
+
+inline void float_to_4bit_nibbles(
+        const float* x,
+        uint8_t* nibbles,
+        idx_t n,
+        int d,
+        int M2,
+        const ScalarQuantizer& sq) {
+    const bool is_uniform = is_uniform_range(sq.qtype) ||
+            sq.qtype == ScalarQuantizer::QT_4bit_uniform;
+    const float* vmin_arr = nullptr;
+    const float* vdiff_arr = nullptr;
+    float vmin_s = 0, inv_vdiff_s = 0;
+
+    if (is_uniform) {
+        float vdiff;
+        get_uniform_range(sq, vmin_s, vdiff);
+        inv_vdiff_s = (vdiff > 0) ? (1.0f / vdiff) : 0;
+    } else {
+        vmin_arr = sq.trained.data();
+        vdiff_arr = sq.trained.data() + d;
+    }
+
+    const int half = M2 / 2;
+    for (idx_t i = 0; i < n; i++) {
+        const float* xi = x + i * d;
+        uint8_t* dst = nibbles + i * half;
+        memset(dst, 0, half);
+        for (int m = 0; m + 1 < d; m += 2) {
+            float f0, f1;
+            if (is_uniform) {
+                f0 = (xi[m] - vmin_s) * inv_vdiff_s;
+                f1 = (xi[m + 1] - vmin_s) * inv_vdiff_s;
+            } else {
+                f0 = (vdiff_arr[m] > 0) ? (xi[m] - vmin_arr[m]) / vdiff_arr[m]
+                                        : 0;
+                f1 = (vdiff_arr[m + 1] > 0)
+                        ? (xi[m + 1] - vmin_arr[m + 1]) / vdiff_arr[m + 1]
+                        : 0;
+            }
+            uint8_t lo = (uint8_t)std::min(15, std::max(0, (int)(f0 * 15.0f)));
+            uint8_t hi = (uint8_t)std::min(15, std::max(0, (int)(f1 * 15.0f)));
+            dst[m / 2] = lo | (hi << 4);
+        }
+        if (d & 1) {
+            float f;
+            if (is_uniform) {
+                f = (xi[d - 1] - vmin_s) * inv_vdiff_s;
+            } else {
+                f = (vdiff_arr[d - 1] > 0)
+                        ? (xi[d - 1] - vmin_arr[d - 1]) / vdiff_arr[d - 1]
+                        : 0;
+            }
+            dst[(d - 1) / 2] =
+                    (uint8_t)std::min(15, std::max(0, (int)(f * 15.0f)));
+        }
+    }
+}
+
+} // namespace sq_fastscan
+} // namespace faiss

data/vendor/faiss/faiss/index_factory.cpp

@@ -35,17 +35,20 @@
 #include <faiss/IndexIVFPQR.h>
 #include <faiss/IndexIVFRaBitQ.h>
 #include <faiss/IndexIVFRaBitQFastScan.h>
+#include <faiss/IndexIVFSQFastScan.h>
 #include <faiss/IndexIVFSpectralHash.h>
 #include <faiss/IndexLSH.h>
 #include <faiss/IndexLattice.h>
 #include <faiss/IndexNSG.h>
 #include <faiss/IndexPQ.h>
 #include <faiss/IndexPQFastScan.h>
+#include <faiss/IndexPiPNN.h>
 #include <faiss/IndexPreTransform.h>
 #include <faiss/IndexRaBitQ.h>
 #include <faiss/IndexRaBitQFastScan.h>
 #include <faiss/IndexRefine.h>
 #include <faiss/IndexRowwiseMinMax.h>
+#include <faiss/IndexSQFastScan.h>
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/VectorTransform.h>
 
@@ -177,6 +180,11 @@ std::map<std::string, ScalarQuantizer::QuantizerType> sq_types = {
 };
 const std::string sq_pattern =
         "(SQ0|SQ4|SQ8|SQ6|SQfp16|SQbf16|SQ8_direct_signed|SQ8_direct|SQtqmse1|SQtqmse2|SQtqmse3|SQtqmse4|SQtqmse8|SQtq2|SQtq3|SQtq4|SQtq5)";
+const std::string sq_fs_pattern =
+        "(SQ4|SQ6|SQ8|SQ8_direct_signed|SQ8_direct)fs(_[0-9]+)?";
+// All SQ types with "fs" suffix for IndexIVFSQFastScan, optional _bbs
+const std::string ivf_sq_fs_pattern =
+        "(SQ0|SQ4|SQ8|SQ6|SQfp16|SQbf16|SQ8_direct_signed|SQ8_direct|SQtqmse1|SQtqmse2|SQtqmse3|SQtqmse4|SQtqmse8|SQtq2|SQtq3|SQtq4|SQtq5)fs(_[0-9]+)?";
 
 std::map<std::string, AdditiveQuantizer::Search_type_t> aq_search_type = {
         {"_Nfloat", AdditiveQuantizer::ST_norm_float},
@@ -310,6 +318,11 @@ Index* parse_coarse_quantizer(
         int R = std::stoi(sm[2]);
         return new IndexNSGFlat(d, R, mt);
     }
+    if (match("IVF([0-9]+[kM]?)_PiPNN([0-9]*)")) {
+        nlist = parse_nlist(sm[1].str());
+        int R = mres_to_int(sm[2], 64);
+        return new IndexPiPNNFlat(d, R, mt);
+    }
 #ifdef FAISS_ENABLE_SVS
     if (match("IVF([0-9]+[kM]?)_SVSVamana([0-9]*)(_.+)?")) {
         nlist = parse_nlist(sm[1].str());
@@ -389,6 +402,17 @@ IndexIVF* parse_IndexIVF(
         return new IndexIVFFlatPanorama(
                 get_q(), d, nlist, nlevels, mt, own_il, bs);
     }
+    if (match(ivf_sq_fs_pattern)) {
+        int bbs = mres_to_int(sm[2], 32, 1);
+        return new IndexIVFSQFastScan(
+                get_q(),
+                d,
+                nlist,
+                sq_types[sm[1].str()],
+                mt,
+                bbs,
+                /*by_residual=*/true);
+    }
     if (match(sq_pattern)) {
         return new IndexIVFScalarQuantizer(
                 get_q(),
@@ -637,6 +661,30 @@ IndexNSG* parse_IndexNSG(
     return nullptr;
 }
 
+/***************************************************************
+ * Parse IndexPiPNN
+ */
+
+IndexPiPNN* parse_IndexPiPNN(
+        const std::string& code_string,
+        int d,
+        MetricType mt,
+        int pipnn_R) {
+    std::smatch sm;
+    auto match = [&sm, &code_string](const std::string& pattern) {
+        return re_match(code_string, pattern, sm);
+    };
+
+    if (match("Flat|")) {
+        return new IndexPiPNNFlat(d, pipnn_R, mt);
+    }
+    if (match(sq_pattern)) {
+        return new IndexPiPNNSQ(d, sq_types[sm[1].str()], pipnn_R, mt);
+    }
+
+    return nullptr;
+}
+
 #ifdef FAISS_ENABLE_SVS
 /***************************************************************
  * Parse IndexSVS
@@ -831,6 +879,12 @@ Index* parse_other_indexes(
         return new IndexLattice(d, M, nbit, r2);
     }
 
+    // IndexSQFastScan (must be checked before IndexScalarQuantizer)
+    if (match(sq_fs_pattern)) {
+        int bbs = mres_to_int(sm[2], 32, 1);
+        return new IndexSQFastScan(d, sq_types[sm[1].str()], metric, bbs);
+    }
+
     // IndexScalarQuantizer
     if (match(sq_pattern)) {
         return new IndexScalarQuantizer(d, sq_types[description], metric);
@@ -1156,6 +1210,28 @@ std::unique_ptr<Index> index_factory_sub(
         return std::unique_ptr<Index>(index);
     }
 
+    // PiPNN variants (separator "," or "_", as for HNSW and NSG)
+    if (re_match(description, "PiPNN([0-9]*)([,_].*)?", sm)) {
+        int pipnn_R = mres_to_int(sm[1], 64);
+        // We also accept empty code string (synonym of Flat)
+        std::string code_string =
+                sm[2].length() > 0 ? sm[2].str().substr(1) : "";
+        if (verbose) {
+            printf("parsing PiPNN string %s code_string=%s pipnn_R=%d\n",
+                   description.c_str(),
+                   code_string.c_str(),
+                   pipnn_R);
+        }
+
+        IndexPiPNN* index = parse_IndexPiPNN(code_string, d, metric, pipnn_R);
+        FAISS_THROW_IF_NOT_FMT(
+                index,
+                "could not parse PiPNN code description %s in %s",
+                code_string.c_str(),
+                description.c_str());
+        return std::unique_ptr<Index>(index);
+    }
+
     // IndexRowwiseMinMax, fp32 version
     if (description.compare(0, 7, "MinMax,") == 0) {
         size_t comma = description.find(',');

data/vendor/faiss/faiss/invlists/InvertedLists.cpp

@@ -127,7 +127,9 @@ size_t InvertedLists::copy_subset_to(
         } else if (subset_type == SUBSET_TYPE_ID_MOD) {
             for (size_t i = 0; i < n; i++) {
                 idx_t id = ids_in[i];
-                if (id % a1 == a2) {
+                // C++ gives a negative remainder, so a plain `id % a1` never
+                // matches for a negative id. Fold it into [0, a1) instead.
+                if (((id % a1) + a1) % a1 == a2) {
                     oivf.add_entry(
                             list_no,
                             get_single_id(list_no, i),

data/vendor/faiss/faiss/svs/IndexSVSVamana.h

@@ -90,8 +90,9 @@ inline svs_runtime::StorageKind to_svs_storage_kind(SVSStorageKind kind) {
 }
 
 struct IndexSVSVamana : Index {
-    size_t graph_max_degree;
-    size_t prune_to;
+    // Same defaults as IndexSVSVamana(d, /*degree=*/64).
+    size_t graph_max_degree = 64;
+    size_t prune_to = 60;
     float alpha = 1.2;
     size_t search_window_size = 10;
     size_t search_buffer_capacity = 10;

data/vendor/faiss/faiss/utils/AlignedTable.h

@@ -117,6 +117,8 @@ template <class T, int A = 32>
 struct AlignedTable {
     AlignedTableTightAlloc<T, A> tab;
     size_t numel = 0;
+    // If true, resize() reallocates only when n is larger than the capacity.
+    bool keep_capacity = false;
 
     static size_t round_capacity(size_t n) {
         if (n == 0) {
@@ -141,7 +143,10 @@ struct AlignedTable {
     }
 
     void resize(size_t n) {
-        tab.resize(round_capacity(n));
+        const size_t capacity = round_capacity(n);
+        if (!keep_capacity || capacity > tab.numel) {
+            tab.resize(capacity);
+        }
         numel = n;
     }
 

data/vendor/faiss/faiss/utils/Heap.cpp

@@ -194,10 +194,11 @@ void merge_knn_results(
             const idx_t* I_in = all_labels + i * k;
             int heap_size = 0;
 
-            // push the first element of each shard (if not -1)
+            // push the first element of each shard (if not -1). A caller id
+            // may be negative, so only -1 marks a missing result.
             for (long s = 0; s < static_cast<long>(nshard); s++) {
                 pointer[s] = 0;
-                if (I_in[stride * s] >= 0) {
+                if (I_in[stride * s] != -1) {
                     heap_push<C>(
                             ++heap_size,
                             heap_vals,
@@ -221,7 +222,7 @@ void merge_knn_results(
                 // pop from shard, advance pointer for this shard
                 heap_pop<C>(heap_size--, heap_vals, shard_ids);
                 p++;
-                if (static_cast<size_t>(p) < k && I_in[stride * s + p] >= 0) {
+                if (static_cast<size_t>(p) < k && I_in[stride * s + p] != -1) {
                     heap_push<C>(
                             ++heap_size,
                             heap_vals,

data/vendor/faiss/faiss/utils/approx_topk_hamming/approx_topk_hamming.h

@@ -21,19 +21,30 @@ using simd8uint32 = simd8uint32_tpl<SINGLE_SIMD_LEVEL_256>;
 using simd16uint16 = simd16uint16_tpl<SINGLE_SIMD_LEVEL_256>;
 
 // HeapWithBucketsForHamming32 uses simd8uint32 under the
-// hood.
-
-template <typename C, uint32_t NBUCKETS, uint32_t N, typename HammingComputerT>
+// hood. SL is the SIMD level of the file that uses it, so each level compiles
+// its own copy under its own name.
+
+template <
+        typename C,
+        uint32_t NBUCKETS,
+        uint32_t N,
+        typename HammingComputerT,
+        SIMDLevel SL>
 struct HeapWithBucketsForHamming32 {
     // this case was not implemented yet.
 };
 
-template <uint32_t NBUCKETS, uint32_t N, typename HammingComputerT>
+template <
+        uint32_t NBUCKETS,
+        uint32_t N,
+        typename HammingComputerT,
+        SIMDLevel SL>
 struct HeapWithBucketsForHamming32<
         CMax<int, int64_t>,
         NBUCKETS,
         N,
-        HammingComputerT> {
+        HammingComputerT,
+        SL> {
     static constexpr uint32_t NBUCKETS_8 = NBUCKETS / 8;
     static_assert(
             (NBUCKETS) > 0 && ((NBUCKETS % 8) == 0),

data/vendor/faiss/faiss/utils/bf16.h

@@ -14,6 +14,10 @@
 #include <immintrin.h>
 #endif
 
+#if defined(__riscv_vector)
+#include <riscv_vector.h>
+#endif
+
 namespace faiss {
 
 namespace {
@@ -38,6 +42,8 @@ inline float decode_bf16(const uint16_t v) {
     return fp.as_f32;
 }
 
+namespace {
+
 inline void encode_bf16_simd(const float* src, uint16_t* dst, size_t n) {
     size_t i = 0;
 #ifdef __AVX512BF16__
@@ -61,10 +67,40 @@ inline void decode_bf16_simd(const uint16_t* src, float* dst, size_t n) {
         w = _mm512_slli_epi32(w, 16);
         _mm512_storeu_ps(dst + i, _mm512_castsi512_ps(w));
     }
+#elif defined(__riscv_vector)
+    if (n == 1) {
+        dst[0] = decode_bf16(src[0]);
+        return;
+    }
+    if (n < 16) {
+        for (; i < n; ++i) {
+            dst[i] = decode_bf16(src[i]);
+        }
+        return;
+    }
+    const size_t vlmax = __riscv_vsetvlmax_e16m1();
+    for (; i + vlmax <= n; i += vlmax) {
+        vuint16m1_t v = __riscv_vle16_v_u16m1(src + i, vlmax);
+        vuint32m2_t w = __riscv_vzext_vf2_u32m2(v, vlmax);
+        w = __riscv_vsll_vx_u32m2(w, 16, vlmax);
+        __riscv_vse32_v_f32m2(
+                dst + i, __riscv_vreinterpret_v_u32m2_f32m2(w), vlmax);
+    }
+    if (i < n) {
+        const size_t vl = __riscv_vsetvl_e16m1(n - i);
+        vuint16m1_t v = __riscv_vle16_v_u16m1(src + i, vl);
+        vuint32m2_t w = __riscv_vzext_vf2_u32m2(v, vl);
+        w = __riscv_vsll_vx_u32m2(w, 16, vl);
+        __riscv_vse32_v_f32m2(
+                dst + i, __riscv_vreinterpret_v_u32m2_f32m2(w), vl);
+        i += vl;
+    }
 #endif
     for (; i < n; i++) {
         dst[i] = decode_bf16(src[i]);
     }
 }
 
+} // namespace
+
 } // namespace faiss