faiss/faiss/IndexBinaryIVF.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.
 */

// -*- c++ -*-

#include <faiss/IndexBinaryIVF.h>

#include <omp.h>
#include <cinttypes>
#include <cstdio>

#include <algorithm>
#include <limits>
#include <memory>

#include <faiss/IndexFlat.h>
#include <faiss/IndexLSH.h>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/utils/hamming.h>
#include <faiss/utils/utils.h>

#include <faiss/impl/simd_dispatch.h>

// Scalar (NONE) fallback for dynamic dispatch
#define THE_SIMD_LEVEL SIMDLevel::NONE
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/impl/binary_hamming/IndexBinaryIVF_impl.h>
#include <faiss/utils/hamming_distance/hamming_computer-generic.h>
#undef THE_SIMD_LEVEL

namespace faiss {

IndexBinaryIVF::IndexBinaryIVF(
        IndexBinary* quantizer_,
        size_t d_,
        size_t nlist_,
        bool own_invlists_)
        : IndexBinary(d_),
          invlists(
                  own_invlists_ ? new ArrayInvertedLists(nlist_, code_size)
                                : nullptr),
          own_invlists(own_invlists_),
          quantizer(quantizer_),
          nlist(nlist_) {
    FAISS_THROW_IF_NOT(d_ == static_cast<size_t>(quantizer_->d));
    is_trained = quantizer_->is_trained &&
            (static_cast<size_t>(quantizer_->ntotal) == nlist_);
    cp.niter = 10;
}

IndexBinaryIVF::IndexBinaryIVF() = default;

void IndexBinaryIVF::add(idx_t n, const uint8_t* x) {
    add_with_ids(n, x, nullptr);
}

void IndexBinaryIVF::add_with_ids(
        idx_t n,
        const uint8_t* x,
        const idx_t* xids) {
    add_core(n, x, xids, nullptr);
}

void IndexBinaryIVF::add_core(
        idx_t n,
        const uint8_t* x,
        const idx_t* xids,
        const idx_t* precomputed_idx) {
    FAISS_THROW_IF_NOT(is_trained);
    FAISS_THROW_IF_NOT_MSG(invlists, "invlists not initialized");
    direct_map.check_can_add(xids);

    const idx_t* idx;

    std::unique_ptr<idx_t[]> scoped_idx;

    if (precomputed_idx) {
        idx = precomputed_idx;
    } else {
        scoped_idx.reset(new idx_t[n]);
        quantizer->assign(n, x, scoped_idx.get());
        idx = scoped_idx.get();
    }

    idx_t n_add = 0;
    for (idx_t i = 0; i < n; i++) {
        idx_t id = xids ? xids[i] : ntotal + i;
        idx_t list_no = idx[i];

        if (list_no < 0) {
            direct_map.add_single_id(id, -1, 0);
        } else {
            const uint8_t* xi = x + i * code_size;
            size_t offset = invlists->add_entry(list_no, id, xi);

            direct_map.add_single_id(id, list_no, offset);
        }

        n_add++;
    }
    if (verbose) {
        printf("IndexBinaryIVF::add_with_ids: added "
               "%" PRId64 " / %" PRId64 " vectors\n",
               n_add,
               n);
    }
    ntotal += n_add;
}

void IndexBinaryIVF::make_direct_map(bool b) {
    if (b) {
        direct_map.set_type(DirectMap::Array, invlists, ntotal);
    } else {
        direct_map.set_type(DirectMap::NoMap, invlists, ntotal);
    }
}

void IndexBinaryIVF::set_direct_map_type(DirectMap::Type type) {
    direct_map.set_type(type, invlists, ntotal);
}

void IndexBinaryIVF::search(
        idx_t n,
        const uint8_t* x,
        idx_t k,
        int32_t* distances,
        idx_t* labels,
        const SearchParameters* params_in) const {
    FAISS_THROW_IF_NOT(k > 0);
    FAISS_THROW_IF_NOT_MSG(invlists, "IVF index has no inverted lists");
    const IVFSearchParameters* params = nullptr;
    if (params_in) {
        params = dynamic_cast<const IVFSearchParameters*>(params_in);
        FAISS_THROW_IF_NOT_MSG(
                params, "IndexBinaryIVF params have incorrect type");
        FAISS_THROW_IF_MSG(
                params->sel, "IDSelector is not supported for IndexBinaryIVF");
    }
    const size_t nprobe_ =
            std::min(nlist, params ? params->nprobe : this->nprobe);
    FAISS_THROW_IF_NOT(nprobe_ > 0);

    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe_]);
    std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe_]);

    double t0 = getmillisecs();
    quantizer->search(
            n,
            x,
            nprobe_,
            coarse_dis.get(),
            idx.get(),
            params ? params->quantizer_params : nullptr);
    indexIVF_stats.quantization_time += getmillisecs() - t0;

    t0 = getmillisecs();
    invlists->prefetch_lists(idx.get(), static_cast<int>(n * nprobe_));

    search_preassigned(
            n,
            x,
            k,
            idx.get(),
            coarse_dis.get(),
            distances,
            labels,
            false,
            params);
    indexIVF_stats.search_time += getmillisecs() - t0;
}

void IndexBinaryIVF::reconstruct(idx_t key, uint8_t* recons) const {
    idx_t lo = direct_map.get(key);
    reconstruct_from_offset(lo_listno(lo), lo_offset(lo), recons);
}

void IndexBinaryIVF::reconstruct_n(idx_t i0, idx_t ni, uint8_t* recons) const {
    FAISS_THROW_IF_NOT(ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));

    for (size_t list_no = 0; list_no < nlist; list_no++) {
        size_t list_size = invlists->list_size(list_no);
        const idx_t* idlist = invlists->get_ids(list_no);

        for (size_t offset = 0; offset < list_size; offset++) {
            idx_t id = idlist[offset];
            if (!(id >= i0 && id < i0 + ni)) {
                continue;
            }

            uint8_t* reconstructed = recons + (id - i0) * d;
            reconstruct_from_offset(list_no, offset, reconstructed);
        }
    }
}

void IndexBinaryIVF::search_and_reconstruct(
        idx_t n,
        const uint8_t* __restrict x,
        idx_t k,
        int32_t* __restrict distances,
        idx_t* __restrict labels,
        uint8_t* __restrict recons,
        const SearchParameters* params) const {
    FAISS_THROW_IF_NOT_MSG(
            !params, "search params not supported for this index");
    const size_t nprobe_2 = std::min(nlist, this->nprobe);
    FAISS_THROW_IF_NOT(k > 0);
    FAISS_THROW_IF_NOT(nprobe_2 > 0);

    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe_2]);
    std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe_2]);

    quantizer->search(n, x, nprobe_2, coarse_dis.get(), idx.get());

    invlists->prefetch_lists(idx.get(), static_cast<int>(n * nprobe_2));

    // search_preassigned() with `store_pairs` enabled to obtain the list_no
    // and offset into `codes` for reconstruction
    search_preassigned(
            n,
            x,
            k,
            idx.get(),
            coarse_dis.get(),
            distances,
            labels,
            /* store_pairs */ true);
    for (idx_t i = 0; i < n; ++i) {
        for (idx_t j = 0; j < k; ++j) {
            idx_t ij = i * k + j;
            idx_t key = labels[ij];
            uint8_t* reconstructed = recons + ij * d;
            if (key < 0) {
                // Fill with NaNs
                memset(reconstructed, -1, sizeof(*reconstructed) * d);
            } else {
                int list_no = key >> 32;
                int offset = key & 0xffffffff;

                // Update label to the actual id
                labels[ij] = invlists->get_single_id(list_no, offset);

                reconstruct_from_offset(list_no, offset, reconstructed);
            }
        }
    }
}

void IndexBinaryIVF::reconstruct_from_offset(
        idx_t list_no,
        idx_t offset,
        uint8_t* recons) const {
    memcpy(recons,
           InvertedLists::ScopedCodes(invlists, list_no, offset).get(),
           code_size);
}

void IndexBinaryIVF::reset() {
    direct_map.clear();
    invlists->reset();
    ntotal = 0;
}

size_t IndexBinaryIVF::remove_ids(const IDSelector& sel) {
    size_t nremove = direct_map.remove_ids(sel, invlists);
    ntotal -= nremove;
    return nremove;
}

void IndexBinaryIVF::train(idx_t n, const uint8_t* x) {
    if (verbose) {
        printf("Training quantizer\n");
    }

    if (quantizer->is_trained &&
        (static_cast<size_t>(quantizer->ntotal) == nlist)) {
        if (verbose) {
            printf("IVF quantizer does not need training.\n");
        }
    } else {
        if (verbose) {
            printf("Training quantizer on %" PRId64 " vectors in %dD\n", n, d);
        }

        Clustering clus(d, static_cast<int>(nlist), cp);
        quantizer->reset();

        IndexFlatL2 index_tmp(d);

        if (clustering_index && verbose) {
            printf("using clustering_index of dimension %d to do the clustering\n",
                   clustering_index->d);
        }

        // LSH codec that is able to convert the binary vectors to floats.
        IndexLSH codec(d, d, false, false);

        clus.train_encoded(
                n, x, &codec, clustering_index ? *clustering_index : index_tmp);

        // convert clusters to binary
        std::unique_ptr<uint8_t[]> x_b(new uint8_t[clus.k * code_size]);
        real_to_binary(d * clus.k, clus.centroids.data(), x_b.get());

        quantizer->add(clus.k, x_b.get());
        quantizer->is_trained = true;
    }

    is_trained = true;
}

void IndexBinaryIVF::check_compatible_for_merge(
        const IndexBinary& otherIndex) const {
    auto other = dynamic_cast<const IndexBinaryIVF*>(&otherIndex);
    FAISS_THROW_IF_NOT(other);
    FAISS_THROW_IF_NOT(other->d == d);
    FAISS_THROW_IF_NOT(other->nlist == nlist);
    FAISS_THROW_IF_NOT(other->code_size == code_size);
    FAISS_THROW_IF_NOT_MSG(
            direct_map.no() && other->direct_map.no(),
            "direct map copy not implemented");
    FAISS_THROW_IF_NOT_MSG(
            typeid(*this) == typeid(*other),
            "can only merge indexes of the same type");
}

void IndexBinaryIVF::merge_from(IndexBinary& otherIndex, idx_t add_id) {
    // minimal sanity checks
    check_compatible_for_merge(otherIndex);
    auto other = static_cast<IndexBinaryIVF*>(&otherIndex);
    invlists->merge_from(other->invlists, add_id);
    ntotal += other->ntotal;
    other->ntotal = 0;
}

void IndexBinaryIVF::replace_invlists(InvertedLists* il, bool own) {
    FAISS_THROW_IF_NOT(
            il->nlist == nlist &&
            il->code_size == static_cast<size_t>(code_size));
    if (own_invlists) {
        delete invlists;
    }
    invlists = il;
    own_invlists = own;
}

// IVFBinaryScannerL2, search_knn_hamming_count, BlockSearch,
// BlockSearchVariableK, search_knn_hamming_per_invlist are now in
// impl/binary_hamming/IndexBinaryIVF_impl.h (compiled per-ISA)

namespace {

void search_knn_hamming_heap(
        const IndexBinaryIVF* ivf,
        size_t n,
        const uint8_t* __restrict x,
        idx_t k,
        const idx_t* __restrict keys,
        const int32_t* __restrict coarse_dis,
        int32_t* __restrict distances,
        idx_t* __restrict labels,
        bool store_pairs,
        const IVFSearchParameters* params) {
    idx_t nprobe = params ? params->nprobe : ivf->nprobe;
    nprobe = std::min((idx_t)ivf->nlist, nprobe);
    idx_t max_codes = params ? params->max_codes : ivf->max_codes;
    const idx_t unlimited_list_size = std::numeric_limits<idx_t>::max();
    if (max_codes == 0) {
        max_codes = unlimited_list_size;
    }
    MetricType metric_type = ivf->metric_type;

    // almost verbatim copy from IndexIVF::search_preassigned

    size_t nlistv = 0, ndis = 0, nheap = 0;
    using HeapForIP = CMin<int32_t, idx_t>;
    using HeapForL2 = CMax<int32_t, idx_t>;

#pragma omp parallel if (n > 1) reduction(+ : nlistv, ndis, nheap)
    {
        std::unique_ptr<BinaryInvertedListScanner> scanner(
                ivf->get_InvertedListScanner(store_pairs));

#pragma omp for
        for (idx_t i = 0; i < static_cast<idx_t>(n); i++) {
            const uint8_t* xi = x + i * ivf->code_size;
            scanner->set_query(xi);

            const idx_t* keysi = keys + i * nprobe;
            int32_t* simi = distances + k * i;
            idx_t* idxi = labels + k * i;

            if (metric_type == METRIC_INNER_PRODUCT) {
                heap_heapify<HeapForIP>(k, simi, idxi);
            } else {
                heap_heapify<HeapForL2>(k, simi, idxi);
            }

            size_t nscan = 0;

            for (idx_t ik = 0; ik < nprobe; ik++) {
                idx_t key = keysi[ik]; /* select the list  */
                if (key < 0) {
                    // not enough centroids for multiprobe
                    continue;
                }
                FAISS_THROW_IF_NOT_FMT(
                        key < (idx_t)ivf->nlist,
                        "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
                        key,
                        static_cast<size_t>(ik),
                        ivf->nlist);

                scanner->set_list(key, coarse_dis[i * nprobe + ik]);

                nlistv++;

                size_t list_size = ivf->invlists->list_size(key);
                size_t list_size_max = static_cast<size_t>(max_codes) - nscan;
                if (list_size > list_size_max) {
                    list_size = list_size_max;
                }
                InvertedLists::ScopedCodes scodes(ivf->invlists, key);
                std::unique_ptr<InvertedLists::ScopedIds> sids;
                const idx_t* ids = nullptr;

                if (!store_pairs) {
                    sids = std::make_unique<InvertedLists::ScopedIds>(
                            ivf->invlists, key);
                    ids = sids->get();
                }

                nheap += scanner->scan_codes(
                        list_size, scodes.get(), ids, simi, idxi, k);

                nscan += list_size;
                if (nscan >= static_cast<size_t>(max_codes)) {
                    break;
                }
            }

            ndis += nscan;
            if (metric_type == METRIC_INNER_PRODUCT) {
                heap_reorder<HeapForIP>(k, simi, idxi);
            } else {
                heap_reorder<HeapForL2>(k, simi, idxi);
            }

        } // parallel for
    } // parallel

    indexIVF_stats.nq += n;
    indexIVF_stats.nlist += nlistv;
    indexIVF_stats.ndis += ndis;
    indexIVF_stats.nheap_updates += nheap;
}

} // anonymous namespace

// The remaining template code (search_knn_hamming_count,
// search_knn_hamming_per_invlist, etc.) has been moved to
// impl/binary_hamming/IndexBinaryIVF_impl.h

BinaryInvertedListScanner* IndexBinaryIVF::get_InvertedListScanner(
        bool store_pairs) const {
    return with_simd_level([&]<SIMDLevel SL>() {
        return make_binary_ivf_scanner_fixSL<SL>(code_size, store_pairs);
    });
}

void IndexBinaryIVF::search_preassigned(
        idx_t n,
        const uint8_t* x,
        idx_t k,
        const idx_t* cidx,
        const int32_t* cdis,
        int32_t* dis,
        idx_t* idx,
        bool store_pairs,
        const IVFSearchParameters* params) const {
    if (per_invlist_search) {
        with_simd_level([&]<SIMDLevel SL>() {
            search_knn_hamming_per_invlist_fixSL<SL>(
                    code_size,
                    this,
                    n,
                    x,
                    k,
                    cidx,
                    cdis,
                    dis,
                    idx,
                    store_pairs,
                    params);
        });
    } else if (use_heap) {
        search_knn_hamming_heap(
                this, n, x, k, cidx, cdis, dis, idx, store_pairs, params);
    } else {
        with_simd_level([&]<SIMDLevel SL>() {
            search_knn_hamming_count_fixSL<SL>(
                    code_size,
                    store_pairs,
                    this,
                    n,
                    x,
                    cidx,
                    k,
                    dis,
                    idx,
                    params);
        });
    }
}

void IndexBinaryIVF::range_search(
        idx_t n,
        const uint8_t* __restrict x,
        int radius,
        RangeSearchResult* __restrict res,
        const SearchParameters* params) const {
    FAISS_THROW_IF_NOT_MSG(
            !params, "search params not supported for this index");
    const size_t nprobe_2 = std::min(nlist, this->nprobe);
    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe_2]);
    std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe_2]);

    double t0 = getmillisecs();
    quantizer->search(n, x, nprobe_2, coarse_dis.get(), idx.get());
    indexIVF_stats.quantization_time += getmillisecs() - t0;

    t0 = getmillisecs();
    invlists->prefetch_lists(idx.get(), static_cast<int>(n * nprobe_2));

    range_search_preassigned(n, x, radius, idx.get(), coarse_dis.get(), res);

    indexIVF_stats.search_time += getmillisecs() - t0;
}

void IndexBinaryIVF::range_search_preassigned(
        idx_t n,
        const uint8_t* __restrict x,
        int radius,
        const idx_t* __restrict assign,
        const int32_t* __restrict /* centroid_dis */,
        RangeSearchResult* __restrict res) const {
    const size_t nprobe_2 = std::min(nlist, this->nprobe);
    bool store_pairs = false;
    size_t nlistv = 0, ndis = 0;

    std::vector<RangeSearchPartialResult*> all_pres(omp_get_max_threads());

#pragma omp parallel reduction(+ : nlistv, ndis)
    {
        RangeSearchPartialResult pres(res);
        std::unique_ptr<BinaryInvertedListScanner> scanner(
                get_InvertedListScanner(store_pairs));
        FAISS_THROW_IF_NOT(scanner.get());

        all_pres[omp_get_thread_num()] = &pres;

        auto scan_list_func = [&](size_t i, size_t ik, RangeQueryResult& qres) {
            idx_t key = assign[i * nprobe_2 + ik]; /* select the list  */
            if (key < 0) {
                return;
            }
            FAISS_THROW_IF_NOT_FMT(
                    key < (idx_t)nlist,
                    "Invalid key=%" PRId64 " at ik=%zd nlist=%zd\n",
                    key,
                    ik,
                    nlist);
            const size_t list_size = invlists->list_size(key);

            if (list_size == 0) {
                return;
            }

            InvertedLists::ScopedCodes scodes(invlists, key);
            InvertedLists::ScopedIds ids(invlists, key);

            scanner->set_list(key, assign[i * nprobe_2 + ik]);
            nlistv++;
            ndis += list_size;
            scanner->scan_codes_range(
                    list_size, scodes.get(), ids.get(), radius, qres);
        };

#pragma omp for
        for (idx_t i = 0; i < n; i++) {
            scanner->set_query(x + i * code_size);

            RangeQueryResult& qres = pres.new_result(i);

            for (size_t ik = 0; ik < nprobe_2; ik++) {
                scan_list_func(i, ik, qres);
            }
        }

        pres.finalize();
    }
    indexIVF_stats.nq += n;
    indexIVF_stats.nlist += nlistv;
    indexIVF_stats.ndis += ndis;
}

IndexBinaryIVF::~IndexBinaryIVF() {
    if (own_invlists) {
        delete invlists;
    }

    if (own_fields) {
        delete quantizer;
    }
}

} // namespace faiss