faiss 0.5.2 → 0.6.0

data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp

@@ -8,13 +8,11 @@
 #include <faiss/IndexAdditiveQuantizer.h>
 
 #include <algorithm>
-#include <cmath>
 #include <cstring>
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResidualQuantizer.h>
 #include <faiss/impl/ResultHandler.h>
-#include <faiss/utils/distances.h>
 #include <faiss/utils/extra_distances.h>
 
 namespace faiss {
@@ -189,17 +187,14 @@ void search_with_LUT(
 FlatCodesDistanceComputer* IndexAdditiveQuantizer::
         get_FlatCodesDistanceComputer() const {
     if (aq->search_type == AdditiveQuantizer::ST_decompress) {
-        if (metric_type == METRIC_L2) {
-            using VD = VectorDistance<METRIC_L2>;
-            VD vd = {size_t(d), metric_arg};
-            return new AQDistanceComputerDecompress<VD>(*this, vd);
-        } else if (metric_type == METRIC_INNER_PRODUCT) {
-            using VD = VectorDistance<METRIC_INNER_PRODUCT>;
-            VD vd = {size_t(d), metric_arg};
-            return new AQDistanceComputerDecompress<VD>(*this, vd);
-        } else {
-            FAISS_THROW_MSG("unsupported metric");
-        }
+        return with_VectorDistance(
+                d,
+                metric_type,
+                metric_arg,
+                [&](auto vd) -> FlatCodesDistanceComputer* {
+                    return new AQDistanceComputerDecompress<decltype(vd)>(
+                            *this, vd);
+                });
     } else {
         if (metric_type == METRIC_INNER_PRODUCT) {
             return new AQDistanceComputerLUT<
@@ -242,17 +237,17 @@ void IndexAdditiveQuantizer::search(
             !params, "search params not supported for this index");
 
     if (aq->search_type == AdditiveQuantizer::ST_decompress) {
-        if (metric_type == METRIC_L2) {
-            using VD = VectorDistance<METRIC_L2>;
-            VD vd = {size_t(d), metric_arg};
-            HeapBlockResultHandler<VD::C> rh(n, distances, labels, k);
-            search_with_decompress(*this, x, vd, rh);
-        } else if (metric_type == METRIC_INNER_PRODUCT) {
-            using VD = VectorDistance<METRIC_INNER_PRODUCT>;
-            VD vd = {size_t(d), metric_arg};
-            HeapBlockResultHandler<VD::C> rh(n, distances, labels, k);
-            search_with_decompress(*this, x, vd, rh);
-        }
+        with_VectorDistance(d, metric_type, metric_arg, [&](auto vd) {
+            if constexpr (decltype(vd)::is_similarity) {
+                HeapBlockResultHandler<CMin<float, idx_t>> rh(
+                        n, distances, labels, k);
+                search_with_decompress(*this, x, vd, rh);
+            } else {
+                HeapBlockResultHandler<CMax<float, idx_t>> rh(
+                        n, distances, labels, k);
+                search_with_decompress(*this, x, vd, rh);
+            }
+        });
     } else {
         if (metric_type == METRIC_INNER_PRODUCT) {
             HeapBlockResultHandler<CMin<float, idx_t>> rh(

data/vendor/faiss/faiss/IndexBinary.cpp

@@ -12,6 +12,7 @@
 
 #include <cinttypes>
 #include <cstring>
+#include <typeinfo>
 
 namespace faiss {
 

data/vendor/faiss/faiss/IndexBinaryHNSW.cpp

@@ -22,6 +22,7 @@
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/VisitedTable.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/random.h>
@@ -205,10 +206,14 @@ void IndexBinaryHNSW::search(
         idx_t k,
         int32_t* distances,
         idx_t* labels,
-        const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+        const SearchParameters* params_in) const {
     FAISS_THROW_IF_NOT(k > 0);
+    const SearchParametersHNSW* params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const SearchParametersHNSW*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(
+                params, "IndexBinaryHNSW params have incorrect type");
+    }
 
     // we use the buffer for distances as float but convert them back
     // to int in the end
@@ -231,7 +236,7 @@ void IndexBinaryHNSW::search(
             // as the index parameter. This state does not get used in the
             // search function, as it is merely there to to enable Panorama
             // execution for IndexHNSWFlatPanorama.
-            hnsw.search(*dis, nullptr, res, vt);
+            hnsw.search(*dis, nullptr, res, vt, params_in);
             res.end();
         }
     }

data/vendor/faiss/faiss/IndexBinaryIVF.cpp

@@ -14,6 +14,7 @@
 #include <cstdio>
 
 #include <algorithm>
+#include <limits>
 #include <memory>
 
 #include <faiss/IndexFlat.h>
@@ -120,25 +121,46 @@ void IndexBinaryIVF::search(
         idx_t k,
         int32_t* distances,
         idx_t* labels,
-        const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+        const SearchParameters* params_in) const {
     FAISS_THROW_IF_NOT(k > 0);
+    const IVFSearchParameters* params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(
+                params, "IndexBinaryIVF params have incorrect type");
+        FAISS_THROW_IF_MSG(
+                params->sel, "IDSelector is not supported for IndexBinaryIVF");
+    }
+    const size_t nprobe =
+            std::min(nlist, params ? params->nprobe : this->nprobe);
     FAISS_THROW_IF_NOT(nprobe > 0);
 
-    const size_t nprobe_2 = std::min(nlist, this->nprobe);
-    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe_2]);
-    std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe_2]);
+    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
+    std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);
 
     double t0 = getmillisecs();
-    quantizer->search(n, x, nprobe_2, coarse_dis.get(), idx.get());
+    quantizer->search(
+            n,
+            x,
+            nprobe,
+            coarse_dis.get(),
+            idx.get(),
+            params ? params->quantizer_params : nullptr);
     indexIVF_stats.quantization_time += getmillisecs() - t0;
 
     t0 = getmillisecs();
-    invlists->prefetch_lists(idx.get(), n * nprobe_2);
+    invlists->prefetch_lists(idx.get(), n * nprobe);
 
     search_preassigned(
-            n, x, k, idx.get(), coarse_dis.get(), distances, labels, false);
+            n,
+            x,
+            k,
+            idx.get(),
+            coarse_dis.get(),
+            distances,
+            labels,
+            false,
+            params);
     indexIVF_stats.search_time += getmillisecs() - t0;
 }
 
@@ -389,6 +411,10 @@ void search_knn_hamming_heap(
     idx_t nprobe = params ? params->nprobe : ivf->nprobe;
     nprobe = std::min((idx_t)ivf->nlist, nprobe);
     idx_t max_codes = params ? params->max_codes : ivf->max_codes;
+    const idx_t unlimited_list_size = std::numeric_limits<idx_t>::max();
+    if (max_codes == 0) {
+        max_codes = unlimited_list_size;
+    }
     MetricType metric_type = ivf->metric_type;
 
     // almost verbatim copy from IndexIVF::search_preassigned
@@ -437,6 +463,10 @@ void search_knn_hamming_heap(
                 nlistv++;
 
                 size_t list_size = ivf->invlists->list_size(key);
+                size_t list_size_max = max_codes - nscan;
+                if (list_size > list_size_max) {
+                    list_size = list_size_max;
+                }
                 InvertedLists::ScopedCodes scodes(ivf->invlists, key);
                 std::unique_ptr<InvertedLists::ScopedIds> sids;
                 const idx_t* ids = nullptr;
@@ -451,7 +481,7 @@ void search_knn_hamming_heap(
                         list_size, scodes.get(), ids, simi, idxi, k);
 
                 nscan += list_size;
-                if (max_codes && nscan >= max_codes) {
+                if (nscan >= max_codes) {
                     break;
                 }
             }
@@ -525,6 +555,10 @@ void search_knn_hamming_count(
 
             nlistv++;
             size_t list_size = ivf->invlists->list_size(key);
+            size_t list_size_max = max_codes - nscan;
+            if (list_size > list_size_max) {
+                list_size = list_size_max;
+            }
             InvertedLists::ScopedCodes scodes(ivf->invlists, key);
             const uint8_t* list_vecs = scodes.get();
             const idx_t* ids =
@@ -541,7 +575,7 @@ void search_knn_hamming_count(
             }
 
             nscan += list_size;
-            if (max_codes && nscan >= max_codes) {
+            if (nscan >= max_codes) {
                 break;
             }
         }

data/vendor/faiss/faiss/IndexFastScan.cpp

@@ -20,11 +20,8 @@
 #include <faiss/impl/pq4_fast_scan.h>
 #include <faiss/impl/simd_result_handlers.h>
 #include <faiss/utils/hamming.h>
-#include <faiss/utils/utils.h>
-
-#include <faiss/impl/pq4_fast_scan.h>
-#include <faiss/impl/simd_result_handlers.h>
 #include <faiss/utils/quantize_lut.h>
+#include <faiss/utils/utils.h>
 
 namespace faiss {
 
@@ -84,7 +81,8 @@ void IndexFastScan::add(idx_t n, const float* x) {
     compute_codes(tmp_codes.get(), n, x);
 
     ntotal2 = roundup(ntotal + n, bbs);
-    size_t new_size = ntotal2 * M2 / 2; // assume nbits = 4
+    size_t n_blocks = ntotal2 / bbs;
+    size_t new_size = n_blocks * get_block_stride();
     size_t old_size = codes.size();
     if (new_size > old_size) {
         codes.resize(new_size);
@@ -92,7 +90,15 @@ void IndexFastScan::add(idx_t n, const float* x) {
     }
 
     pq4_pack_codes_range(
-            tmp_codes.get(), M, ntotal, ntotal + n, bbs, M2, codes.get());
+            tmp_codes.get(),
+            M,
+            ntotal,
+            ntotal + n,
+            bbs,
+            M2,
+            codes.get(),
+            0,
+            get_block_stride());
 
     ntotal += n;
 }
@@ -101,17 +107,25 @@ CodePacker* IndexFastScan::get_CodePacker() const {
     return new CodePackerPQ4(M, bbs);
 }
 
+size_t IndexFastScan::get_block_stride() const {
+    std::unique_ptr<CodePacker> packer(get_CodePacker());
+    FAISS_THROW_IF_NOT_MSG(
+            packer->nvec == static_cast<size_t>(bbs),
+            "CodePacker must pack bbs vectors per block for fast-scan");
+    return packer->block_size;
+}
+
 size_t IndexFastScan::remove_ids(const IDSelector& sel) {
     idx_t j = 0;
     std::vector<uint8_t> buffer(code_size);
-    CodePackerPQ4 packer(M, bbs);
+    std::unique_ptr<CodePacker> packer(get_CodePacker());
     for (idx_t i = 0; i < ntotal; i++) {
         if (sel.is_member(i)) {
             // should be removed
         } else {
             if (i > j) {
-                packer.unpack_1(codes.data(), i, buffer.data());
-                packer.pack_1(buffer.data(), j, codes.data());
+                packer->unpack_1(codes.data(), i, buffer.data());
+                packer->pack_1(buffer.data(), j, codes.data());
             }
             j++;
         }
@@ -120,8 +134,7 @@ size_t IndexFastScan::remove_ids(const IDSelector& sel) {
     if (nremove > 0) {
         ntotal = j;
         ntotal2 = roundup(ntotal, bbs);
-        size_t new_size = ntotal2 * M2 / 2;
-        codes.resize(new_size);
+        codes.resize(ntotal2 / bbs * get_block_stride());
     }
     return nremove;
 }
@@ -143,13 +156,14 @@ void IndexFastScan::merge_from(Index& otherIndex, idx_t add_id) {
     check_compatible_for_merge(otherIndex);
     IndexFastScan* other = static_cast<IndexFastScan*>(&otherIndex);
     ntotal2 = roundup(ntotal + other->ntotal, bbs);
-    codes.resize(ntotal2 * M2 / 2);
+    codes.resize(ntotal2 / bbs * get_block_stride());
     std::vector<uint8_t> buffer(code_size);
-    CodePackerPQ4 packer(M, bbs);
+    std::unique_ptr<CodePacker> packer(get_CodePacker());
+    std::unique_ptr<CodePacker> other_packer(other->get_CodePacker());
 
     for (int i = 0; i < other->ntotal; i++) {
-        packer.unpack_1(other->codes.data(), i, buffer.data());
-        packer.pack_1(buffer.data(), ntotal + i, codes.data());
+        other_packer->unpack_1(other->codes.data(), i, buffer.data());
+        packer->pack_1(buffer.data(), ntotal + i, codes.data());
     }
     ntotal += other->ntotal;
     other->reset();
@@ -531,7 +545,8 @@ void IndexFastScan::search_implem_12(
                 codes.get(),
                 LUT.get(),
                 *handler.get(),
-                context.norm_scaler);
+                context.norm_scaler,
+                get_block_stride());
     }
     if (!(skip & 8)) {
         handler->end();
@@ -614,7 +629,8 @@ void IndexFastScan::search_implem_14(
                 codes.get(),
                 LUT.get(),
                 *handler.get(),
-                context.norm_scaler);
+                context.norm_scaler,
+                get_block_stride());
     }
     if (!(skip & 8)) {
         handler->end();
@@ -639,11 +655,8 @@ template void IndexFastScan::search_dispatch_implem<false>(
 
 void IndexFastScan::reconstruct(idx_t key, float* recons) const {
     std::vector<uint8_t> code(code_size, 0);
-    BitstringWriter bsw(code.data(), code_size);
-    for (size_t m = 0; m < M; m++) {
-        uint8_t c = pq4_get_packed_element(codes.data(), bbs, M2, key, m);
-        bsw.write(c, nbits);
-    }
+    std::unique_ptr<CodePacker> packer(get_CodePacker());
+    packer->unpack_1(codes.data(), key, code.data());
     sa_decode(1, code.data(), recons);
 }
 

data/vendor/faiss/faiss/IndexFastScan.h

@@ -214,7 +214,16 @@ struct IndexFastScan : Index {
      *
      * @return  pointer to the code packer
      */
-    CodePacker* get_CodePacker() const;
+    virtual CodePacker* get_CodePacker() const;
+
+    /** Get stride in bytes between consecutive SIMD blocks.
+     *
+     * Derived from get_CodePacker()->block_size so that there is a
+     * single source of truth for the block layout.
+     *
+     * @return stride in bytes
+     */
+    size_t get_block_stride() const;
 
     /** Merge another index into this one
      *