faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/scalar_quantizer/scanners.h

@@ -0,0 +1,190 @@
+/*
+ * 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.
+ */
+
+// Private implementation header for per-SIMD scalar quantizer TUs.
+// Do not include in public APIs.
+
+#pragma once
+
+#include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/impl/simdlib/simdlib_dispatch.h>
+#include <faiss/utils/simd_levels.h>
+
+#include <faiss/impl/simd_dispatch.h>
+
+#include <faiss/IndexIVF.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
+#include <faiss/impl/expanded_scanners.h>
+#include <faiss/utils/bf16.h>
+#include <faiss/utils/fp16.h>
+
+#include <faiss/impl/scalar_quantizer/codecs.h>
+#include <faiss/impl/scalar_quantizer/distance_computers.h>
+#include <faiss/impl/scalar_quantizer/quantizers.h>
+#include <faiss/impl/scalar_quantizer/similarities.h>
+
+namespace faiss {
+
+namespace scalar_quantizer {
+
+using QuantizerType = ScalarQuantizer::QuantizerType;
+using SQDistanceComputer = ScalarQuantizer::SQDistanceComputer;
+
+/*******************************************************************
+ * IVFSQScannerIP / IVFSQScannerL2 — moved from anonymous namespace
+ * in ScalarQuantizer.cpp
+ *******************************************************************/
+
+template <class DCClass>
+struct IVFSQScannerIP : InvertedListScanner {
+    DCClass dc;
+    bool by_residual;
+
+    float accu0; /// added to all distances
+
+    IVFSQScannerIP(
+            int d,
+            const std::vector<float>& trained,
+            size_t code_size,
+            bool store_pairs,
+            const IDSelector* sel,
+            bool by_residual)
+            : dc(d, trained), by_residual(by_residual), accu0(0) {
+        this->store_pairs = store_pairs;
+        this->sel = sel;
+        this->code_size = code_size;
+        this->keep_max = true;
+    }
+
+    void set_query(const float* query) override {
+        dc.set_query(query);
+    }
+
+    void set_list(idx_t list_no, float coarse_dis) override {
+        this->list_no = list_no;
+        accu0 = by_residual ? coarse_dis : 0;
+    }
+
+    float distance_to_code(const uint8_t* code) const final {
+        return accu0 + dc.query_to_code(code);
+    }
+
+    size_t scan_codes(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            ResultHandler& handler) const override {
+        return run_scan_codes_fix_C<CMin<float, idx_t>>(
+                *this, list_size, codes, ids, handler);
+    }
+};
+
+template <class DCClass>
+struct IVFSQScannerL2 : InvertedListScanner {
+    DCClass dc;
+
+    bool by_residual;
+    const Index* quantizer;
+    const float* x; /// current query
+
+    std::vector<float> tmp;
+
+    IVFSQScannerL2(
+            int d,
+            const std::vector<float>& trained,
+            size_t code_size,
+            const Index* quantizer,
+            bool store_pairs,
+            const IDSelector* sel,
+            bool by_residual)
+            : dc(d, trained),
+              by_residual(by_residual),
+              quantizer(quantizer),
+              x(nullptr),
+              tmp(d) {
+        this->store_pairs = store_pairs;
+        this->sel = sel;
+        this->code_size = code_size;
+    }
+
+    void set_query(const float* query) override {
+        x = query;
+        if (!quantizer) {
+            dc.set_query(query);
+        }
+    }
+
+    void set_list(idx_t list_no, float /*coarse_dis*/) override {
+        this->list_no = list_no;
+        if (by_residual) {
+            quantizer->compute_residual(x, tmp.data(), list_no);
+            dc.set_query(tmp.data());
+        } else {
+            dc.set_query(x);
+        }
+    }
+
+    float distance_to_code(const uint8_t* code) const final {
+        return dc.query_to_code(code);
+    }
+
+    size_t scan_codes(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            ResultHandler& handler) const override {
+        return run_scan_codes_fix_C<CMax<float, idx_t>>(
+                *this, list_size, codes, ids, handler);
+    }
+};
+
+/*******************************************************************
+ * Forward declaration of inverts list scanner
+ *******************************************************************/
+
+template <SIMDLevel SL>
+InvertedListScanner* sq_select_InvertedListScanner(
+        QuantizerType qtype,
+        MetricType mt,
+        size_t d,
+        size_t code_size,
+        const std::vector<float>& trained,
+        const Index* quantizer,
+        bool store_pairs,
+        const IDSelector* sel,
+        bool by_residual);
+
+/// Scanner for QT_0bit / centroid-only distance: always returns the
+/// coarse distance that was set via set_list().
+struct IVFCoarseDistanceScanner : InvertedListScanner {
+    float coarse_dis = 0;
+
+    IVFCoarseDistanceScanner(
+            bool is_similarity,
+            bool store_pairs,
+            const IDSelector* sel)
+            : InvertedListScanner(store_pairs, sel) {
+        code_size = 0;
+        keep_max = is_similarity;
+    }
+
+    void set_query(const float* /*query_vector*/) override {}
+
+    void set_list(idx_t list_no_in, float coarse_dis_in) override {
+        this->list_no = list_no_in;
+        this->coarse_dis = coarse_dis_in;
+    }
+
+    float distance_to_code(const uint8_t* /*code*/) const override {
+        return coarse_dis;
+    }
+};
+
+} // namespace scalar_quantizer
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/scalar_quantizer/similarities.h

@@ -0,0 +1,94 @@
+/*
+ * 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 <faiss/impl/ScalarQuantizer.h>
+#include <faiss/impl/simdlib/simdlib_dispatch.h>
+#include <faiss/utils/simd_levels.h>
+
+namespace faiss {
+
+namespace scalar_quantizer {
+
+/*******************************************************************
+ * Similarity: gets vector components and computes a similarity wrt. a
+ * query vector stored in the object. The data fields just encapsulate
+ * an accumulator.
+ */
+
+template <SIMDLevel SL>
+struct SimilarityL2 {};
+
+template <>
+struct SimilarityL2<SIMDLevel::NONE> {
+    static constexpr int simdwidth = 1;
+    static constexpr SIMDLevel simd_level = SIMDLevel::NONE;
+    static constexpr MetricType metric_type = METRIC_L2;
+
+    const float *y, *yi;
+
+    explicit SimilarityL2(const float* y_in) : y(y_in), yi(nullptr), accu(0) {}
+
+    /******* scalar accumulator *******/
+
+    float accu;
+
+    FAISS_ALWAYS_INLINE void begin() {
+        accu = 0;
+        yi = y;
+    }
+
+    FAISS_ALWAYS_INLINE void add_component(float x) {
+        float tmp = *yi++ - x;
+        accu += tmp * tmp;
+    }
+
+    FAISS_ALWAYS_INLINE void add_component_2(float x1, float x2) {
+        float tmp = x1 - x2;
+        accu += tmp * tmp;
+    }
+
+    FAISS_ALWAYS_INLINE float result() {
+        return accu;
+    }
+};
+
+template <SIMDLevel SL>
+struct SimilarityIP {};
+
+template <>
+struct SimilarityIP<SIMDLevel::NONE> {
+    static constexpr int simdwidth = 1;
+    static constexpr SIMDLevel simd_level = SIMDLevel::NONE;
+    static constexpr MetricType metric_type = METRIC_INNER_PRODUCT;
+    const float *y, *yi;
+
+    float accu;
+
+    explicit SimilarityIP(const float* y_in) : y(y_in), yi(nullptr), accu(0) {}
+
+    FAISS_ALWAYS_INLINE void begin() {
+        accu = 0;
+        yi = y;
+    }
+
+    FAISS_ALWAYS_INLINE void add_component(float x) {
+        accu += *yi++ * x;
+    }
+
+    FAISS_ALWAYS_INLINE void add_component_2(float x1, float x2) {
+        accu += x1 * x2;
+    }
+
+    FAISS_ALWAYS_INLINE float result() {
+        return accu;
+    }
+};
+
+} // namespace scalar_quantizer
+} // namespace faiss