faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/scalar_quantizer/sq-dispatch.h

@@ -0,0 +1,388 @@
+/*
+ * 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/scalar_quantizer/codecs.h>
+#include <faiss/impl/scalar_quantizer/distance_computers.h>
+#include <faiss/impl/scalar_quantizer/quantizers.h>
+#include <faiss/impl/scalar_quantizer/scanners.h>
+#include <faiss/impl/scalar_quantizer/similarities.h>
+
+#ifndef THE_LEVEL_TO_DISPATCH
+#error "THE_LEVEL_TO_DISPATCH should be set on input to this header"
+#endif
+
+namespace faiss {
+
+namespace scalar_quantizer {
+
+// Define SL as alias for THE_LEVEL_TO_DISPATCH for use in this file
+constexpr SIMDLevel SL = THE_LEVEL_TO_DISPATCH;
+
+// Returns true if dimension d is compatible with the given SIMD level
+template <SIMDLevel SL2>
+constexpr bool is_dimension_compatible(size_t d) {
+    if constexpr (SL2 == SIMDLevel::AVX512) {
+        return d % 16 == 0;
+    } else if constexpr (SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
+        return d % 8 == 0;
+    } else {
+        return true; // SIMDLevel::NONE has no alignment requirements
+    }
+}
+
+/*******************************************************************
+ * sq_select_quantizer: the big switch returning SQuantizer*
+ *******************************************************************/
+
+template <>
+ScalarQuantizer::SQuantizer* sq_select_quantizer<THE_LEVEL_TO_DISPATCH>(
+        QuantizerType qtype,
+        size_t d,
+        const std::vector<float>& trained) {
+    // Return nullptr for incompatible dimensions in SIMD cases
+    if constexpr (SL != SIMDLevel::NONE) {
+        if (!is_dimension_compatible<SL>(d)) {
+            return nullptr;
+        }
+    }
+    switch (qtype) {
+        case ScalarQuantizer::QT_8bit:
+            return new QuantizerTemplate<
+                    Codec8bit<SL>,
+                    QuantizerTemplateScaling::NON_UNIFORM,
+                    SL>(d, trained);
+        case ScalarQuantizer::QT_6bit:
+            return new QuantizerTemplate<
+                    Codec6bit<SL>,
+                    QuantizerTemplateScaling::NON_UNIFORM,
+                    SL>(d, trained);
+        case ScalarQuantizer::QT_4bit:
+            return new QuantizerTemplate<
+                    Codec4bit<SL>,
+                    QuantizerTemplateScaling::NON_UNIFORM,
+                    SL>(d, trained);
+        case ScalarQuantizer::QT_8bit_uniform:
+            return new QuantizerTemplate<
+                    Codec8bit<SL>,
+                    QuantizerTemplateScaling::UNIFORM,
+                    SL>(d, trained);
+        case ScalarQuantizer::QT_4bit_uniform:
+            return new QuantizerTemplate<
+                    Codec4bit<SL>,
+                    QuantizerTemplateScaling::UNIFORM,
+                    SL>(d, trained);
+        case ScalarQuantizer::QT_fp16:
+            return new QuantizerFP16<SL>(d, trained);
+        case ScalarQuantizer::QT_bf16:
+            return new QuantizerBF16<SL>(d, trained);
+        case ScalarQuantizer::QT_8bit_direct:
+            return new Quantizer8bitDirect<SL>(d, trained);
+        case ScalarQuantizer::QT_8bit_direct_signed:
+            return new Quantizer8bitDirectSigned<SL>(d, trained);
+        case ScalarQuantizer::QT_0bit:
+            FAISS_THROW_MSG(
+                    "QT_0bit does not support standalone quantization, use IndexIVFScalarQuantizer");
+        case ScalarQuantizer::QT_1bit_tqmse:
+            return new QuantizerTurboQuantMSE<1, SL>(d, trained);
+        case ScalarQuantizer::QT_2bit_tqmse:
+            return new QuantizerTurboQuantMSE<2, SL>(d, trained);
+        case ScalarQuantizer::QT_3bit_tqmse:
+            return new QuantizerTurboQuantMSE<3, SL>(d, trained);
+        case ScalarQuantizer::QT_4bit_tqmse:
+            return new QuantizerTurboQuantMSE<4, SL>(d, trained);
+        case ScalarQuantizer::QT_8bit_tqmse:
+            return new QuantizerTurboQuantMSE<8, SL>(d, trained);
+        default:
+            FAISS_THROW_MSG("unknown qtype");
+    }
+}
+
+/*******************************************************************
+ * select_distance_computer_body: helper for sq_select_distance_computer
+ *******************************************************************/
+
+template <class Sim, SIMDLevel SL2>
+SQDistanceComputer* select_distance_computer_body(
+        ScalarQuantizer::QuantizerType qtype,
+        size_t d,
+        const std::vector<float>& trained) {
+    // Return nullptr for incompatible dimensions in SIMD cases
+    if constexpr (SL2 != SIMDLevel::NONE) {
+        if (!is_dimension_compatible<SL2>(d)) {
+            return nullptr;
+        }
+    }
+    switch (qtype) {
+        case ScalarQuantizer::QT_8bit_uniform:
+            return new DCTemplate<
+                    QuantizerTemplate<
+                            Codec8bit<SL2>,
+                            QuantizerTemplateScaling::UNIFORM,
+                            SL2>,
+                    Sim,
+                    SL2>(d, trained);
+
+        case ScalarQuantizer::QT_4bit_uniform:
+            return new DCTemplate<
+                    QuantizerTemplate<
+                            Codec4bit<SL2>,
+                            QuantizerTemplateScaling::UNIFORM,
+                            SL2>,
+                    Sim,
+                    SL2>(d, trained);
+
+        case ScalarQuantizer::QT_8bit:
+            return new DCTemplate<
+                    QuantizerTemplate<
+                            Codec8bit<SL2>,
+                            QuantizerTemplateScaling::NON_UNIFORM,
+                            SL2>,
+                    Sim,
+                    SL2>(d, trained);
+
+        case ScalarQuantizer::QT_6bit:
+            return new DCTemplate<
+                    QuantizerTemplate<
+                            Codec6bit<SL2>,
+                            QuantizerTemplateScaling::NON_UNIFORM,
+                            SL2>,
+                    Sim,
+                    SL2>(d, trained);
+
+        case ScalarQuantizer::QT_4bit:
+            return new DCTemplate<
+                    QuantizerTemplate<
+                            Codec4bit<SL2>,
+                            QuantizerTemplateScaling::NON_UNIFORM,
+                            SL2>,
+                    Sim,
+                    SL2>(d, trained);
+
+        case ScalarQuantizer::QT_fp16:
+            return new DCTemplate<QuantizerFP16<SL2>, Sim, SL2>(d, trained);
+
+        case ScalarQuantizer::QT_bf16:
+            return new DCTemplate<QuantizerBF16<SL2>, Sim, SL2>(d, trained);
+
+        case ScalarQuantizer::QT_8bit_direct:
+            if constexpr (SL2 == SIMDLevel::AVX512) {
+                if (d % 32 == 0) {
+                    return new DistanceComputerByte<Sim, SL2>(
+                            static_cast<int>(d), trained);
+                }
+            } else if constexpr (SL2 == SIMDLevel::AVX2) {
+                if (d % 16 == 0) {
+                    return new DistanceComputerByte<Sim, SL2>(
+                            static_cast<int>(d), trained);
+                }
+            }
+            return new DCTemplate<Quantizer8bitDirect<SL2>, Sim, SL2>(
+                    d, trained);
+
+        case ScalarQuantizer::QT_8bit_direct_signed:
+            return new DCTemplate<Quantizer8bitDirectSigned<SL2>, Sim, SL2>(
+                    d, trained);
+        case ScalarQuantizer::QT_0bit:
+            FAISS_THROW_MSG(
+                    "QT_0bit does not support standalone distance computation, use IndexIVFScalarQuantizer");
+        case ScalarQuantizer::QT_1bit_tqmse:
+            return new DCTemplate<QuantizerTurboQuantMSE<1, SL2>, Sim, SL2>(
+                    d, trained);
+        case ScalarQuantizer::QT_2bit_tqmse:
+            return new DCTemplate<QuantizerTurboQuantMSE<2, SL2>, Sim, SL2>(
+                    d, trained);
+        case ScalarQuantizer::QT_3bit_tqmse:
+            return new DCTemplate<QuantizerTurboQuantMSE<3, SL2>, Sim, SL2>(
+                    d, trained);
+        case ScalarQuantizer::QT_4bit_tqmse:
+            return new DCTemplate<QuantizerTurboQuantMSE<4, SL2>, Sim, SL2>(
+                    d, trained);
+        case ScalarQuantizer::QT_8bit_tqmse:
+            return new DCTemplate<QuantizerTurboQuantMSE<8, SL2>, Sim, SL2>(
+                    d, trained);
+        default:
+            FAISS_THROW_MSG("unknown qtype");
+    }
+}
+
+/*******************************************************************
+ * sq_select_distance_computer: returns SQDistanceComputer*
+ *******************************************************************/
+
+template <>
+SQDistanceComputer* sq_select_distance_computer<THE_LEVEL_TO_DISPATCH>(
+        MetricType metric,
+        ScalarQuantizer::QuantizerType qtype,
+        size_t d,
+        const std::vector<float>& trained) {
+    if (metric == METRIC_L2) {
+        return select_distance_computer_body<SimilarityL2<SL>, SL>(
+                qtype, d, trained);
+    } else {
+        return select_distance_computer_body<SimilarityIP<SL>, SL>(
+                qtype, d, trained);
+    }
+}
+
+/*******************************************************************
+ * sq_select_InvertedListScanner: returns InvertedListScanner*
+ *******************************************************************/
+
+template <>
+InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
+        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) {
+    auto scan = [&]<class DCClass>() -> InvertedListScanner* {
+        if constexpr (DCClass::Sim::metric_type == METRIC_L2) {
+            return new IVFSQScannerL2<DCClass>(
+                    int(d),
+                    trained,
+                    code_size,
+                    quantizer,
+                    store_pairs,
+                    sel,
+                    by_residual);
+        } else if constexpr (
+                DCClass::Sim::metric_type == METRIC_INNER_PRODUCT) {
+            return new IVFSQScannerIP<DCClass>(
+                    int(d), trained, code_size, store_pairs, sel, by_residual);
+        } else {
+            FAISS_THROW_MSG("unsupported metric type");
+        }
+    };
+
+    auto select_by_simd_and_metric =
+            [&]<SIMDLevel SL2, class Similarity>() -> InvertedListScanner* {
+        // Return nullptr for incompatible dimensions in SIMD cases
+        if constexpr (SL2 != SIMDLevel::NONE) {
+            if (!is_dimension_compatible<SL2>(d)) {
+                return nullptr;
+            }
+        }
+        switch (qtype) {
+            case ScalarQuantizer::QT_8bit_uniform:
+                return scan.template operator()<DCTemplate<
+                        QuantizerTemplate<
+                                Codec8bit<SL2>,
+                                QuantizerTemplateScaling::UNIFORM,
+                                SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_4bit_uniform:
+                return scan.template operator()<DCTemplate<
+                        QuantizerTemplate<
+                                Codec4bit<SL2>,
+                                QuantizerTemplateScaling::UNIFORM,
+                                SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_8bit:
+                return scan.template operator()<DCTemplate<
+                        QuantizerTemplate<
+                                Codec8bit<SL2>,
+                                QuantizerTemplateScaling::NON_UNIFORM,
+                                SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_4bit:
+                return scan.template operator()<DCTemplate<
+                        QuantizerTemplate<
+                                Codec4bit<SL2>,
+                                QuantizerTemplateScaling::NON_UNIFORM,
+                                SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_6bit:
+                return scan.template operator()<DCTemplate<
+                        QuantizerTemplate<
+                                Codec6bit<SL2>,
+                                QuantizerTemplateScaling::NON_UNIFORM,
+                                SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_fp16:
+                return scan.template
+                operator()<DCTemplate<QuantizerFP16<SL2>, Similarity, SL2>>();
+            case ScalarQuantizer::QT_bf16:
+                return scan.template
+                operator()<DCTemplate<QuantizerBF16<SL2>, Similarity, SL2>>();
+            case ScalarQuantizer::QT_8bit_direct:
+                if constexpr (SL2 == SIMDLevel::AVX512) {
+                    if (d % 32 == 0) {
+                        return scan.template
+                        operator()<DistanceComputerByte<Similarity, SL2>>();
+                    }
+                } else if constexpr (SL2 == SIMDLevel::AVX2) {
+                    if (d % 16 == 0) {
+                        return scan.template
+                        operator()<DistanceComputerByte<Similarity, SL2>>();
+                    }
+                }
+                return scan.template operator()<DCTemplate<
+                        Quantizer8bitDirect<SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_8bit_direct_signed:
+                return scan.template operator()<DCTemplate<
+                        Quantizer8bitDirectSigned<SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_0bit:
+                return new IVFCoarseDistanceScanner(
+                        Similarity::metric_type != METRIC_L2, store_pairs, sel);
+            case ScalarQuantizer::QT_1bit_tqmse:
+                return scan.template operator()<DCTemplate<
+                        QuantizerTurboQuantMSE<1, SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_2bit_tqmse:
+                return scan.template operator()<DCTemplate<
+                        QuantizerTurboQuantMSE<2, SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_3bit_tqmse:
+                return scan.template operator()<DCTemplate<
+                        QuantizerTurboQuantMSE<3, SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_4bit_tqmse:
+                return scan.template operator()<DCTemplate<
+                        QuantizerTurboQuantMSE<4, SL2>,
+                        Similarity,
+                        SL2>>();
+            case ScalarQuantizer::QT_8bit_tqmse:
+                return scan.template operator()<DCTemplate<
+                        QuantizerTurboQuantMSE<8, SL2>,
+                        Similarity,
+                        SL2>>();
+            default:
+                FAISS_THROW_MSG("unknown qtype");
+        }
+    };
+
+    if (mt == METRIC_L2) {
+        return select_by_simd_and_metric
+                .template operator()<SL, SimilarityL2<SL>>();
+    } else if (mt == METRIC_INNER_PRODUCT) {
+        return select_by_simd_and_metric
+                .template operator()<SL, SimilarityIP<SL>>();
+    }
+    FAISS_THROW_MSG("unsupported metric type");
+}
+
+} // namespace scalar_quantizer
+
+} // namespace faiss