faiss 0.5.2 → 0.6.0

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

@@ -0,0 +1,280 @@
+/*
+ * 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/utils/bf16.h>
+#include <faiss/utils/fp16.h>
+#include <faiss/utils/simd_levels.h>
+#include <faiss/utils/simdlib.h>
+
+namespace faiss {
+
+namespace scalar_quantizer {
+
+using QuantizerType = ScalarQuantizer::QuantizerType;
+
+/*******************************************************************
+ * Quantizer: normalizes scalar vector components, then passes them
+ * through a codec
+ *******************************************************************/
+
+enum class QuantizerTemplateScaling { UNIFORM = 0, NON_UNIFORM = 1 };
+
+template <class Codec, QuantizerTemplateScaling SCALING, SIMDLevel SL>
+struct QuantizerTemplate {};
+
+template <class Codec>
+struct QuantizerTemplate<
+        Codec,
+        QuantizerTemplateScaling::UNIFORM,
+        SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
+    const size_t d;
+    const float vmin, vdiff;
+
+    QuantizerTemplate(size_t d, const std::vector<float>& trained)
+            : d(d), vmin(trained[0]), vdiff(trained[1]) {}
+
+    void encode_vector(const float* x, uint8_t* code) const final {
+        for (size_t i = 0; i < d; i++) {
+            float xi = 0;
+            if (vdiff != 0) {
+                xi = (x[i] - vmin) / vdiff;
+                if (xi < 0) {
+                    xi = 0;
+                }
+                if (xi > 1.0) {
+                    xi = 1.0;
+                }
+            }
+            Codec::encode_component(xi, code, i);
+        }
+    }
+
+    void decode_vector(const uint8_t* code, float* x) const final {
+        for (size_t i = 0; i < d; i++) {
+            float xi = Codec::decode_component(code, i);
+            x[i] = vmin + xi * vdiff;
+        }
+    }
+
+    FAISS_ALWAYS_INLINE float reconstruct_component(
+            const uint8_t* code,
+            size_t i) const {
+        float xi = Codec::decode_component(code, i);
+        return vmin + xi * vdiff;
+    }
+};
+
+template <class Codec>
+struct QuantizerTemplate<
+        Codec,
+        QuantizerTemplateScaling::NON_UNIFORM,
+        SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
+    const size_t d;
+    const float *vmin, *vdiff;
+
+    QuantizerTemplate(size_t d, const std::vector<float>& trained)
+            : d(d), vmin(trained.data()), vdiff(trained.data() + d) {}
+
+    void encode_vector(const float* x, uint8_t* code) const final {
+        for (size_t i = 0; i < d; i++) {
+            float xi = 0;
+            if (vdiff[i] != 0) {
+                xi = (x[i] - vmin[i]) / vdiff[i];
+                if (xi < 0) {
+                    xi = 0;
+                }
+                if (xi > 1.0) {
+                    xi = 1.0;
+                }
+            }
+            Codec::encode_component(xi, code, i);
+        }
+    }
+
+    void decode_vector(const uint8_t* code, float* x) const final {
+        for (size_t i = 0; i < d; i++) {
+            float xi = Codec::decode_component(code, i);
+            x[i] = vmin[i] + xi * vdiff[i];
+        }
+    }
+
+    FAISS_ALWAYS_INLINE float reconstruct_component(
+            const uint8_t* code,
+            size_t i) const {
+        float xi = Codec::decode_component(code, i);
+        return vmin[i] + xi * vdiff[i];
+    }
+};
+
+/*******************************************************************
+ * FP16 quantizer
+ *******************************************************************/
+
+template <SIMDLevel SL>
+struct QuantizerFP16;
+
+template <>
+struct QuantizerFP16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
+    const size_t d;
+
+    QuantizerFP16(size_t d, const std::vector<float>& /* unused */) : d(d) {}
+
+    void encode_vector(const float* x, uint8_t* code) const final {
+        for (size_t i = 0; i < d; i++) {
+            ((uint16_t*)code)[i] = encode_fp16(x[i]);
+        }
+    }
+
+    void decode_vector(const uint8_t* code, float* x) const final {
+        for (size_t i = 0; i < d; i++) {
+            x[i] = decode_fp16(((uint16_t*)code)[i]);
+        }
+    }
+
+    FAISS_ALWAYS_INLINE float reconstruct_component(
+            const uint8_t* code,
+            size_t i) const {
+        return decode_fp16(((uint16_t*)code)[i]);
+    }
+};
+
+template <SIMDLevel SL>
+struct QuantizerFP16 : QuantizerFP16<SIMDLevel::NONE> {
+    using QuantizerFP16<SIMDLevel::NONE>::QuantizerFP16;
+};
+
+/*******************************************************************
+ * BF16 quantizer
+ *******************************************************************/
+
+template <SIMDLevel SL>
+struct QuantizerBF16;
+
+template <>
+struct QuantizerBF16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
+    const size_t d;
+
+    QuantizerBF16(size_t d, const std::vector<float>& /* unused */) : d(d) {}
+
+    void encode_vector(const float* x, uint8_t* code) const final {
+        for (size_t i = 0; i < d; i++) {
+            ((uint16_t*)code)[i] = encode_bf16(x[i]);
+        }
+    }
+
+    void decode_vector(const uint8_t* code, float* x) const final {
+        for (size_t i = 0; i < d; i++) {
+            x[i] = decode_bf16(((uint16_t*)code)[i]);
+        }
+    }
+
+    FAISS_ALWAYS_INLINE float reconstruct_component(
+            const uint8_t* code,
+            size_t i) const {
+        return decode_bf16(((uint16_t*)code)[i]);
+    }
+};
+
+template <SIMDLevel SL>
+struct QuantizerBF16 : QuantizerBF16<SIMDLevel::NONE> {
+    using QuantizerBF16<SIMDLevel::NONE>::QuantizerBF16;
+};
+
+/*******************************************************************
+ * 8bit_direct quantizer
+ *******************************************************************/
+
+template <SIMDLevel SL>
+struct Quantizer8bitDirect;
+
+template <>
+struct Quantizer8bitDirect<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
+    const size_t d;
+
+    Quantizer8bitDirect(size_t d, const std::vector<float>& /* unused */)
+            : d(d) {}
+
+    void encode_vector(const float* x, uint8_t* code) const final {
+        for (size_t i = 0; i < d; i++) {
+            code[i] = (uint8_t)x[i];
+        }
+    }
+
+    void decode_vector(const uint8_t* code, float* x) const final {
+        for (size_t i = 0; i < d; i++) {
+            x[i] = code[i];
+        }
+    }
+
+    FAISS_ALWAYS_INLINE float reconstruct_component(
+            const uint8_t* code,
+            size_t i) const {
+        return code[i];
+    }
+};
+
+template <SIMDLevel SL>
+struct Quantizer8bitDirect : Quantizer8bitDirect<SIMDLevel::NONE> {
+    using Quantizer8bitDirect<SIMDLevel::NONE>::Quantizer8bitDirect;
+};
+
+/*******************************************************************
+ * 8bit_direct_signed quantizer
+ *******************************************************************/
+
+template <SIMDLevel SL>
+struct Quantizer8bitDirectSigned;
+
+template <>
+struct Quantizer8bitDirectSigned<SIMDLevel::NONE>
+        : ScalarQuantizer::SQuantizer {
+    const size_t d;
+
+    Quantizer8bitDirectSigned(size_t d, const std::vector<float>& /* unused */)
+            : d(d) {}
+
+    void encode_vector(const float* x, uint8_t* code) const final {
+        for (size_t i = 0; i < d; i++) {
+            code[i] = (uint8_t)(x[i] + 128);
+        }
+    }
+
+    void decode_vector(const uint8_t* code, float* x) const final {
+        for (size_t i = 0; i < d; i++) {
+            x[i] = code[i] - 128;
+        }
+    }
+
+    FAISS_ALWAYS_INLINE float reconstruct_component(
+            const uint8_t* code,
+            size_t i) const {
+        return code[i] - 128;
+    }
+};
+
+template <SIMDLevel SL>
+struct Quantizer8bitDirectSigned : Quantizer8bitDirectSigned<SIMDLevel::NONE> {
+    using Quantizer8bitDirectSigned<SIMDLevel::NONE>::Quantizer8bitDirectSigned;
+};
+
+/*******************************************************************
+ * Selection function
+ *******************************************************************/
+
+// declare for all levels
+template <SIMDLevel SL>
+ScalarQuantizer::SQuantizer* sq_select_quantizer(
+        QuantizerType qtype,
+        size_t d,
+        const std::vector<float>& trained);
+
+} // namespace scalar_quantizer
+
+} // namespace faiss

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

@@ -0,0 +1,164 @@
+/*
+ * 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/utils/simd_levels.h>
+#include <faiss/utils/simdlib.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);
+
+} // 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/utils/simd_levels.h>
+#include <faiss/utils/simdlib.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) : y(y), 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) : y(y), 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