faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.h

@@ -11,7 +11,7 @@
 #include <vector>
 
 #include <faiss/Index.h>
-#include <faiss/utils/approx_topk/mode.h>
+#include <faiss/impl/approx_topk/approx_topk.h>
 
 namespace faiss {
 

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

@@ -0,0 +1,121 @@
+/*
+ * 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 {
+
+/*******************************************************************
+ * Codec: converts between values in [0, 1] and an index in a code
+ * array. The "i" parameter is the vector component index (not byte
+ * index).
+ */
+
+template <SIMDLevel SL>
+struct Codec8bit {};
+
+template <>
+struct Codec8bit<SIMDLevel::NONE> {
+    static FAISS_ALWAYS_INLINE void encode_component(
+            float x,
+            uint8_t* code,
+            size_t i) {
+        code[i] = (int)(255 * x);
+    }
+
+    static FAISS_ALWAYS_INLINE float decode_component(
+            const uint8_t* code,
+            size_t i) {
+        return (code[i] + 0.5f) / 255.0f;
+    }
+};
+
+template <SIMDLevel SL>
+struct Codec4bit {};
+
+template <>
+struct Codec4bit<SIMDLevel::NONE> {
+    static FAISS_ALWAYS_INLINE void encode_component(
+            float x,
+            uint8_t* code,
+            size_t i) {
+        code[i / 2] |= (int)(x * 15.0) << ((i & 1) << 2);
+    }
+
+    static FAISS_ALWAYS_INLINE float decode_component(
+            const uint8_t* code,
+            size_t i) {
+        return (((code[i / 2] >> ((i & 1) << 2)) & 0xf) + 0.5f) / 15.0f;
+    }
+};
+
+template <SIMDLevel SL>
+struct Codec6bit {};
+
+template <>
+struct Codec6bit<SIMDLevel::NONE> {
+    static FAISS_ALWAYS_INLINE void encode_component(
+            float x,
+            uint8_t* code,
+            size_t i) {
+        int bits = (int)(x * 63.0);
+        code += (i >> 2) * 3;
+        switch (i & 3) {
+            case 0:
+                code[0] |= bits;
+                break;
+            case 1:
+                code[0] |= bits << 6;
+                code[1] |= bits >> 2;
+                break;
+            case 2:
+                code[1] |= bits << 4;
+                code[2] |= bits >> 4;
+                break;
+            case 3:
+                code[2] |= bits << 2;
+                break;
+            default:
+                break;
+        }
+    }
+
+    static FAISS_ALWAYS_INLINE float decode_component(
+            const uint8_t* code,
+            size_t i) {
+        uint8_t bits = 0;
+        code += (i >> 2) * 3;
+        switch (i & 3) {
+            case 0:
+                bits = code[0] & 0x3f;
+                break;
+            case 1:
+                bits = code[0] >> 6;
+                bits |= (code[1] & 0xf) << 2;
+                break;
+            case 2:
+                bits = code[1] >> 4;
+                bits |= (code[2] & 3) << 4;
+                break;
+            case 3:
+                bits = code[2] >> 2;
+                break;
+            default:
+                break;
+        }
+        return (bits + 0.5f) / 63.0f;
+    }
+};
+
+} // namespace scalar_quantizer
+} // namespace faiss

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

@@ -0,0 +1,137 @@
+/*
+ * 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 {
+
+using SQDistanceComputer = ScalarQuantizer::SQDistanceComputer;
+
+/*******************************************************************
+ * DistanceComputer: combines a similarity and a quantizer to do
+ * code-to-vector or code-to-code comparisons
+ *******************************************************************/
+
+template <class Quantizer, class Similarity, SIMDLevel SL>
+struct DCTemplate : SQDistanceComputer {};
+
+template <class Quantizer, class Similarity>
+struct DCTemplate<Quantizer, Similarity, SIMDLevel::NONE> : SQDistanceComputer {
+    using Sim = Similarity;
+
+    Quantizer quant;
+
+    DCTemplate(size_t d, const std::vector<float>& trained)
+            : quant(d, trained) {}
+
+    float compute_distance(const float* x, const uint8_t* code) const {
+        Similarity sim(x);
+        sim.begin();
+        for (size_t i = 0; i < quant.d; i++) {
+            float xi = quant.reconstruct_component(code, i);
+            sim.add_component(xi);
+        }
+        return sim.result();
+    }
+
+    float compute_code_distance(const uint8_t* code1, const uint8_t* code2)
+            const {
+        Similarity sim(nullptr);
+        sim.begin();
+        for (size_t i = 0; i < quant.d; i++) {
+            float x1 = quant.reconstruct_component(code1, i);
+            float x2 = quant.reconstruct_component(code2, i);
+            sim.add_component_2(x1, x2);
+        }
+        return sim.result();
+    }
+
+    void set_query(const float* x) final {
+        q = x;
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return compute_code_distance(
+                codes + i * code_size, codes + j * code_size);
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        return compute_distance(q, code);
+    }
+};
+
+/*******************************************************************
+ * DistanceComputerByte: computes distances in the integer domain
+ *******************************************************************/
+
+template <class Similarity, SIMDLevel SL>
+struct DistanceComputerByte : SQDistanceComputer {};
+
+template <class Similarity>
+struct DistanceComputerByte<Similarity, SIMDLevel::NONE> : SQDistanceComputer {
+    using Sim = Similarity;
+
+    int d;
+    std::vector<uint8_t> tmp;
+
+    DistanceComputerByte(int d_in, const std::vector<float>&)
+            : d(d_in), tmp(d_in) {}
+
+    int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
+            const {
+        int accu = 0;
+        for (int i = 0; i < d; i++) {
+            if (Sim::metric_type == METRIC_INNER_PRODUCT) {
+                accu += int(code1[i]) * code2[i];
+            } else {
+                int diff = int(code1[i]) - code2[i];
+                accu += diff * diff;
+            }
+        }
+        return accu;
+    }
+
+    void set_query(const float* x) final {
+        for (int i = 0; i < d; i++) {
+            tmp[i] = int(x[i]);
+        }
+    }
+
+    int compute_distance(const float* x, const uint8_t* code) {
+        set_query(x);
+        return compute_code_distance(tmp.data(), code);
+    }
+
+    float symmetric_dis(idx_t i, idx_t j) override {
+        return compute_code_distance(
+                codes + i * code_size, codes + j * code_size);
+    }
+
+    float query_to_code(const uint8_t* code) const final {
+        return compute_code_distance(tmp.data(), code);
+    }
+};
+
+/*******************************************************************
+ * Selection function
+ *******************************************************************/
+
+template <SIMDLevel SL>
+SQDistanceComputer* sq_select_distance_computer(
+        MetricType metric,
+        ScalarQuantizer::QuantizerType qtype,
+        size_t d,
+        const std::vector<float>& trained);
+
+} // namespace scalar_quantizer
+} // namespace faiss

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

@@ -0,0 +1,371 @@
+/*
+ * 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 <algorithm>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/impl/simdlib/simdlib_dispatch.h>
+#include <faiss/utils/bf16.h>
+#include <faiss/utils/fp16.h>
+#include <faiss/utils/simd_levels.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_in, const std::vector<float>& trained)
+            : d(d_in), 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_in, const std::vector<float>& trained)
+            : d(d_in), vmin(trained.data()), vdiff(trained.data() + d_in) {}
+
+    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];
+    }
+};
+
+/*******************************************************************
+ * TurboQuant MSE quantizer
+ *******************************************************************/
+template <int NBits, SIMDLevel SL>
+struct QuantizerTurboQuantMSE;
+
+template <int NBits>
+struct QuantizerTurboQuantMSE<NBits, SIMDLevel::NONE>
+        : ScalarQuantizer::SQuantizer {
+    static_assert(NBits >= 1 && NBits <= 8);
+
+    static constexpr size_t kCentroidsCount = size_t(1) << NBits;
+    static constexpr uint16_t kIndexMask =
+            static_cast<uint16_t>((1u << NBits) - 1);
+
+    const size_t d;
+    const float* centroids;
+    const float* boundaries;
+
+    QuantizerTurboQuantMSE(size_t d_in, const std::vector<float>& trained)
+            : d(d_in), centroids(nullptr), boundaries(nullptr) {
+        FAISS_THROW_IF_NOT(trained.size() == 2 * kCentroidsCount - 1);
+        centroids = trained.data();
+        boundaries = trained.data() + kCentroidsCount;
+    }
+
+    FAISS_ALWAYS_INLINE uint8_t select_index(float x) const {
+        return static_cast<uint8_t>(
+                std::upper_bound(
+                        boundaries, boundaries + (kCentroidsCount - 1), x) -
+                boundaries);
+    }
+
+    FAISS_ALWAYS_INLINE void encode_index(uint8_t idx, uint8_t* code, size_t i)
+            const {
+        const size_t bit_offset = i * NBits;
+        const size_t byte_offset = bit_offset >> 3;
+        const size_t bit_shift = bit_offset & 7;
+        const uint16_t packed = static_cast<uint16_t>(idx & kIndexMask)
+                << bit_shift;
+        code[byte_offset] |= packed & 0xff;
+        if (bit_shift + NBits > 8) {
+            code[byte_offset + 1] |= packed >> 8;
+        }
+    }
+
+    FAISS_ALWAYS_INLINE uint8_t
+    decode_index(const uint8_t* code, size_t i) const {
+        const size_t bit_offset = i * NBits;
+        const size_t byte_offset = bit_offset >> 3;
+        const size_t bit_shift = bit_offset & 7;
+
+        uint16_t packed = code[byte_offset];
+        if (bit_shift + NBits > 8) {
+            packed |= static_cast<uint16_t>(code[byte_offset + 1]) << 8;
+        }
+        return static_cast<uint8_t>((packed >> bit_shift) & kIndexMask);
+    }
+
+    void encode_vector(const float* x, uint8_t* code) const final {
+        for (size_t i = 0; i < d; i++) {
+            encode_index(select_index(x[i]), code, i);
+        }
+    }
+
+    void decode_vector(const uint8_t* code, float* x) const final {
+        for (size_t i = 0; i < d; i++) {
+            x[i] = centroids[decode_index(code, i)];
+        }
+    }
+
+    FAISS_ALWAYS_INLINE float reconstruct_component(
+            const uint8_t* code,
+            size_t i) const {
+        return centroids[decode_index(code, i)];
+    }
+};
+
+template <int NBits, SIMDLevel SL>
+struct QuantizerTurboQuantMSE : QuantizerTurboQuantMSE<NBits, SIMDLevel::NONE> {
+    using QuantizerTurboQuantMSE<NBits, SIMDLevel::NONE>::
+            QuantizerTurboQuantMSE;
+};
+
+/*******************************************************************
+ * FP16 quantizer
+ *******************************************************************/
+
+template <SIMDLevel SL>
+struct QuantizerFP16;
+
+template <>
+struct QuantizerFP16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
+    const size_t d;
+
+    QuantizerFP16(size_t d_in, const std::vector<float>& /* unused */)
+            : d(d_in) {}
+
+    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_in, const std::vector<float>& /* unused */)
+            : d(d_in) {}
+
+    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_in, const std::vector<float>& /* unused */)
+            : d(d_in) {}
+
+    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_in,
+            const std::vector<float>& /* unused */)
+            : d(d_in) {}
+
+    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