faiss 0.6.0 → 0.6.2

data/vendor/faiss/faiss/IndexScalarQuantizer.cpp

@@ -14,6 +14,9 @@
 
 #include <omp.h>
 
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/expanded_scanners.h>
+
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ScalarQuantizer.h>
@@ -26,10 +29,10 @@ namespace faiss {
  ********************************************************************/
 
 IndexScalarQuantizer::IndexScalarQuantizer(
-        int d,
+        int d_in,
         ScalarQuantizer::QuantizerType qtype,
         MetricType metric)
-        : IndexFlatCodes(0, d, metric), sq(d, qtype) {
+        : IndexFlatCodes(0, d_in, metric), sq(d_in, qtype) {
     is_trained = qtype == ScalarQuantizer::QT_fp16 ||
             qtype == ScalarQuantizer::QT_8bit_direct ||
             qtype == ScalarQuantizer::QT_bf16 ||
@@ -117,21 +120,29 @@ void IndexScalarQuantizer::sa_decode(idx_t n, const uint8_t* bytes, float* x)
  ********************************************************************/
 
 IndexIVFScalarQuantizer::IndexIVFScalarQuantizer(
-        Index* quantizer,
-        size_t d,
-        size_t nlist,
+        Index* quantizer_,
+        size_t d_,
+        size_t nlist_,
         ScalarQuantizer::QuantizerType qtype,
         MetricType metric,
-        bool by_residual,
-        bool own_invlists)
-        : IndexIVF(quantizer, d, nlist, 0, metric, own_invlists), sq(d, qtype) {
+        bool by_residual_,
+        bool own_invlists_)
+        : IndexIVF(quantizer_, d_, nlist_, 0, metric, own_invlists_),
+          sq(d_, qtype) {
     code_size = sq.code_size;
-    this->by_residual = by_residual;
+    this->by_residual = by_residual_;
     if (invlists) {
         // was not known at construction time
         invlists->code_size = code_size;
     }
     is_trained = false;
+    if (qtype == ScalarQuantizer::QT_0bit) {
+        by_residual = false;
+        is_trained = true; // no training needed
+    }
+    if (ScalarQuantizer::TurboQuantRefine::is_turboq_full(qtype)) {
+        by_residual = false;
+    }
 }
 
 IndexIVFScalarQuantizer::IndexIVFScalarQuantizer() : IndexIVF() {
@@ -141,7 +152,7 @@ IndexIVFScalarQuantizer::IndexIVFScalarQuantizer() : IndexIVF() {
 void IndexIVFScalarQuantizer::train_encoder(
         idx_t n,
         const float* x,
-        const idx_t* assign) {
+        const idx_t* /*assign*/) {
     sq.train(n, x);
 }
 
@@ -155,6 +166,19 @@ void IndexIVFScalarQuantizer::encode_vectors(
         const idx_t* list_nos,
         uint8_t* codes,
         bool include_listnos) const {
+    if (sq.code_size == 0) {
+        // QT_0bit: nothing to encode, but handle coarse codes if needed
+        if (include_listnos) {
+            size_t coarse_size = coarse_code_size();
+            for (idx_t i = 0; i < n; i++) {
+                int64_t list_no = list_nos[i];
+                if (list_no >= 0) {
+                    encode_listno(list_no, codes + i * coarse_size);
+                }
+            }
+        }
+        return;
+    }
     std::unique_ptr<ScalarQuantizer::SQuantizer> squant(sq.select_quantizer());
     size_t coarse_size = include_listnos ? coarse_code_size() : 0;
     memset(codes, 0, (code_size + coarse_size) * n);
@@ -185,14 +209,42 @@ void IndexIVFScalarQuantizer::encode_vectors(
 void IndexIVFScalarQuantizer::decode_vectors(
         idx_t n,
         const uint8_t* codes,
-        const idx_t*,
+        const idx_t* list_nos,
         float* x) const {
+    if (sq.code_size == 0 && list_nos) {
+        // QT_0bit: reconstruct centroids if list_nos provided
+        for (idx_t i = 0; i < n; i++) {
+            quantizer->reconstruct(list_nos[i], x + i * d);
+        }
+        return;
+    }
     FAISS_THROW_IF_NOT(is_trained);
-    return sq.decode(codes, x, n);
+    sq.decode(codes, x, n);
+    if (by_residual) {
+        FAISS_THROW_IF_NOT_MSG(
+                list_nos, "decode_vectors with by_residual requires list_nos");
+#pragma omp parallel for if (n > 1000)
+        for (idx_t i = 0; i < n; i++) {
+            std::vector<float> centroid(d);
+            quantizer->reconstruct(list_nos[i], centroid.data());
+            for (size_t j = 0; j < static_cast<size_t>(d); j++) {
+                x[i * d + j] += centroid[j];
+            }
+        }
+    }
 }
 
 void IndexIVFScalarQuantizer::sa_decode(idx_t n, const uint8_t* codes, float* x)
         const {
+    if (sq.code_size == 0) {
+        size_t coarse_size = coarse_code_size();
+        for (idx_t i = 0; i < n; i++) {
+            const uint8_t* code = codes + i * coarse_size;
+            int64_t list_no = decode_listno(code);
+            quantizer->reconstruct(list_no, x + i * d);
+        }
+        return;
+    }
     std::unique_ptr<ScalarQuantizer::SQuantizer> squant(sq.select_quantizer());
     size_t coarse_size = coarse_code_size();
 
@@ -208,7 +260,7 @@ void IndexIVFScalarQuantizer::sa_decode(idx_t n, const uint8_t* codes, float* x)
             squant->decode_vector(code + coarse_size, xi);
             if (by_residual) {
                 quantizer->reconstruct(list_no, residual.data());
-                for (size_t j = 0; j < d; j++) {
+                for (size_t j = 0; j < static_cast<size_t>(d); j++) {
                     xi[j] += residual[j];
                 }
             }
@@ -223,6 +275,23 @@ void IndexIVFScalarQuantizer::add_core(
         const idx_t* coarse_idx,
         void* inverted_list_context) {
     FAISS_THROW_IF_NOT(is_trained);
+    if (sq.code_size == 0) {
+        // QT_0bit: just add IDs with empty codes
+        uint8_t dummy_code = 0;
+        DirectMapAdd dm_add(direct_map, n, xids);
+        for (idx_t i = 0; i < n; i++) {
+            int64_t list_no = coarse_idx[i];
+            if (list_no >= 0) {
+                int64_t id = xids ? xids[i] : ntotal + i;
+                size_t ofs = invlists->add_entry(list_no, id, &dummy_code);
+                dm_add.add(i, list_no, ofs);
+            } else {
+                dm_add.add(i, -1, 0);
+            }
+        }
+        ntotal += n;
+        return;
+    }
 
     std::unique_ptr<ScalarQuantizer::SQuantizer> squant(sq.select_quantizer());
 
@@ -236,7 +305,7 @@ void IndexIVFScalarQuantizer::add_core(
         int rank = omp_get_thread_num();
 
         // each thread takes care of a subset of lists
-        for (size_t i = 0; i < n; i++) {
+        for (idx_t i = 0; i < n; i++) {
             int64_t list_no = coarse_idx[i];
             if (list_no >= 0 && list_no % nt == rank) {
                 int64_t id = xids ? xids[i] : ntotal + i;
@@ -267,27 +336,88 @@ void IndexIVFScalarQuantizer::add_core(
 InvertedListScanner* IndexIVFScalarQuantizer::get_InvertedListScanner(
         bool store_pairs,
         const IDSelector* sel,
-        const IVFSearchParameters*) const {
-    return sq.select_InvertedListScanner(
-            metric_type, quantizer, store_pairs, sel, by_residual);
+        const IVFSearchParameters* search_params) const {
+    if (!ScalarQuantizer::TurboQuantRefine::is_turboq_full(sq.qtype)) {
+        return sq.select_InvertedListScanner(
+                metric_type, quantizer, store_pairs, sel, by_residual);
+    }
+
+    // TurboQ full types: create a TQ-specific scanner that supports
+    // search params (qb, int_qjl) and pre-screening.
+    uint8_t tq_qb = 0;
+    bool tq_int_qjl = false;
+    if (auto* tp = dynamic_cast<const IVFSQTurboQSearchParameters*>(
+                search_params)) {
+        tq_qb = tp->qb;
+        tq_int_qjl = tp->int_qjl;
+    }
+    using TurboQDC = ScalarQuantizer::TurboQuantRefine::DistanceComputer;
+    auto* dc = static_cast<TurboQDC*>(sq.get_distance_computer(metric_type));
+    dc->configure(tq_qb, tq_int_qjl);
+
+    struct TQScanner : InvertedListScanner {
+        std::unique_ptr<TurboQDC> dc;
+
+        explicit TQScanner(
+                TurboQDC* dc_in,
+                bool store_pairs_in,
+                const IDSelector* sel_in,
+                bool keep_max_in)
+                : InvertedListScanner(store_pairs_in, sel_in), dc(dc_in) {
+            this->keep_max = keep_max_in;
+        }
+
+        void set_query(const float* query) override {
+            dc->set_query(query);
+        }
+
+        void set_list(idx_t list_no_in, float) override {
+            this->list_no = list_no_in;
+        }
+
+        float distance_to_code(const uint8_t* code) const final {
+            return dc->distance_to_code(code);
+        }
+
+        size_t scan_codes(
+                size_t list_size,
+                const uint8_t* codes,
+                const idx_t* ids,
+                ResultHandler& handler) const override {
+            dc->set_prescreen_threshold(&handler.threshold, !keep_max);
+            size_t nup = run_scan_codes(*this, list_size, codes, ids, handler);
+            dc->clear_prescreen_threshold();
+            return nup;
+        }
+    };
+
+    auto* scanner = new TQScanner(
+            dc, store_pairs, sel, is_similarity_metric(metric_type));
+    scanner->code_size = code_size;
+    return scanner;
 }
 
 void IndexIVFScalarQuantizer::reconstruct_from_offset(
         int64_t list_no,
         int64_t offset,
         float* recons) const {
-    const uint8_t* code = invlists->get_single_code(list_no, offset);
+    if (sq.code_size == 0) {
+        // QT_0bit: reconstruct from centroid
+        quantizer->reconstruct(list_no, recons);
+        return;
+    }
+    InvertedLists::ScopedCodes sc(invlists, list_no, offset);
 
     if (by_residual) {
         std::vector<float> centroid(d);
         quantizer->reconstruct(list_no, centroid.data());
 
-        sq.decode(code, recons, 1);
+        sq.decode(sc.get(), recons, 1);
         for (int i = 0; i < d; ++i) {
             recons[i] += centroid[i];
         }
     } else {
-        sq.decode(code, recons, 1);
+        sq.decode(sc.get(), recons, 1);
     }
 }
 

data/vendor/faiss/faiss/IndexScalarQuantizer.h

@@ -63,6 +63,16 @@ struct IndexScalarQuantizer : IndexFlatCodes {
  * distances are computed.
  */
 
+/// Search parameters for TurboQuant full types (QT_*_tq).
+struct IVFSQTurboQSearchParameters : IVFSearchParameters {
+    /// Query quantization bits for integer MSE pre-screening.
+    /// 0 = float path (default), 1-8 = integer popcount path.
+    uint8_t qb = 0;
+
+    /// Also use integer popcount for QJL stage (requires qb > 0).
+    bool int_qjl = false;
+};
+
 struct IndexIVFScalarQuantizer : IndexIVF {
     ScalarQuantizer sq;
 

data/vendor/faiss/faiss/IndexShards.cpp

@@ -46,27 +46,28 @@ void translate_labels(int64_t n, idx_t* labels, int64_t translation) {
 
 template <typename IndexT>
 IndexShardsTemplate<IndexT>::IndexShardsTemplate(
-        idx_t d,
+        idx_t d_,
         bool threaded,
-        bool successive_ids)
-        : ThreadedIndex<IndexT>(d, threaded), successive_ids(successive_ids) {
+        bool successive_ids_)
+        : ThreadedIndex<IndexT>(static_cast<int>(d_), threaded),
+          successive_ids(successive_ids_) {
     sync_d(this);
 }
 
 template <typename IndexT>
 IndexShardsTemplate<IndexT>::IndexShardsTemplate(
-        int d,
+        int d_,
         bool threaded,
-        bool successive_ids)
-        : ThreadedIndex<IndexT>(d, threaded), successive_ids(successive_ids) {
+        bool successive_ids_)
+        : ThreadedIndex<IndexT>(d_, threaded), successive_ids(successive_ids_) {
     sync_d(this);
 }
 
 template <typename IndexT>
 IndexShardsTemplate<IndexT>::IndexShardsTemplate(
         bool threaded,
-        bool successive_ids)
-        : ThreadedIndex<IndexT>(threaded), successive_ids(successive_ids) {
+        bool successive_ids_)
+        : ThreadedIndex<IndexT>(threaded), successive_ids(successive_ids_) {
     sync_d(this);
 }
 
@@ -202,7 +203,7 @@ void IndexShardsTemplate<IndexT>::search(
         const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
 
-    int64_t nshard = this->count();
+    int nshard = this->count();
 
     std::vector<distance_t> all_distances(nshard * k * n);
     std::vector<idx_t> all_labels(nshard * k * n);

data/vendor/faiss/faiss/IndexShardsIVF.cpp

@@ -43,13 +43,14 @@ void translate_labels(int64_t n, idx_t* labels, int64_t translation) {
  ************************************************************/
 
 IndexShardsIVF::IndexShardsIVF(
-        Index* quantizer,
-        size_t nlist,
+        Index* quantizer_,
+        size_t nlist_,
         bool threaded,
-        bool successive_ids)
-        : IndexShardsTemplate<Index>(quantizer->d, threaded, successive_ids),
-          Level1Quantizer(quantizer, nlist) {
-    is_trained = quantizer->is_trained && quantizer->ntotal == nlist;
+        bool successive_ids_)
+        : IndexShardsTemplate<Index>(quantizer_->d, threaded, successive_ids_),
+          Level1Quantizer(quantizer_, nlist_) {
+    is_trained = quantizer_->is_trained &&
+            quantizer_->ntotal == static_cast<idx_t>(nlist_);
 }
 
 void IndexShardsIVF::addIndex(Index* index) {
@@ -73,11 +74,11 @@ void IndexShardsIVF::train(idx_t n, const component_t* x) {
     for (size_t i = 0; i < indices_.size(); i++) {
         Index* index = indices_[i].first;
         auto index_ivf = dynamic_cast<IndexIVFInterface*>(index);
-        Index* quantizer = index_ivf->quantizer;
-        if (!quantizer->is_trained) {
-            quantizer->train(nlist, centroids.data());
+        Index* sub_quantizer = index_ivf->quantizer;
+        if (!sub_quantizer->is_trained) {
+            sub_quantizer->train(nlist, centroids.data());
         }
-        quantizer->add(nlist, centroids.data());
+        sub_quantizer->add(nlist, centroids.data());
         // finish training
         index->train(n, x);
     }
@@ -132,9 +133,9 @@ void IndexShardsIVF::add_with_ids(
         }
         ids = aids.data();
     }
-    idx_t d = this->d;
+    idx_t cur_d = this->d;
 
-    auto fn = [n, ids, x, nshard, d, Iq](int no, Index* index) {
+    auto fn = [n, ids, x, nshard, cur_d, Iq](int no, Index* index) {
         idx_t i0 = (idx_t)no * n / nshard;
         idx_t i1 = ((idx_t)no + 1) * n / nshard;
         auto index_ivf = dynamic_cast<IndexIVF*>(index);
@@ -144,7 +145,10 @@ void IndexShardsIVF::add_with_ids(
         }
 
         index_ivf->add_core(
-                i1 - i0, x + i0 * d, ids ? ids + i0 : nullptr, Iq.data() + i0);
+                i1 - i0,
+                x + i0 * cur_d,
+                ids ? ids + i0 : nullptr,
+                Iq.data() + i0);
 
         if (index->verbose) {
             printf("end add shard %d on %" PRId64 " points\n", no, i1 - i0);
@@ -179,7 +183,7 @@ void IndexShardsIVF::search(
 
     quantizer->search(n, x, nprobe, Dq.data(), Iq.data());
 
-    int64_t nshard = this->count();
+    int nshard = this->count();
 
     std::vector<distance_t> all_distances(nshard * k * n);
     std::vector<idx_t> all_labels(nshard * k * n);
@@ -199,7 +203,9 @@ void IndexShardsIVF::search(
 
         auto index = dynamic_cast<const IndexIVFInterface*>(indexIn);
 
-        FAISS_THROW_IF_NOT_MSG(index->nprobe == nprobe, "inconsistent nprobe");
+        FAISS_THROW_IF_NOT_MSG(
+                index->nprobe == static_cast<size_t>(nprobe),
+                "inconsistent nprobe");
 
         index->search_preassigned(
                 n,

data/vendor/faiss/faiss/MatrixStats.cpp

@@ -52,7 +52,7 @@ void MatrixStats::PerDimStats::compute_mean_std() {
     if (var < 0) {
         var = 0;
     }
-    stddev = sqrt(var);
+    stddev = std::sqrt(var);
 }
 
 void MatrixStats::do_comment(const char* fmt, ...) {
@@ -67,7 +67,8 @@ void MatrixStats::do_comment(const char* fmt, ...) {
     buf += size;
 }
 
-MatrixStats::MatrixStats(size_t n, size_t d, const float* x) : n(n), d(d) {
+MatrixStats::MatrixStats(size_t n_in, size_t d_in, const float* x)
+        : n(n_in), d(d_in) {
     std::vector<char> comment_buf(10000);
     buf = comment_buf.data();
     nbuf = comment_buf.size();
@@ -162,8 +163,8 @@ MatrixStats::MatrixStats(size_t n, size_t d, const float* x) : n(n), d(d) {
     }
 
     { // norm stats
-        min_norm2 = sqrt(min_norm2);
-        max_norm2 = sqrt(max_norm2);
+        min_norm2 = std::sqrt(min_norm2);
+        max_norm2 = std::sqrt(max_norm2);
         do_comment(
                 "range of L2 norms=[%g, %g] (%zd null vectors)\n",
                 min_norm2,

data/vendor/faiss/faiss/MetaIndexes.cpp

@@ -27,8 +27,8 @@ namespace faiss {
  * IndexSplitVectors implementation
  *******************************************************/
 
-IndexSplitVectors::IndexSplitVectors(idx_t d, bool threaded)
-        : Index(d), own_fields(false), threaded(threaded), sum_d(0) {}
+IndexSplitVectors::IndexSplitVectors(idx_t d_in, bool threaded_in)
+        : Index(d_in), own_fields(false), threaded(threaded_in), sum_d(0) {}
 
 void IndexSplitVectors::add_sub_index(Index* index) {
     sub_indexes.push_back(index);
@@ -44,7 +44,7 @@ void IndexSplitVectors::sync_with_sub_indexes() {
     metric_type = index0->metric_type;
     is_trained = index0->is_trained;
     ntotal = index0->ntotal;
-    for (int i = 1; i < sub_indexes.size(); i++) {
+    for (size_t i = 1; i < sub_indexes.size(); i++) {
         Index* index = sub_indexes[i];
         FAISS_THROW_IF_NOT(metric_type == index->metric_type);
         FAISS_THROW_IF_NOT(ntotal == index->ntotal);
@@ -88,7 +88,7 @@ void IndexSplitVectors::search(
                            n);
                 }
                 const Index* sub_index = index->sub_indexes[no];
-                int64_t sub_d = sub_index->d, d = index->d;
+                int64_t sub_d = sub_index->d;
                 idx_t ofs = 0;
                 for (int i = 0; i < no; i++) {
                     ofs += index->sub_indexes[i]->d;
@@ -107,17 +107,18 @@ void IndexSplitVectors::search(
             };
 
     if (!threaded) {
-        for (int i = 0; i < nshard; i++) {
-            query_func(i);
+        for (int64_t i = 0; i < nshard; i++) {
+            query_func(static_cast<int>(i));
         }
     } else {
         std::vector<std::unique_ptr<WorkerThread>> threads;
         std::vector<std::future<bool>> v;
 
-        for (int i = 0; i < nshard; i++) {
+        for (int64_t i = 0; i < nshard; i++) {
             threads.emplace_back(new WorkerThread());
             WorkerThread* wt = threads.back().get();
-            v.emplace_back(wt->add([i, query_func]() { query_func(i); }));
+            v.emplace_back(wt->add(
+                    [i, query_func]() { query_func(static_cast<int>(i)); }));
         }
 
         // Blocking wait for completion
@@ -155,7 +156,7 @@ void IndexSplitVectors::reset() {
 
 IndexSplitVectors::~IndexSplitVectors() {
     if (own_fields) {
-        for (int s = 0; s < sub_indexes.size(); s++) {
+        for (size_t s = 0; s < sub_indexes.size(); s++) {
             delete sub_indexes[s];
         }
     }
@@ -166,12 +167,12 @@ IndexSplitVectors::~IndexSplitVectors() {
  */
 
 IndexRandom::IndexRandom(
-        idx_t d,
-        idx_t ntotal,
-        int64_t seed,
-        MetricType metric_type)
-        : Index(d, metric_type), seed(seed) {
-    this->ntotal = ntotal;
+        idx_t d_in,
+        idx_t ntotal_in,
+        int64_t seed_in,
+        MetricType metric_type_in)
+        : Index(d_in, metric_type_in), seed(seed_in) {
+    this->ntotal = ntotal_in;
     is_trained = true;
 }
 
@@ -216,7 +217,8 @@ void IndexRandom::search(
                 perm[j] = j;
             }
             for (int j = 0; j < k; j++) {
-                std::swap(perm[j], perm[rng.rand_int(ntotal)]);
+                std::swap(
+                        perm[j], perm[rng.rand_int(static_cast<int>(ntotal))]);
                 I[j] = perm[j];
             }
         }
@@ -234,7 +236,7 @@ void IndexRandom::search(
 
 void IndexRandom::reconstruct(idx_t key, float* recons) const {
     RandomGenerator rng(seed + 123332 + key);
-    for (size_t i = 0; i < d; i++) {
+    for (int i = 0; i < d; i++) {
         recons[i] = rng.rand_float();
     }
 }

data/vendor/faiss/faiss/MetaIndexes.h

@@ -22,7 +22,7 @@ namespace faiss {
  * used to distribute a MultiIndexQuantizer
  */
 struct IndexSplitVectors : Index {
-    bool own_fields;
+    bool own_fields = false;
     bool threaded;
     std::vector<Index*> sub_indexes;
     idx_t sum_d; /// sum of dimensions seen so far

data/vendor/faiss/faiss/MetricType.h

@@ -24,14 +24,15 @@ namespace faiss {
 /// (brute-force) indices supporting additional metric types for vector
 /// comparison.
 ///
-/// NOTE: when adding or removing values, update metric_type_from_int() below.
+/// NOTE: when adding or removing values, update metric_type_from_int()
+///       and metric_type_count() below.
 enum MetricType {
-    METRIC_INNER_PRODUCT = 0, ///< maximum inner product search
-    METRIC_L2 = 1,            ///< squared L2 search
-    METRIC_L1,                ///< L1 (aka cityblock)
-    METRIC_Linf,              ///< infinity distance
-    METRIC_Lp,                ///< L_p distance, p is given by a faiss::Index
-                              /// metric_arg
+    METRIC_INNER_PRODUCT, ///< maximum inner product search
+    METRIC_L2,            ///< squared L2 search
+    METRIC_L1,            ///< L1 (aka cityblock)
+    METRIC_Linf,          ///< infinity distance
+    METRIC_Lp,            ///< L_p distance, p is given by a faiss::Index
+                          /// metric_arg
 
     /// some additional metrics defined in scipy.spatial.distance
     METRIC_Canberra = 20,
@@ -68,6 +69,12 @@ inline MetricType metric_type_from_int(int x) {
     return static_cast<MetricType>(x);
 }
 
+/// Count of entries in the MetricType enum.
+constexpr size_t metric_type_count() {
+    return (METRIC_Lp - METRIC_INNER_PRODUCT) + 1 +
+            (METRIC_GOWER - METRIC_Canberra) + 1;
+}
+
 } // namespace faiss
 
 #endif