faiss 0.6.0 → 0.6.1

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

@@ -1,230 +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,
-        size_t block_stride) {
-    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 += block_stride;
-    }
-}
-
-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,
-        size_t block_stride) {
-    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, block_stride); \
-        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,
-        size_t block_stride) {
-    if (scaler) {
-        pq4_accumulate_loop_fixed_scaler(
-                nq, nb, bbs, nsq, codes, LUT, res, *scaler, block_stride);
-    } else {
-        DummyScaler dscaler;
-        pq4_accumulate_loop_fixed_scaler(
-                nq, nb, bbs, nsq, codes, LUT, res, dscaler, block_stride);
-    }
-}
-
-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,
-           size_t block_stride) {
-        pq4_accumulate_loop_fixed_handler(
-                nq, nb, bbs, nsq, codes, LUT, res, scaler, block_stride);
-    }
-};
-
-} // 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,
-        size_t block_stride) {
-    Run_pq4_accumulate_loop consumer;
-    dispatch_SIMDResultHandler(
-            res, consumer, nq, nb, bbs, nsq, codes, LUT, scaler, block_stride);
-}
-
-} // 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

data/vendor/faiss/faiss/utils/approx_topk/avx2-inl.h

@@ -1,196 +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 <immintrin.h>
-
-#include <limits>
-
-#include <faiss/impl/FaissAssert.h>
-#include <faiss/utils/Heap.h>
-
-namespace faiss {
-
-template <typename C, uint32_t NBUCKETS, uint32_t N>
-struct HeapWithBuckets {
-    // this case was not implemented yet.
-};
-
-template <uint32_t NBUCKETS, uint32_t N>
-struct HeapWithBuckets<CMax<float, int>, NBUCKETS, N> {
-    static constexpr uint32_t NBUCKETS_8 = NBUCKETS / 8;
-    static_assert(
-            (NBUCKETS) > 0 && ((NBUCKETS % 8) == 0),
-            "Number of buckets needs to be 8, 16, 24, ...");
-
-    static void addn(
-            // number of elements
-            const uint32_t n,
-            // distances. It is assumed to have n elements.
-            const float* const __restrict distances,
-            // number of best elements to keep
-            const uint32_t k,
-            // output distances
-            float* const __restrict bh_val,
-            // output indices, each being within [0, n) range
-            int32_t* const __restrict bh_ids) {
-        // forward a call to bs_addn with 1 beam
-        bs_addn(1, n, distances, k, bh_val, bh_ids);
-    }
-
-    static void bs_addn(
-            // beam_size parameter of Beam Search algorithm
-            const uint32_t beam_size,
-            // number of elements per beam
-            const uint32_t n_per_beam,
-            // distances. It is assumed to have (n_per_beam * beam_size)
-            // elements.
-            const float* const __restrict distances,
-            // number of best elements to keep
-            const uint32_t k,
-            // output distances
-            float* const __restrict bh_val,
-            // output indices, each being within [0, n_per_beam * beam_size)
-            // range
-            int32_t* const __restrict bh_ids) {
-        // // Basically, the function runs beam_size iterations.
-        // // Every iteration NBUCKETS * N elements are added to a regular heap.
-        // // So, maximum number of added elements is beam_size * NBUCKETS * N.
-        // // This number is expected to be less or equal than k.
-        // FAISS_THROW_IF_NOT_FMT(
-        //         beam_size * NBUCKETS * N >= k,
-        //         "Cannot pick %d elements, only %d. "
-        //         "Check the function and template arguments values.",
-        //         k,
-        //         beam_size * NBUCKETS * N);
-
-        using C = CMax<float, int>;
-
-        // main loop
-        for (uint32_t beam_index = 0; beam_index < beam_size; beam_index++) {
-            __m256 min_distances_i[NBUCKETS_8][N];
-            __m256i min_indices_i[NBUCKETS_8][N];
-
-            for (uint32_t j = 0; j < NBUCKETS_8; j++) {
-                for (uint32_t p = 0; p < N; p++) {
-                    min_distances_i[j][p] =
-                            _mm256_set1_ps(std::numeric_limits<float>::max());
-                    min_indices_i[j][p] =
-                            _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
-                }
-            }
-
-            __m256i current_indices = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
-            __m256i indices_delta = _mm256_set1_epi32(NBUCKETS);
-
-            const uint32_t nb = (n_per_beam / NBUCKETS) * NBUCKETS;
-
-            // put the data into buckets
-            for (uint32_t ip = 0; ip < nb; ip += NBUCKETS) {
-                for (uint32_t j = 0; j < NBUCKETS_8; j++) {
-                    const __m256 distances_reg = _mm256_loadu_ps(
-                            distances + j * 8 + ip + n_per_beam * beam_index);
-
-                    // loop. Compiler should get rid of unneeded ops
-                    __m256 distance_candidate = distances_reg;
-                    __m256i indices_candidate = current_indices;
-
-                    for (uint32_t p = 0; p < N; p++) {
-                        const __m256 comparison = _mm256_cmp_ps(
-                                min_distances_i[j][p],
-                                distance_candidate,
-                                _CMP_LE_OS);
-
-                        // // blend seems to be slower than min
-                        // const __m256 min_distances_new = _mm256_blendv_ps(
-                        //         distance_candidate,
-                        //         min_distances_i[j][p],
-                        //         comparison);
-                        const __m256 min_distances_new = _mm256_min_ps(
-                                distance_candidate, min_distances_i[j][p]);
-                        const __m256i min_indices_new =
-                                _mm256_castps_si256(_mm256_blendv_ps(
-                                        _mm256_castsi256_ps(indices_candidate),
-                                        _mm256_castsi256_ps(
-                                                min_indices_i[j][p]),
-                                        comparison));
-
-                        // // blend seems to be slower than min
-                        // const __m256 max_distances_new = _mm256_blendv_ps(
-                        //         min_distances_i[j][p],
-                        //         distance_candidate,
-                        //         comparison);
-                        const __m256 max_distances_new = _mm256_max_ps(
-                                min_distances_i[j][p], distances_reg);
-                        const __m256i max_indices_new =
-                                _mm256_castps_si256(_mm256_blendv_ps(
-                                        _mm256_castsi256_ps(
-                                                min_indices_i[j][p]),
-                                        _mm256_castsi256_ps(indices_candidate),
-                                        comparison));
-
-                        distance_candidate = max_distances_new;
-                        indices_candidate = max_indices_new;
-
-                        min_distances_i[j][p] = min_distances_new;
-                        min_indices_i[j][p] = min_indices_new;
-                    }
-                }
-
-                current_indices =
-                        _mm256_add_epi32(current_indices, indices_delta);
-            }
-
-            // fix the indices
-            for (uint32_t j = 0; j < NBUCKETS_8; j++) {
-                const __m256i offset =
-                        _mm256_set1_epi32(n_per_beam * beam_index + j * 8);
-                for (uint32_t p = 0; p < N; p++) {
-                    min_indices_i[j][p] =
-                            _mm256_add_epi32(min_indices_i[j][p], offset);
-                }
-            }
-
-            // merge every bucket into the regular heap
-            for (uint32_t p = 0; p < N; p++) {
-                for (uint32_t j = 0; j < NBUCKETS_8; j++) {
-                    int32_t min_indices_scalar[8];
-                    float min_distances_scalar[8];
-
-                    _mm256_storeu_si256(
-                            (__m256i*)min_indices_scalar, min_indices_i[j][p]);
-                    _mm256_storeu_ps(
-                            min_distances_scalar, min_distances_i[j][p]);
-
-                    // this exact way is needed to maintain the order as if the
-                    // input elements were pushed to the heap sequentially
-                    for (size_t j8 = 0; j8 < 8; j8++) {
-                        const auto value = min_distances_scalar[j8];
-                        const auto index = min_indices_scalar[j8];
-                        if (C::cmp2(bh_val[0], value, bh_ids[0], index)) {
-                            heap_replace_top<C>(
-                                    k, bh_val, bh_ids, value, index);
-                        }
-                    }
-                }
-            }
-
-            // process leftovers
-            for (uint32_t ip = nb; ip < n_per_beam; ip++) {
-                const int32_t index = ip + n_per_beam * beam_index;
-                const float value = distances[index];
-
-                if (C::cmp(bh_val[0], value)) {
-                    heap_replace_top<C>(k, bh_val, bh_ids, value, index);
-                }
-            }
-        }
-    }
-};
-
-} // namespace faiss

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

@@ -1,34 +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
-
-/// Represents the mode of use of approximate top-k computations
-/// that allows to trade accuracy vs speed. So, every options
-/// besides EXACT_TOPK increases the speed.
-///
-/// B represents the number of buckets.
-/// D is the number of min-k elements to track within every bucket.
-///
-/// Default option is EXACT_TOPK.
-/// APPROX_TOPK_BUCKETS_B16_D2 is worth starting from, if you'd like
-/// to experiment a bit.
-///
-/// It seems that only the limited number of combinations are
-/// meaningful, because of the limited supply of SIMD registers.
-/// Also, certain combinations, such as B32_D1 and B16_D1, were concluded
-/// to be not very precise in benchmarks, so they were not introduced.
-///
-/// TODO: Consider d-ary SIMD heap.
-
-enum ApproxTopK_mode_t : int {
-    EXACT_TOPK = 0,
-    APPROX_TOPK_BUCKETS_B32_D2 = 1,
-    APPROX_TOPK_BUCKETS_B8_D3 = 2,
-    APPROX_TOPK_BUCKETS_B16_D2 = 3,
-    APPROX_TOPK_BUCKETS_B8_D2 = 4,
-};

data/vendor/faiss/faiss/utils/distances_fused/avx512.h

@@ -1,36 +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.
- */
-
-// AVX512 might be not used, but this version provides ~2x speedup
-// over AVX2 kernel, say, for training PQx10 or PQx12, and speeds up
-// additional cases with larger dimensionalities.
-
-#pragma once
-
-#include <faiss/impl/ResultHandler.h>
-#include <faiss/impl/platform_macros.h>
-
-#include <faiss/utils/Heap.h>
-
-#ifdef __AVX512F__
-
-namespace faiss {
-
-// Returns true if the fused kernel is available and the data was processed.
-// Returns false if the fused kernel is not available.
-bool exhaustive_L2sqr_fused_cmax_AVX512(
-        const float* x,
-        const float* y,
-        size_t d,
-        size_t nx,
-        size_t ny,
-        Top1BlockResultHandler<CMax<float, int64_t>>& res,
-        const float* y_norms);
-
-} // namespace faiss
-
-#endif