faiss/faiss/IndexAdditiveQuantizerFastScan.cpp

/*
 * 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/IndexAdditiveQuantizerFastScan.h>

#include <cassert>
#include <memory>

#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/LocalSearchQuantizer.h>
#include <faiss/impl/ResidualQuantizer.h>
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
#include <faiss/impl/fast_scan/fast_scan.h>
#include <faiss/utils/quantize_lut.h>
#include <faiss/utils/utils.h>

namespace faiss {

inline size_t roundup(size_t a, size_t b) {
    return (a + b - 1) / b * b;
}

IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan(
        AdditiveQuantizer* aq,
        MetricType metric,
        int bbs) {
    init(aq, metric, bbs);
}

void IndexAdditiveQuantizerFastScan::init(
        AdditiveQuantizer* aq_init,
        MetricType metric,
        int bbs) {
    FAISS_THROW_IF_NOT(aq_init != nullptr);
    FAISS_THROW_IF_NOT(!aq_init->nbits.empty());
    FAISS_THROW_IF_NOT(aq_init->nbits[0] == 4);
    if (metric == METRIC_INNER_PRODUCT) {
        FAISS_THROW_IF_NOT_MSG(
                aq_init->search_type == AdditiveQuantizer::ST_LUT_nonorm,
                "Search type must be ST_LUT_nonorm for IP metric");
    } else {
        FAISS_THROW_IF_NOT_MSG(
                aq_init->search_type == AdditiveQuantizer::ST_norm_lsq2x4 ||
                        aq_init->search_type ==
                                AdditiveQuantizer::ST_norm_rq2x4,
                "Search type must be lsq2x4 or rq2x4 for L2 metric");
    }

    this->aq = aq_init;
    if (metric == METRIC_L2) {
        M = aq_init->M + 2; // 2x4 bits AQ
    } else {
        M = aq_init->M;
    }
    init_fastscan(aq_init->d, M, 4, metric, bbs);

    max_train_points = 1024 * ksub * M;
}

IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan()
        : IndexFastScan() {
    is_trained = false;
    aq = nullptr;
}

IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan(
        const IndexAdditiveQuantizer& orig,
        int bbs) {
    init(orig.aq, orig.metric_type, bbs);

    ntotal = orig.ntotal;
    is_trained = orig.is_trained;
    orig_codes = orig.codes.data();

    ntotal2 = roundup(ntotal, bbs);
    codes.resize(ntotal2 * M2 / 2);
    pq4_pack_codes(orig_codes, ntotal, M, ntotal2, bbs, M2, codes.get());
}

IndexAdditiveQuantizerFastScan::~IndexAdditiveQuantizerFastScan() = default;

void IndexAdditiveQuantizerFastScan::train(idx_t n, const float* x_in) {
    if (is_trained) {
        return;
    }

    const int seed = 0x12345;
    size_t nt = n;
    const float* x = fvecs_maybe_subsample(
            d, &nt, max_train_points, x_in, verbose, seed);
    n = nt;
    if (verbose) {
        printf("training additive quantizer on %zd vectors\n", nt);
    }

    aq->verbose = verbose;
    aq->train(n, x);
    if (metric_type == METRIC_L2) {
        estimate_norm_scale(n, x);
    }

    is_trained = true;
}

void IndexAdditiveQuantizerFastScan::estimate_norm_scale(
        idx_t n,
        const float* x_in) {
    FAISS_THROW_IF_NOT(metric_type == METRIC_L2);

    constexpr int seed = 0x980903;
    constexpr size_t max_points_estimated = 65536;
    size_t ns = n;
    const float* x = fvecs_maybe_subsample(
            d, &ns, max_points_estimated, x_in, verbose, seed);
    n = ns;
    std::unique_ptr<float[]> del_x;
    if (x != x_in) {
        del_x.reset((float*)x);
    }

    std::vector<float> dis_tables(n * M * ksub);
    FastScanDistancePostProcessing empty_context;
    compute_float_LUT(dis_tables.data(), n, x, empty_context);

    // here we compute the mean of scales for each query
    // TODO: try max of scales
    double scale = 0;

#pragma omp parallel for reduction(+ : scale)
    for (idx_t i = 0; i < n; i++) {
        const float* lut = dis_tables.data() + i * M * ksub;
        scale += quantize_lut::aq_estimate_norm_scale(M, ksub, 2, lut);
    }
    scale /= n;
    norm_scale = (int)std::roundf(std::max(scale, 1.0));

    if (verbose) {
        printf("estimated norm scale: %lf\n", scale);
        printf("rounded norm scale: %d\n", norm_scale);
    }
}

void IndexAdditiveQuantizerFastScan::compute_codes(
        uint8_t* tmp_codes,
        idx_t n,
        const float* x) const {
    aq->compute_codes(x, tmp_codes, n);
}

void IndexAdditiveQuantizerFastScan::compute_float_LUT(
        float* lut,
        idx_t n,
        const float* x,
        const FastScanDistancePostProcessing&) const {
    if (metric_type == METRIC_INNER_PRODUCT) {
        aq->compute_LUT(n, x, lut, 1.0f);
    } else {
        // compute inner product look-up tables
        const size_t ip_dim12 = aq->M * ksub;
        const size_t norm_dim12 = 2 * ksub;
        std::vector<float> ip_lut(n * ip_dim12);
        aq->compute_LUT(n, x, ip_lut.data(), -2.0f);

        // copy and rescale norm look-up tables
        auto norm_tabs = aq->norm_tabs;
        if (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2) {
            for (size_t i = 0; i < norm_tabs.size(); i++) {
                norm_tabs[i] /= norm_scale;
            }
        }
        const float* norm_lut = norm_tabs.data();
        FAISS_THROW_IF_NOT(norm_tabs.size() == norm_dim12);

        // combine them
        for (idx_t i = 0; i < n; i++) {
            memcpy(lut, ip_lut.data() + i * ip_dim12, ip_dim12 * sizeof(*lut));
            lut += ip_dim12;
            memcpy(lut, norm_lut, norm_dim12 * sizeof(*lut));
            lut += norm_dim12;
        }
    }
}

void IndexAdditiveQuantizerFastScan::search(
        idx_t n,
        const float* x,
        idx_t k,
        float* distances,
        idx_t* labels,
        const SearchParameters* params) const {
    FAISS_THROW_IF_NOT_MSG(
            !params, "search params not supported for this index");
    FAISS_THROW_IF_NOT(k > 0);
    bool rescale = (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2);
    if (!rescale) {
        IndexFastScan::search(n, x, k, distances, labels);
        return;
    }

    FastScanDistancePostProcessing context;
    context.pq2x4_scale = norm_scale;
    if (metric_type == METRIC_L2) {
        search_dispatch_implem<true>(n, x, k, distances, labels, context);
    } else {
        search_dispatch_implem<false>(n, x, k, distances, labels, context);
    }
}

void IndexAdditiveQuantizerFastScan::sa_decode(
        idx_t n,
        const uint8_t* bytes,
        float* x) const {
    aq->decode(bytes, x, n);
}

/**************************************************************************************
 * IndexResidualQuantizerFastScan
 **************************************************************************************/

IndexResidualQuantizerFastScan::IndexResidualQuantizerFastScan(
        int d,        ///< dimensionality of the input vectors
        size_t M,     ///< number of subquantizers
        size_t nbits, ///< number of bit per subvector index
        MetricType metric,
        Search_type_t search_type,
        int bbs)
        : rq(d, M, nbits, search_type) {
    init(&rq, metric, bbs);
}

IndexResidualQuantizerFastScan::IndexResidualQuantizerFastScan() {
    aq = &rq;
}

/**************************************************************************************
 * IndexLocalSearchQuantizerFastScan
 **************************************************************************************/

IndexLocalSearchQuantizerFastScan::IndexLocalSearchQuantizerFastScan(
        int d,
        size_t M,     ///< number of subquantizers
        size_t nbits, ///< number of bit per subvector index
        MetricType metric,
        Search_type_t search_type,
        int bbs)
        : lsq(d, M, nbits, search_type) {
    init(&lsq, metric, bbs);
}

IndexLocalSearchQuantizerFastScan::IndexLocalSearchQuantizerFastScan() {
    aq = &lsq;
}

/**************************************************************************************
 * IndexProductResidualQuantizerFastScan
 **************************************************************************************/

IndexProductResidualQuantizerFastScan::IndexProductResidualQuantizerFastScan(
        int d,          ///< dimensionality of the input vectors
        size_t nsplits, ///< number of residual quantizers
        size_t Msub,    ///< number of subquantizers per RQ
        size_t nbits,   ///< number of bit per subvector index
        MetricType metric,
        Search_type_t search_type,
        int bbs)
        : prq(d, nsplits, Msub, nbits, search_type) {
    init(&prq, metric, bbs);
}

IndexProductResidualQuantizerFastScan::IndexProductResidualQuantizerFastScan() {
    aq = &prq;
}

/**************************************************************************************
 * IndexProductLocalSearchQuantizerFastScan
 **************************************************************************************/

IndexProductLocalSearchQuantizerFastScan::
        IndexProductLocalSearchQuantizerFastScan(
                int d,          ///< dimensionality of the input vectors
                size_t nsplits, ///< number of local search quantizers
                size_t Msub,    ///< number of subquantizers per LSQ
                size_t nbits,   ///< number of bit per subvector index
                MetricType metric,
                Search_type_t search_type,
                int bbs)
        : plsq(d, nsplits, Msub, nbits, search_type) {
    init(&plsq, metric, bbs);
}

IndexProductLocalSearchQuantizerFastScan::
        IndexProductLocalSearchQuantizerFastScan() {
    aq = &plsq;
}

} // namespace faiss