faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/code_distance/code_distance.h

@@ -1,186 +0,0 @@
-/*
- * 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/platform_macros.h>
-
-// This directory contains functions to compute a distance
-// from a given PQ code to a query vector, given that the
-// distances to a query vector for pq.M codebooks are precomputed.
-//
-// The code was originally the part of IndexIVFPQ.cpp.
-// The baseline implementation can be found in
-//   code_distance-generic.h, distance_single_code_generic().
-
-// The reason for this somewhat unusual structure is that
-// custom implementations may need to fall off to generic
-// implementation in certain cases. So, say, avx2 header file
-// needs to reference the generic header file. This is
-// why the names of the functions for custom implementations
-// have this _generic or _avx2 suffix.
-
-#ifdef __AVX2__
-
-#include <faiss/impl/code_distance/code_distance-avx2.h>
-
-namespace faiss {
-
-template <typename PQDecoderT>
-inline float distance_single_code(
-        // number of subquantizers
-        const size_t M,
-        // number of bits per quantization index
-        const size_t nbits,
-        // precomputed distances, layout (M, ksub)
-        const float* sim_table,
-        // the code
-        const uint8_t* code) {
-    return distance_single_code_avx2<PQDecoderT>(M, nbits, sim_table, code);
-}
-
-template <typename PQDecoderT>
-inline void distance_four_codes(
-        // number of subquantizers
-        const size_t M,
-        // number of bits per quantization index
-        const size_t nbits,
-        // precomputed distances, layout (M, ksub)
-        const float* sim_table,
-        // codes
-        const uint8_t* __restrict code0,
-        const uint8_t* __restrict code1,
-        const uint8_t* __restrict code2,
-        const uint8_t* __restrict code3,
-        // computed distances
-        float& result0,
-        float& result1,
-        float& result2,
-        float& result3) {
-    distance_four_codes_avx2<PQDecoderT>(
-            M,
-            nbits,
-            sim_table,
-            code0,
-            code1,
-            code2,
-            code3,
-            result0,
-            result1,
-            result2,
-            result3);
-}
-
-} // namespace faiss
-
-#elif defined(__ARM_FEATURE_SVE)
-
-#include <faiss/impl/code_distance/code_distance-sve.h>
-
-namespace faiss {
-
-template <typename PQDecoderT>
-inline float distance_single_code(
-        // the product quantizer
-        const size_t M,
-        // number of bits per quantization index
-        const size_t nbits,
-        // precomputed distances, layout (M, ksub)
-        const float* sim_table,
-        // the code
-        const uint8_t* code) {
-    return distance_single_code_sve<PQDecoderT>(M, nbits, sim_table, code);
-}
-
-template <typename PQDecoderT>
-inline void distance_four_codes(
-        // the product quantizer
-        const size_t M,
-        // number of bits per quantization index
-        const size_t nbits,
-        // precomputed distances, layout (M, ksub)
-        const float* sim_table,
-        // codes
-        const uint8_t* __restrict code0,
-        const uint8_t* __restrict code1,
-        const uint8_t* __restrict code2,
-        const uint8_t* __restrict code3,
-        // computed distances
-        float& result0,
-        float& result1,
-        float& result2,
-        float& result3) {
-    distance_four_codes_sve<PQDecoderT>(
-            M,
-            nbits,
-            sim_table,
-            code0,
-            code1,
-            code2,
-            code3,
-            result0,
-            result1,
-            result2,
-            result3);
-}
-
-} // namespace faiss
-
-#else
-
-#include <faiss/impl/code_distance/code_distance-generic.h>
-
-namespace faiss {
-
-template <typename PQDecoderT>
-inline float distance_single_code(
-        // number of subquantizers
-        const size_t M,
-        // number of bits per quantization index
-        const size_t nbits,
-        // precomputed distances, layout (M, ksub)
-        const float* sim_table,
-        // the code
-        const uint8_t* code) {
-    return distance_single_code_generic<PQDecoderT>(M, nbits, sim_table, code);
-}
-
-template <typename PQDecoderT>
-inline void distance_four_codes(
-        // number of subquantizers
-        const size_t M,
-        // number of bits per quantization index
-        const size_t nbits,
-        // precomputed distances, layout (M, ksub)
-        const float* sim_table,
-        // codes
-        const uint8_t* __restrict code0,
-        const uint8_t* __restrict code1,
-        const uint8_t* __restrict code2,
-        const uint8_t* __restrict code3,
-        // computed distances
-        float& result0,
-        float& result1,
-        float& result2,
-        float& result3) {
-    distance_four_codes_generic<PQDecoderT>(
-            M,
-            nbits,
-            sim_table,
-            code0,
-            code1,
-            code2,
-            code3,
-            result0,
-            result1,
-            result2,
-            result3);
-}
-
-} // namespace faiss
-
-#endif

data/vendor/faiss/faiss/impl/pq4_fast_scan.h

@@ -1,216 +0,0 @@
-/*
- * 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 <cstdint>
-#include <cstdlib>
-
-#include <faiss/impl/CodePacker.h>
-
-/** PQ4 SIMD packing and accumulation functions
- *
- * The basic kernel accumulates nq query vectors with bbs = nb * 2 * 16 vectors
- * and produces an output matrix for that. It is interesting for nq * nb <= 4,
- * otherwise register spilling becomes too large.
- *
- * The implementation of these functions is spread over 3 cpp files to reduce
- * parallel compile times. Templates are instantiated explicitly.
- */
-
-namespace faiss {
-
-struct NormTableScaler;
-struct SIMDResultHandler;
-
-/** Pack codes for consumption by the SIMD kernels.
- *  The unused bytes are set to 0.
- *
- * @param codes   input codes, size (ntotal, ceil(M / 2))
- * @param ntotal  number of input codes
- * @param nb      output number of codes (ntotal rounded up to a multiple of
- *                bbs)
- * @param nsq      number of sub-quantizers (=M rounded up to a multiple of 2)
- * @param bbs     size of database blocks (multiple of 32)
- * @param blocks  output array, size nb * nsq / 2.
- * @param code_stride  optional stride between consecutive codes (0 = use
-default (M + 1) / 2)
- */
-void pq4_pack_codes(
-        const uint8_t* codes,
-        size_t ntotal,
-        size_t M,
-        size_t nb,
-        size_t bbs,
-        size_t nsq,
-        uint8_t* blocks,
-        size_t code_stride = 0);
-
-/** Same as pack_codes but write in a given range of the output,
- * leaving the rest untouched. Assumes allocated entries are 0 on input.
- *
- * @param codes   input codes, size (i1 - i0, ceil(M / 2))
- * @param i0      first output code to write
- * @param i1      last output code to write
- * @param blocks  output array, size at least ceil(i1 / bbs) * bbs * nsq / 2
- * @param code_stride  optional stride between consecutive codes (0 = use
- * default (M + 1) / 2)
- */
-void pq4_pack_codes_range(
-        const uint8_t* codes,
-        size_t M,
-        size_t i0,
-        size_t i1,
-        size_t bbs,
-        size_t nsq,
-        uint8_t* blocks,
-        size_t code_stride = 0);
-
-/** get a single element from a packed codes table
- *
- * @param vector_id        vector id
- * @param sq       subquantizer (< nsq)
- */
-uint8_t pq4_get_packed_element(
-        const uint8_t* data,
-        size_t bbs,
-        size_t nsq,
-        size_t vector_id,
-        size_t sq);
-
-/** set a single element "code" into a packed codes table
- *
- * @param vector_id       vector id
- * @param sq       subquantizer (< nsq)
- */
-void pq4_set_packed_element(
-        uint8_t* data,
-        uint8_t code,
-        size_t bbs,
-        size_t nsq,
-        size_t vector_id,
-        size_t sq);
-
-/** CodePacker API for the PQ4 fast-scan */
-struct CodePackerPQ4 : CodePacker {
-    size_t nsq;
-
-    CodePackerPQ4(size_t nsq, size_t bbs);
-
-    void pack_1(const uint8_t* flat_code, size_t offset, uint8_t* block)
-            const final;
-    void unpack_1(const uint8_t* block, size_t offset, uint8_t* flat_code)
-            const final;
-};
-
-/** Pack Look-up table for consumption by the kernel.
- *
- * @param nq      number of queries
- * @param nsq     number of sub-quantizers (multiple of 2)
- * @param src     input array, size (nq, 16)
- * @param dest    output array, size (nq, 16)
- */
-void pq4_pack_LUT(int nq, int nsq, const uint8_t* src, uint8_t* dest);
-
-/** Loop over database elements and accumulate results into result handler
- *
- * @param nq      number of queries
- * @param nb      number of database elements
- * @param bbs     size of database blocks (multiple of 32)
- * @param nsq     number of sub-quantizers (multiple of 2)
- * @param codes   packed codes array
- * @param LUT     packed look-up table
- * @param scaler  scaler to scale the encoded norm
- */
-void pq4_accumulate_loop(
-        int nq,
-        size_t nb,
-        int bbs,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        SIMDResultHandler& res,
-        const NormTableScaler* scaler);
-
-/* qbs versions, supported only for bbs=32.
- *
- * The kernel function runs the kernel for *several* query blocks
- * and bbs database vectors. The sizes of the blocks are encoded in qbs as
- * base-16 digits.
- *
- * For example, qbs = 0x1223 means that the kernel will be run 4 times, the
- * first time with 3 query vectors, second time with 2 query vectors, then 2
- * vectors again and finally with 1 query vector. The output block will thus be
- * nq = 3 + 2 + 2 + 1 = 6 queries. For a given total block size, the optimal
- * decomposition into sub-blocks (measured empirically) is given by
- * preferred_qbs().
- */
-
-/* compute the number of queries from a base-16 decomposition */
-int pq4_qbs_to_nq(int qbs);
-
-/** return the preferred decomposition in blocks for a nb of queries. */
-int pq4_preferred_qbs(int nq);
-
-/** Pack Look-up table for consumption by the kernel.
- *
- * @param qbs     4-bit encoded number of query blocks, the total number of
- *                queries handled (nq) is deduced from it
- * @param nsq     number of sub-quantizers (multiple of 2)
- * @param src     input array, size (nq, 16)
- * @param dest    output array, size (nq, 16)
- * @return nq
- */
-int pq4_pack_LUT_qbs(int fqbs, int nsq, const uint8_t* src, uint8_t* dest);
-
-/** Same as pq4_pack_LUT_qbs, except the source vectors are remapped with q_map
- */
-int pq4_pack_LUT_qbs_q_map(
-        int qbs,
-        int nsq,
-        const uint8_t* src,
-        const int* q_map,
-        uint8_t* dest);
-
-/** Run accumulation loop.
- *
- * @param qbs     4-bit encoded number of queries
- * @param nb      number of database codes (multiple of bbs)
- * @param nsq     number of sub-quantizers
- * @param codes   encoded database vectors (packed)
- * @param LUT     look-up table (packed)
- * @param res     call-back for the results
- * @param scaler  scaler to scale the encoded norm
- */
-void pq4_accumulate_loop_qbs(
-        int qbs,
-        size_t nb,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        SIMDResultHandler& res,
-        const NormTableScaler* scaler = nullptr);
-
-/** Wrapper of pq4_accumulate_loop_qbs using simple StoreResultHandler
- *  and DummyScaler
- *
- * @param nq      number of queries
- * @param ntotal2 number of database elements (multiple of 32)
- * @param nsq     number of sub-quantizers (muliple of 2)
- * @param codes   packed codes array
- * @param LUT     packed look-up table
- * @param accu    array to store the results
- */
-void accumulate_to_mem(
-        int nq,
-        size_t ntotal2,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        uint16_t* accu);
-
-} // namespace faiss

data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp

@@ -1,224 +0,0 @@
-/*
- * 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/impl/pq4_fast_scan.h>
-
-#include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/LookupTableScaler.h>
-#include <faiss/impl/simd_result_handlers.h>
-
-namespace faiss {
-
-using namespace simd_result_handlers;
-
-/***************************************************************
- * accumulation functions
- ***************************************************************/
-
-namespace {
-
-/*
- * The computation kernel
- * It accumulates results for NQ queries and BB * 32 database elements
- * writes results in a ResultHandler
- */
-
-template <int NQ, int BB, class ResultHandler, class Scaler>
-void kernel_accumulate_block(
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        ResultHandler& res,
-        const Scaler& scaler) {
-    // distance accumulators
-    simd16uint16 accu[NQ][BB][4];
-
-    for (int q = 0; q < NQ; q++) {
-        for (int b = 0; b < BB; b++) {
-            accu[q][b][0].clear();
-            accu[q][b][1].clear();
-            accu[q][b][2].clear();
-            accu[q][b][3].clear();
-        }
-    }
-
-    for (int sq = 0; sq < nsq - scaler.nscale; sq += 2) {
-        simd32uint8 lut_cache[NQ];
-        for (int q = 0; q < NQ; q++) {
-            lut_cache[q] = simd32uint8(LUT);
-            LUT += 32;
-        }
-
-        for (int b = 0; b < BB; b++) {
-            simd32uint8 c = simd32uint8(codes);
-            codes += 32;
-            simd32uint8 mask(15);
-            simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
-            simd32uint8 clo = c & mask;
-
-            for (int q = 0; q < NQ; q++) {
-                simd32uint8 lut = lut_cache[q];
-                simd32uint8 res0 = lut.lookup_2_lanes(clo);
-                simd32uint8 res1 = lut.lookup_2_lanes(chi);
-
-                accu[q][b][0] += simd16uint16(res0);
-                accu[q][b][1] += simd16uint16(res0) >> 8;
-
-                accu[q][b][2] += simd16uint16(res1);
-                accu[q][b][3] += simd16uint16(res1) >> 8;
-            }
-        }
-    }
-
-    for (int sq = 0; sq < scaler.nscale; sq += 2) {
-        simd32uint8 lut_cache[NQ];
-        for (int q = 0; q < NQ; q++) {
-            lut_cache[q] = simd32uint8(LUT);
-            LUT += 32;
-        }
-
-        for (int b = 0; b < BB; b++) {
-            simd32uint8 c = simd32uint8(codes);
-            codes += 32;
-            simd32uint8 mask(15);
-            simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
-            simd32uint8 clo = c & mask;
-
-            for (int q = 0; q < NQ; q++) {
-                simd32uint8 lut = lut_cache[q];
-
-                simd32uint8 res0 = scaler.lookup(lut, clo);
-                accu[q][b][0] += scaler.scale_lo(res0); // handle vectors 0..7
-                accu[q][b][1] += scaler.scale_hi(res0); // handle vectors 8..15
-
-                simd32uint8 res1 = scaler.lookup(lut, chi);
-                accu[q][b][2] += scaler.scale_lo(res1); // handle vectors 16..23
-                accu[q][b][3] +=
-                        scaler.scale_hi(res1); //  handle vectors 24..31
-            }
-        }
-    }
-
-    for (int q = 0; q < NQ; q++) {
-        for (int b = 0; b < BB; b++) {
-            accu[q][b][0] -= accu[q][b][1] << 8;
-            simd16uint16 dis0 = combine2x2(accu[q][b][0], accu[q][b][1]);
-
-            accu[q][b][2] -= accu[q][b][3] << 8;
-            simd16uint16 dis1 = combine2x2(accu[q][b][2], accu[q][b][3]);
-
-            res.handle(q, b, dis0, dis1);
-        }
-    }
-}
-
-template <int NQ, int BB, class ResultHandler, class Scaler>
-void accumulate_fixed_blocks(
-        size_t nb,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        ResultHandler& res,
-        const Scaler& scaler) {
-    constexpr int bbs = 32 * BB;
-    for (size_t j0 = 0; j0 < nb; j0 += bbs) {
-        FixedStorageHandler<NQ, 2 * BB> res2;
-        kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2, scaler);
-        res.set_block_origin(0, j0);
-        res2.to_other_handler(res);
-        codes += bbs * nsq / 2;
-    }
-}
-
-template <class ResultHandler, class Scaler>
-void pq4_accumulate_loop_fixed_scaler(
-        int nq,
-        size_t nb,
-        int bbs,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        ResultHandler& res,
-        const Scaler& scaler) {
-    FAISS_THROW_IF_NOT(is_aligned_pointer(codes));
-    FAISS_THROW_IF_NOT(is_aligned_pointer(LUT));
-    FAISS_THROW_IF_NOT(bbs % 32 == 0);
-    FAISS_THROW_IF_NOT(nb % bbs == 0);
-
-#define DISPATCH(NQ, BB)                                                   \
-    case NQ * 1000 + BB:                                                   \
-        accumulate_fixed_blocks<NQ, BB>(nb, nsq, codes, LUT, res, scaler); \
-        break
-
-    switch (nq * 1000 + bbs / 32) {
-        DISPATCH(1, 1);
-        DISPATCH(1, 2);
-        DISPATCH(1, 3);
-        DISPATCH(1, 4);
-        DISPATCH(1, 5);
-        DISPATCH(2, 1);
-        DISPATCH(2, 2);
-        DISPATCH(3, 1);
-        DISPATCH(4, 1);
-        default:
-            FAISS_THROW_FMT("nq=%d bbs=%d not instantiated", nq, bbs);
-    }
-#undef DISPATCH
-}
-
-template <class ResultHandler>
-void pq4_accumulate_loop_fixed_handler(
-        int nq,
-        size_t nb,
-        int bbs,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        ResultHandler& res,
-        const NormTableScaler* scaler) {
-    if (scaler) {
-        pq4_accumulate_loop_fixed_scaler(
-                nq, nb, bbs, nsq, codes, LUT, res, *scaler);
-    } else {
-        DummyScaler dscaler;
-        pq4_accumulate_loop_fixed_scaler(
-                nq, nb, bbs, nsq, codes, LUT, res, dscaler);
-    }
-}
-
-struct Run_pq4_accumulate_loop {
-    template <class ResultHandler>
-    void f(ResultHandler& res,
-           int nq,
-           size_t nb,
-           int bbs,
-           int nsq,
-           const uint8_t* codes,
-           const uint8_t* LUT,
-           const NormTableScaler* scaler) {
-        pq4_accumulate_loop_fixed_handler(
-                nq, nb, bbs, nsq, codes, LUT, res, scaler);
-    }
-};
-
-} // anonymous namespace
-
-void pq4_accumulate_loop(
-        int nq,
-        size_t nb,
-        int bbs,
-        int nsq,
-        const uint8_t* codes,
-        const uint8_t* LUT,
-        SIMDResultHandler& res,
-        const NormTableScaler* scaler) {
-    Run_pq4_accumulate_loop consumer;
-    dispatch_SIMDResultHandler(
-            res, consumer, nq, nb, bbs, nsq, codes, LUT, scaler);
-}
-
-} // namespace faiss

data/vendor/faiss/faiss/utils/approx_topk/approx_topk.h

@@ -1,84 +0,0 @@
-/*
- * 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.
- */
-
-// This file contains an implementation of approximate top-k search
-// using heap. It was initially created for a beam search.
-//
-// The core idea is the following.
-// Say we need to find beam_size indices with the minimal distance
-// values. It is done via heap (priority_queue) using the following
-// pseudocode:
-//
-//   def baseline():
-//     distances = np.empty([beam_size * n], dtype=float)
-//     indices = np.empty([beam_size * n], dtype=int)
-//
-//     heap = Heap(max_heap_size=beam_size)
-//
-//     for i in range(0, beam_size * n):
-//         heap.push(distances[i], indices[i])
-//
-// Basically, this is what heap_addn() function from utils/Heap.h does.
-//
-// The following scheme can be used for approximate beam search.
-// Say, we need to find elements with min distance.
-// Basically, we split n elements of every beam into NBUCKETS buckets
-// and track the index with the minimal distance for every bucket.
-// This can be effectively SIMD-ed and significantly lowers the number
-// of operations, but yields approximate results for beam_size >= 2.
-//
-//  def approximate_v1():
-//    distances = np.empty([beam_size * n], dtype=float)
-//    indices = np.empty([beam_size * n], dtype=int)
-//
-//    heap = Heap(max_heap_size=beam_size)
-//
-//    for beam in range(0, beam_size):
-//      # The value of 32 is just an example.
-//      # The value may be varied: the larger the value is,
-//      #  the slower and the more precise vs baseline beam search is
-//      NBUCKETS = 32
-//
-//     local_min_distances = [HUGE_VALF] * NBUCKETS
-//     local_min_indices = [0] * NBUCKETS
-//
-//      for i in range(0, n / NBUCKETS):
-//        for j in range(0, NBUCKETS):
-//          idx = beam * n + i * NBUCKETS + j
-//          if distances[idx] < local_min_distances[j]:
-//            local_min_distances[j] = distances[idx]
-//            local_min_indices[j] = indices[idx]
-//
-//    for j in range(0, NBUCKETS):
-//      heap.push(local_min_distances[j], local_min_indices[j])
-//
-// The accuracy can be improved by tracking min-2 elements for every
-// bucket. Such a min-2 implementation with NBUCKETS buckets provides
-// better accuracy than top-1 implementation with 2 * NBUCKETS buckets.
-// Min-3 is also doable. One can use min-N approach, but I'm not sure
-// whether min-4 and above are practical, because of the lack of SIMD
-// registers (unless AVX-512 version is used).
-//
-// C++ template for top-N implementation is provided. The code
-// assumes that indices[idx] == idx. One can write a code that lifts
-// such an assumption easily.
-//
-// Currently, the code that tracks elements with min distances is implemented
-//    (Max Heap). Min Heap option can be added easily.
-
-#pragma once
-
-#include <faiss/impl/platform_macros.h>
-
-// the list of available modes is in the following file
-#include <faiss/utils/approx_topk/mode.h>
-
-#ifdef __AVX2__
-#include <faiss/utils/approx_topk/avx2-inl.h>
-#else
-#include <faiss/utils/approx_topk/generic.h>
-#endif