faiss 0.5.3 → 0.6.0

data/vendor/faiss/faiss/impl/ScalarQuantizer.h

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 #pragma once
 
 #include <faiss/impl/AuxIndexStructures.h>

data/vendor/faiss/faiss/impl/VisitedTable.cpp

@@ -0,0 +1,42 @@
+/*
+ * 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.
+ */
+
+#include <faiss/impl/VisitedTable.h>
+
+#include <cstring>
+
+namespace faiss {
+
+// The vector strategy is faster for get()/set(), but O(size) to initialize.
+// advance() is O(1) except every 250 calls, which are O(size).
+// The hash set strategy is a constant factor slower for get()/set(),
+// but O(1) to construct and O(visits) to advance.
+// A size of ~1M seems to be the threshold where the hash set wins.
+size_t visited_table_hashset_threshold = 500000;
+
+VisitedTable::VisitedTable(size_t size, std::optional<bool> use_hashset)
+        : visno(use_hashset.value_or(size >= visited_table_hashset_threshold)
+                        ? 0
+                        : 1) {
+    if (visno != 0) {
+        visited.resize(size, 0);
+    }
+}
+
+void VisitedTable::advance() {
+    if (visno == 0) {
+        visited_set.clear();
+    } else if (visno < 254) {
+        // 254 rather than 255 because sometimes we use visno and visno+1
+        ++visno;
+    } else {
+        memset(visited.data(), 0, sizeof(visited[0]) * visited.size());
+        visno = 1;
+    }
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/VisitedTable.h

@@ -0,0 +1,69 @@
+/*
+ * 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.
+ */
+
+#ifndef FAISS_VISITED_TABLE_H
+#define FAISS_VISITED_TABLE_H
+
+#include <stdint.h>
+
+#include <optional>
+#include <unordered_set>
+#include <vector>
+
+#include <faiss/impl/platform_macros.h>
+#include <faiss/utils/prefetch.h>
+
+namespace faiss {
+
+FAISS_API extern size_t visited_table_hashset_threshold;
+
+/// A fast, reusable Visited Set for graph search algorithms.
+struct VisitedTable {
+    std::vector<uint8_t> visited;
+    std::unordered_set<size_t> visited_set;
+    uint8_t visno; // 0 if using visited_set, 1..250 if using vector.
+
+    // If use_hashset is nullopt, the use of a hashset will be determined by
+    // size >= visited_table_hashset_threshold.
+    explicit VisitedTable(
+            size_t size,
+            std::optional<bool> use_hashset = std::nullopt);
+
+    /// set flag #no to true, return whether this changed it.
+    bool set(size_t no) {
+        if (visno == 0) {
+            return visited_set.insert(no).second;
+        } else if (visited[no] == visno) {
+            return false;
+        } else {
+            visited[no] = visno;
+            return true;
+        }
+    }
+
+    /// get flag #no
+    bool get(size_t no) const {
+        if (visno == 0) {
+            return visited_set.count(no) != 0;
+        } else {
+            return visited[no] == visno;
+        }
+    }
+
+    void prefetch(size_t no) const {
+        if (visno != 0) {
+            prefetch_L2(&visited[no]);
+        }
+    }
+
+    /// reset all flags to false
+    void advance();
+};
+
+} // namespace faiss
+
+#endif

data/vendor/faiss/faiss/impl/expanded_scanners.h

@@ -0,0 +1,158 @@
+/*
+ * 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 <cstdio>
+
+#include <faiss/IndexIVF.h>
+#include <faiss/impl/ResultHandler.h>
+
+/* This is the inner loop of the inverted list scanners. The default version
+ * that is defined in IndexIVF.cpp works fine but it cannot inline the distance
+ * computation code by calling one or another of the run_scan_codes_* variants
+ * with the exact ScannerType and by setting distance_to_code to be a final
+ * function, the code can be inlined. The speed difference matters for very
+ * small distance computations (eg. SQ or Flat) */
+
+namespace faiss {
+
+namespace {
+
+template <class ScannerType, typename C, bool store_pairs, bool use_sel>
+size_t run_scan_codes1(
+        const ScannerType& scanner,
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) {
+    size_t nup = 0;
+    size_t list_no = scanner.list_no;
+    size_t code_size = scanner.code_size;
+    const IDSelector* sel = scanner.sel;
+    float threshold = handler.threshold;
+    for (size_t j = 0; j < list_size; j++) {
+        if (use_sel) {
+            int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+            // skip code without computing distance
+            if (!sel->is_member(id)) {
+                codes += code_size;
+                continue;
+            }
+        }
+
+        float dis = scanner.distance_to_code(codes); // will be inlined if final
+        if (C::cmp(threshold, dis)) {
+            int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+            handler.add_result(dis, id);
+            threshold = handler.threshold;
+            nup++;
+        }
+        codes += code_size;
+    }
+
+    return nup;
+}
+
+/*****************************************************************************
+ * The following functions dispatch runtime parameters to templates, with
+ * possibly some already-fixed templates.
+ */
+
+template <bool store_pairs, bool use_sel, class ScannerType>
+size_t run_scan_codes_fix_store_pairs_fix_use_sel(
+        const ScannerType& scanner,
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) {
+    if (scanner.keep_max) {
+        return run_scan_codes1<
+                ScannerType,
+                CMin<float, idx_t>,
+                store_pairs,
+                use_sel>(scanner, list_size, codes, ids, handler);
+    } else {
+        return run_scan_codes1<
+                ScannerType,
+                CMax<float, idx_t>,
+                store_pairs,
+                use_sel>(scanner, list_size, codes, ids, handler);
+    }
+}
+
+template <class C, bool use_sel, class ScannerType>
+size_t run_scan_codes_fix_C_fix_use_sel(
+        const ScannerType& scanner,
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) {
+    if (scanner.store_pairs) {
+        return run_scan_codes1<ScannerType, C, true, use_sel>(
+                scanner, list_size, codes, ids, handler);
+    } else {
+        return run_scan_codes1<ScannerType, C, false, use_sel>(
+                scanner, list_size, codes, ids, handler);
+    }
+}
+
+template <class C, class ScannerType>
+size_t run_scan_codes_fix_C(
+        const ScannerType& scanner,
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) {
+    if (scanner.sel) {
+        if (scanner.store_pairs) {
+            return run_scan_codes1<ScannerType, C, true, true>(
+                    scanner, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes1<ScannerType, C, false, true>(
+                    scanner, list_size, codes, ids, handler);
+        }
+    } else {
+        if (scanner.store_pairs) {
+            return run_scan_codes1<ScannerType, C, true, false>(
+                    scanner, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes1<ScannerType, C, false, false>(
+                    scanner, list_size, codes, ids, handler);
+        }
+    }
+}
+
+template <class ScannerType>
+size_t run_scan_codes(
+        const ScannerType& scanner,
+        size_t list_size,
+        const uint8_t* codes,
+        const idx_t* ids,
+        ResultHandler& handler) {
+    if (scanner.sel == nullptr) {
+        if (scanner.store_pairs) {
+            return run_scan_codes_fix_store_pairs_fix_use_sel<true, false>(
+                    scanner, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes_fix_store_pairs_fix_use_sel<false, false>(
+                    scanner, list_size, codes, ids, handler);
+        }
+    } else {
+        if (scanner.store_pairs) {
+            return run_scan_codes_fix_store_pairs_fix_use_sel<true, true>(
+                    scanner, list_size, codes, ids, handler);
+        } else {
+            return run_scan_codes_fix_store_pairs_fix_use_sel<false, true>(
+                    scanner, list_size, codes, ids, handler);
+        }
+    }
+}
+
+} // anonymous namespace
+
+} // namespace faiss