diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 3616986c0..aeef0ed6c 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -12,6 +12,7 @@ set(FAISS_SIMD_AVX2_SRC impl/fast_scan/impl-avx2.cpp impl/pq_code_distance/pq_code_distance-avx2.cpp impl/scalar_quantizer/sq-avx2.cpp + impl/approx_topk/avx2.cpp utils/simd_impl/distances_avx2.cpp ) set(FAISS_SIMD_AVX512_SRC @@ -23,6 +24,7 @@ set(FAISS_SIMD_AVX512_SRC set(FAISS_SIMD_NEON_SRC impl/fast_scan/impl-neon.cpp impl/scalar_quantizer/sq-neon.cpp + impl/approx_topk/neon.cpp utils/simd_impl/distances_aarch64.cpp ) set(FAISS_SIMD_SVE_SRC @@ -315,10 +317,11 @@ set(FAISS_HEADERS utils/distances_fused/avx512.h utils/distances_fused/distances_fused.h utils/distances_fused/simdlib_based.h - utils/approx_topk/approx_topk.h - utils/approx_topk/avx2-inl.h - utils/approx_topk/generic.h - utils/approx_topk/mode.h + impl/approx_topk/approx_topk.h + impl/approx_topk/generic.h + impl/approx_topk/rq_beam_search_tab.h + impl/approx_topk/rq_beam_search_tab-inl.h + impl/approx_topk/simdlib256-inl.h utils/approx_topk_hamming/approx_topk_hamming.h utils/transpose/transpose-avx2-inl.h utils/transpose/transpose-avx512-inl.h diff --git a/faiss/IndexBinaryFlat.h b/faiss/IndexBinaryFlat.h index 0ce43f3e9..b3d88b9f8 100644 --- a/faiss/IndexBinaryFlat.h +++ b/faiss/IndexBinaryFlat.h @@ -14,8 +14,8 @@ #include +#include #include -#include namespace faiss { diff --git a/faiss/impl/LocalSearchQuantizer.cpp b/faiss/impl/LocalSearchQuantizer.cpp index 7ed2b8c88..f36653e35 100644 --- a/faiss/impl/LocalSearchQuantizer.cpp +++ b/faiss/impl/LocalSearchQuantizer.cpp @@ -22,7 +22,7 @@ #include #include -#include +#include // this is needed for prefetching diff --git a/faiss/impl/ResidualQuantizer.h b/faiss/impl/ResidualQuantizer.h index 6a70b822e..73132d259 100644 --- a/faiss/impl/ResidualQuantizer.h +++ b/faiss/impl/ResidualQuantizer.h @@ -13,7 +13,7 @@ #include #include -#include +#include namespace faiss { diff --git a/faiss/impl/approx_topk/approx_topk.h b/faiss/impl/approx_topk/approx_topk.h new file mode 100644 index 000000000..5d7550278 --- /dev/null +++ b/faiss/impl/approx_topk/approx_topk.h @@ -0,0 +1,215 @@ +/* + * 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. + */ + +// Approximate top-k search using bucketed heaps with SIMD acceleration. +// +// HeapWithBuckets splits n elements into NBUCKETS buckets, +// tracks the top-N per bucket using SIMD, then merges into a regular heap. +// This trades a small accuracy loss for significant speedup. +// +// 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. +// +// Dispatch: +// AVX2 / ARM_NEON → HeapWithBucketsCMaxFloat (simdlib256-inl.h) +// NONE (scalar) → HeapWithBucketsGenericCMaxFloat (generic.h) +// +// The SIMD definitions live in simdlib256-inl.h (only included by per-ISA +// .cpp files). Common TUs see only declarations here, so no extern template +// suppression is needed. + +#pragma once + +#include + +#include +#include +#include +#include +#include + +// ----------------------------------------------------------------------- +// ApproxTopK_mode_t +// ----------------------------------------------------------------------- + +/// 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. + +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, +}; + +// ----------------------------------------------------------------------- +// HeapWithBuckets dispatch +// ----------------------------------------------------------------------- + +namespace faiss { + +// Primary template: declared, bs_addn NOT defined here. +// The out-of-line definition lives in simdlib256-inl.h, which is only +// included by the per-ISA .cpp files (avx2.cpp, neon.cpp). +template +struct HeapWithBucketsCMaxFloat { + static_assert( + (NBUCKETS) > 0 && ((NBUCKETS % 8) == 0), + "Number of buckets needs to be 8, 16, 24, ..."); + + static void addn( + const uint32_t n, + const float* const __restrict distances, + const uint32_t k, + float* const __restrict bh_val, + int32_t* const __restrict bh_ids) { + bs_addn(1, n, distances, k, bh_val, bh_ids); + } + + // Declared but not defined — resolved at link time from avx2.cpp/neon.cpp. + static void bs_addn( + const uint32_t beam_size, + const uint32_t n_per_beam, + const float* const __restrict distances, + const uint32_t k, + float* const __restrict bh_val, + int32_t* const __restrict bh_ids); +}; + +// NONE specialization: delegates to the scalar generic implementation. +template +struct HeapWithBucketsCMaxFloat { + static void addn( + const uint32_t n, + const float* const __restrict distances, + const uint32_t k, + float* const __restrict bh_val, + int32_t* const __restrict bh_ids) { + bs_addn(1, n, distances, k, bh_val, bh_ids); + } + + static void bs_addn( + const uint32_t beam_size, + const uint32_t n_per_beam, + const float* const __restrict distances, + const uint32_t k, + float* const __restrict bh_val, + int32_t* const __restrict bh_ids) { + HeapWithBucketsGenericCMaxFloat::bs_addn( + beam_size, n_per_beam, distances, k, bh_val, bh_ids); + } +}; + +// Primary template — not implemented for arbitrary comparators. +template +struct HeapWithBuckets { + static_assert( + sizeof(C) == 0, + "HeapWithBuckets: unsupported comparator type"); +}; + +// Partial specialization for CMax that dispatches +// via with_simd_level_256bit. +template +struct HeapWithBuckets, NBUCKETS, N> { + static void addn( + const uint32_t n, + const float* const __restrict distances, + const uint32_t k, + float* const __restrict bh_val, + int32_t* const __restrict bh_ids) { + bs_addn(1, n, distances, k, bh_val, bh_ids); + } + + static void bs_addn( + const uint32_t beam_size, + const uint32_t n_per_beam, + const float* const __restrict distances, + const uint32_t k, + float* const __restrict bh_val, + int32_t* const __restrict bh_ids) { + with_simd_level_256bit([&]() { + HeapWithBucketsCMaxFloat::bs_addn( + beam_size, n_per_beam, distances, k, bh_val, bh_ids); + }); + } +}; + +} // namespace faiss diff --git a/faiss/impl/approx_topk/avx2.cpp b/faiss/impl/approx_topk/avx2.cpp new file mode 100644 index 000000000..85022e8d4 --- /dev/null +++ b/faiss/impl/approx_topk/avx2.cpp @@ -0,0 +1,68 @@ +/* + * 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. + */ + +// Explicit template instantiations of HeapWithBucketsCMaxFloat and +// accum_and_*_tab for SIMDLevel::AVX2. + +#ifdef COMPILE_SIMD_AVX2 + +#include +#include +#include + +namespace faiss { + +template struct HeapWithBucketsCMaxFloat<8, 3, SIMDLevel::AVX2>; +template struct HeapWithBucketsCMaxFloat<8, 2, SIMDLevel::AVX2>; +template struct HeapWithBucketsCMaxFloat<16, 2, SIMDLevel::AVX2>; +template struct HeapWithBucketsCMaxFloat<16, 1, SIMDLevel::AVX2>; +template struct HeapWithBucketsCMaxFloat<32, 2, SIMDLevel::AVX2>; + +#define INSTANTIATE_ACCUM_TAB(M) \ + template void accum_and_store_tab( \ + size_t, \ + const float* __restrict, \ + const uint64_t* __restrict, \ + const int32_t* __restrict, \ + size_t, \ + size_t, \ + size_t, \ + float* __restrict); \ + template void accum_and_add_tab( \ + size_t, \ + const float* __restrict, \ + const uint64_t* __restrict, \ + const int32_t* __restrict, \ + size_t, \ + size_t, \ + size_t, \ + float* __restrict); \ + template void accum_and_finalize_tab( \ + const float* __restrict, \ + const uint64_t* __restrict, \ + const int32_t* __restrict, \ + size_t, \ + size_t, \ + size_t, \ + const float* __restrict, \ + const float* __restrict, \ + float* __restrict); + +INSTANTIATE_ACCUM_TAB(1) +INSTANTIATE_ACCUM_TAB(2) +INSTANTIATE_ACCUM_TAB(3) +INSTANTIATE_ACCUM_TAB(4) +INSTANTIATE_ACCUM_TAB(5) +INSTANTIATE_ACCUM_TAB(6) +INSTANTIATE_ACCUM_TAB(7) +INSTANTIATE_ACCUM_TAB(8) + +#undef INSTANTIATE_ACCUM_TAB + +} // namespace faiss + +#endif // COMPILE_SIMD_AVX2 diff --git a/faiss/utils/approx_topk/generic.h b/faiss/impl/approx_topk/generic.h similarity index 89% rename from faiss/utils/approx_topk/generic.h rename to faiss/impl/approx_topk/generic.h index 5bc813d4d..cad6391a8 100644 --- a/faiss/utils/approx_topk/generic.h +++ b/faiss/impl/approx_topk/generic.h @@ -16,16 +16,12 @@ namespace faiss { -// This is the implementation of the idea and it is very slow, -// because a compiler is unable to vectorize it properly. - -template -struct HeapWithBuckets { - // this case was not implemented yet. -}; +// Scalar (generic) implementation of HeapWithBuckets. +// This is correct but slow because a compiler is unable to +// vectorize it properly. Used as the SIMDLevel::NONE fallback. template -struct HeapWithBuckets, NBUCKETS, N> { +struct HeapWithBucketsGenericCMaxFloat { static void addn( // number of elements const uint32_t n, @@ -135,4 +131,15 @@ struct HeapWithBuckets, NBUCKETS, N> { } }; +// Legacy name kept for backward compatibility (used when +// approx_topk.h is not included). +template +struct HeapWithBucketsGeneric { + // not implemented for arbitrary C +}; + +template +struct HeapWithBucketsGeneric, NBUCKETS, N> + : HeapWithBucketsGenericCMaxFloat {}; + } // namespace faiss diff --git a/faiss/impl/approx_topk/neon.cpp b/faiss/impl/approx_topk/neon.cpp new file mode 100644 index 000000000..81ebdd2cf --- /dev/null +++ b/faiss/impl/approx_topk/neon.cpp @@ -0,0 +1,68 @@ +/* + * 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. + */ + +// Explicit template instantiations of HeapWithBucketsCMaxFloat and +// accum_and_*_tab for SIMDLevel::ARM_NEON. + +#ifdef COMPILE_SIMD_ARM_NEON + +#include +#include +#include + +namespace faiss { + +template struct HeapWithBucketsCMaxFloat<8, 3, SIMDLevel::ARM_NEON>; +template struct HeapWithBucketsCMaxFloat<8, 2, SIMDLevel::ARM_NEON>; +template struct HeapWithBucketsCMaxFloat<16, 2, SIMDLevel::ARM_NEON>; +template struct HeapWithBucketsCMaxFloat<16, 1, SIMDLevel::ARM_NEON>; +template struct HeapWithBucketsCMaxFloat<32, 2, SIMDLevel::ARM_NEON>; + +#define INSTANTIATE_ACCUM_TAB(M) \ + template void accum_and_store_tab( \ + size_t, \ + const float* __restrict, \ + const uint64_t* __restrict, \ + const int32_t* __restrict, \ + size_t, \ + size_t, \ + size_t, \ + float* __restrict); \ + template void accum_and_add_tab( \ + size_t, \ + const float* __restrict, \ + const uint64_t* __restrict, \ + const int32_t* __restrict, \ + size_t, \ + size_t, \ + size_t, \ + float* __restrict); \ + template void accum_and_finalize_tab( \ + const float* __restrict, \ + const uint64_t* __restrict, \ + const int32_t* __restrict, \ + size_t, \ + size_t, \ + size_t, \ + const float* __restrict, \ + const float* __restrict, \ + float* __restrict); + +INSTANTIATE_ACCUM_TAB(1) +INSTANTIATE_ACCUM_TAB(2) +INSTANTIATE_ACCUM_TAB(3) +INSTANTIATE_ACCUM_TAB(4) +INSTANTIATE_ACCUM_TAB(5) +INSTANTIATE_ACCUM_TAB(6) +INSTANTIATE_ACCUM_TAB(7) +INSTANTIATE_ACCUM_TAB(8) + +#undef INSTANTIATE_ACCUM_TAB + +} // namespace faiss + +#endif // COMPILE_SIMD_ARM_NEON diff --git a/faiss/impl/approx_topk/rq_beam_search_tab-inl.h b/faiss/impl/approx_topk/rq_beam_search_tab-inl.h new file mode 100644 index 000000000..867b9ba2b --- /dev/null +++ b/faiss/impl/approx_topk/rq_beam_search_tab-inl.h @@ -0,0 +1,169 @@ +/* + * 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. + */ + +// Definitions of the SIMDLevel-templatized accum_and_*_tab functions. +// Only included by per-ISA .cpp files (avx2.cpp, neon.cpp). +// Do NOT include this from common translation units. +// +// Common TUs include rq_beam_search_tab.h (declarations only). + +#pragma once + +#include +#include + +#include +#include + +namespace faiss { + +template +void accum_and_store_tab( + const size_t m_offset, + const float* const __restrict codebook_cross_norms, + const uint64_t* const __restrict codebook_offsets, + const int32_t* const __restrict codes_i, + const size_t b, + const size_t ldc, + const size_t K, + float* const __restrict output) { + using simd_float = simd8float32_tpl; + + const float* cbs[M]; + for (size_t ij = 0; ij < M; ij++) { + const size_t code = static_cast(codes_i[b * m_offset + ij]); + cbs[ij] = &codebook_cross_norms[(codebook_offsets[ij] + code) * ldc]; + } + + const size_t K8 = (K / (8 * NK)) * (8 * NK); + + for (size_t kk = 0; kk < K8; kk += 8 * NK) { + simd_float regs[NK]; + for (size_t ik = 0; ik < NK; ik++) { + regs[ik] = simd_float(cbs[0] + kk + ik * 8); + } + + for (size_t ij = 1; ij < M; ij++) { + for (size_t ik = 0; ik < NK; ik++) { + regs[ik] += simd_float(cbs[ij] + kk + ik * 8); + } + } + + for (size_t ik = 0; ik < NK; ik++) { + regs[ik].storeu(output + kk + ik * 8); + } + } + + for (size_t kk = K8; kk < K; kk++) { + float reg = cbs[0][kk]; + for (size_t ij = 1; ij < M; ij++) { + reg += cbs[ij][kk]; + } + output[kk] = reg; + } +} + +template +void accum_and_add_tab( + const size_t m_offset, + const float* const __restrict codebook_cross_norms, + const uint64_t* const __restrict codebook_offsets, + const int32_t* const __restrict codes_i, + const size_t b, + const size_t ldc, + const size_t K, + float* const __restrict output) { + using simd_float = simd8float32_tpl; + + const float* cbs[M]; + for (size_t ij = 0; ij < M; ij++) { + const size_t code = static_cast(codes_i[b * m_offset + ij]); + cbs[ij] = &codebook_cross_norms[(codebook_offsets[ij] + code) * ldc]; + } + + const size_t K8 = (K / (8 * NK)) * (8 * NK); + + for (size_t kk = 0; kk < K8; kk += 8 * NK) { + simd_float regs[NK]; + for (size_t ik = 0; ik < NK; ik++) { + regs[ik] = simd_float(cbs[0] + kk + ik * 8); + } + + for (size_t ij = 1; ij < M; ij++) { + for (size_t ik = 0; ik < NK; ik++) { + regs[ik] += simd_float(cbs[ij] + kk + ik * 8); + } + } + + for (size_t ik = 0; ik < NK; ik++) { + simd_float existing(output + kk + ik * 8); + existing += regs[ik]; + existing.storeu(output + kk + ik * 8); + } + } + + for (size_t kk = K8; kk < K; kk++) { + float reg = cbs[0][kk]; + for (size_t ij = 1; ij < M; ij++) { + reg += cbs[ij][kk]; + } + output[kk] += reg; + } +} + +template +void accum_and_finalize_tab( + const float* const __restrict codebook_cross_norms, + const uint64_t* const __restrict codebook_offsets, + const int32_t* const __restrict codes_i, + const size_t b, + const size_t ldc, + const size_t K, + const float* const __restrict distances_i, + const float* const __restrict cd_common, + float* const __restrict output) { + using simd_float = simd8float32_tpl; + + const float* cbs[M]; + for (size_t ij = 0; ij < M; ij++) { + const size_t code = static_cast(codes_i[b * M + ij]); + cbs[ij] = &codebook_cross_norms[(codebook_offsets[ij] + code) * ldc]; + } + + const size_t K8 = (K / (8 * NK)) * (8 * NK); + + for (size_t kk = 0; kk < K8; kk += 8 * NK) { + simd_float regs[NK]; + for (size_t ik = 0; ik < NK; ik++) { + regs[ik] = simd_float(cbs[0] + kk + ik * 8); + } + + for (size_t ij = 1; ij < M; ij++) { + for (size_t ik = 0; ik < NK; ik++) { + regs[ik] += simd_float(cbs[ij] + kk + ik * 8); + } + } + + simd_float two(2.0f); + for (size_t ik = 0; ik < NK; ik++) { + simd_float common_v(cd_common + kk + ik * 8); + common_v = fmadd(two, regs[ik], common_v); + common_v += simd_float(distances_i[b]); + common_v.storeu(output + b * K + kk + ik * 8); + } + } + + for (size_t kk = K8; kk < K; kk++) { + float reg = cbs[0][kk]; + for (size_t ij = 1; ij < M; ij++) { + reg += cbs[ij][kk]; + } + output[b * K + kk] = distances_i[b] + cd_common[kk] + 2 * reg; + } +} + +} // namespace faiss diff --git a/faiss/impl/approx_topk/rq_beam_search_tab.h b/faiss/impl/approx_topk/rq_beam_search_tab.h new file mode 100644 index 000000000..40f92daf9 --- /dev/null +++ b/faiss/impl/approx_topk/rq_beam_search_tab.h @@ -0,0 +1,117 @@ +/* + * 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. + */ + +/// @file rq_beam_search_tab.h +/// @brief Declarations for SIMDLevel-templatized codebook accumulation +/// functions. +/// +/// These functions accumulate codebook cross-norm tables for beam search +/// encoding in the Residual Quantizer. They compute the distance +/// contributions from previously encoded codebooks using SIMD-accelerated +/// register accumulation. +/// +/// Definitions are in rq_beam_search_tab-inl.h (only included by per-ISA +/// .cpp files). The common TU only sees these declarations, so no extern +/// template suppression is needed — the linker resolves to the explicit +/// instantiations in avx2.cpp / neon.cpp. + +#pragma once + +#include +#include + +#include + +namespace faiss { + +/// Accumulate cross-norms for M codebooks and store the result. +/// +/// Loads M codebook rows (selected by codes_i) and sums them using +/// NK×8-wide SIMD chunks, writing the result to output. Used to +/// initialize the temporary buffer in the m≥8 path. +/// +/// @tparam M number of codebook rows to accumulate +/// @tparam NK number of 8-float SIMD chunks per loop iteration +/// @tparam SL SIMD level (AVX2, ARM_NEON, etc.) +/// @param m_offset stride between beam entries in codes_i +/// @param codebook_cross_norms cross-norm table, shape (total_codes, ldc) +/// @param codebook_offsets per-codebook offset into cross-norm table +/// @param codes_i code indices for the current query +/// @param b beam index +/// @param ldc leading dimension of cross-norm table (≥ K) +/// @param K number of centroids in the current codebook +/// @param output output buffer, size K (overwritten) +template +void accum_and_store_tab( + size_t m_offset, + const float* __restrict codebook_cross_norms, + const uint64_t* __restrict codebook_offsets, + const int32_t* __restrict codes_i, + size_t b, + size_t ldc, + size_t K, + float* __restrict output); + +/// Accumulate cross-norms for M codebooks and add to existing output. +/// +/// Like accum_and_store_tab, but adds the accumulated result to the +/// existing values in output (output[k] += sum). Used for subsequent +/// chunks of 8 codebooks in the m≥8 path. +/// +/// @tparam M number of codebook rows to accumulate +/// @tparam NK number of 8-float SIMD chunks per loop iteration +/// @tparam SL SIMD level (AVX2, ARM_NEON, etc.) +/// @param m_offset stride between beam entries in codes_i +/// @param codebook_cross_norms cross-norm table +/// @param codebook_offsets per-codebook offset +/// @param codes_i code indices +/// @param b beam index +/// @param ldc leading dimension of cross-norm table +/// @param K number of centroids +/// @param output output buffer, size K (accumulated into) +template +void accum_and_add_tab( + size_t m_offset, + const float* __restrict codebook_cross_norms, + const uint64_t* __restrict codebook_offsets, + const int32_t* __restrict codes_i, + size_t b, + size_t ldc, + size_t K, + float* __restrict output); + +/// Accumulate cross-norms for M codebooks and finalize distances. +/// +/// Accumulates M codebook rows, then computes the final centroid distance: +/// output[b*K + k] = distances_i[b] + cd_common[k] + 2 * sum[k] +/// Used for m=1..7 where the entire accumulation fits in registers. +/// +/// @tparam M number of codebook rows to accumulate (equals m) +/// @tparam NK number of 8-float SIMD chunks per loop iteration +/// @tparam SL SIMD level (AVX2, ARM_NEON, etc.) +/// @param codebook_cross_norms cross-norm table +/// @param codebook_offsets per-codebook offset +/// @param codes_i code indices (stride is M) +/// @param b beam index +/// @param ldc leading dimension of cross-norm table +/// @param K number of centroids +/// @param distances_i per-beam input distances, size beam_size +/// @param cd_common common distance term, size K +/// @param output output centroid distances (b*K offset) +template +void accum_and_finalize_tab( + const float* __restrict codebook_cross_norms, + const uint64_t* __restrict codebook_offsets, + const int32_t* __restrict codes_i, + size_t b, + size_t ldc, + size_t K, + const float* __restrict distances_i, + const float* __restrict cd_common, + float* __restrict output); + +} // namespace faiss diff --git a/faiss/impl/approx_topk/simdlib256-inl.h b/faiss/impl/approx_topk/simdlib256-inl.h new file mode 100644 index 000000000..8b93b4cb2 --- /dev/null +++ b/faiss/impl/approx_topk/simdlib256-inl.h @@ -0,0 +1,146 @@ +/* + * 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. + */ + +// Out-of-line definition of HeapWithBucketsCMaxFloat::bs_addn using +// simdlib types. Only included by per-ISA .cpp files (avx2.cpp, neon.cpp). +// Do NOT include this from common translation units. + +#pragma once + +#include +#include + +#include +#include +#include +#include + +namespace faiss { + +// Element-wise max of two simd8float32 vectors, implemented via +// cmplt_min_max_fast (which computes both min and max). +template +inline simd8float32_tpl simd8float32_max( + simd8float32_tpl a, + simd8float32_tpl b) { + simd8float32_tpl min_val, max_val; + simd8uint32_tpl dummy(0u), dmin, dmax; + cmplt_min_max_fast(a, dummy, b, dummy, min_val, dmin, max_val, dmax); + return max_val; +} + +template +void HeapWithBucketsCMaxFloat::bs_addn( + const uint32_t beam_size, + const uint32_t n_per_beam, + const float* const __restrict distances, + const uint32_t k, + float* const __restrict bh_val, + int32_t* const __restrict bh_ids) { + using C = CMax; + using simd_float = simd8float32_tpl; + using simd_uint = simd8uint32_tpl; + + for (uint32_t beam_index = 0; beam_index < beam_size; beam_index++) { + simd_float min_distances_i[NBUCKETS / 8][N]; + simd_uint 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] = + simd_float(std::numeric_limits::max()); + min_indices_i[j][p] = simd_uint(0u, 1u, 2u, 3u, 4u, 5u, 6u, 7u); + } + } + + simd_uint current_indices(0u, 1u, 2u, 3u, 4u, 5u, 6u, 7u); + simd_uint indices_delta(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 simd_float distances_reg( + distances + j * 8 + ip + n_per_beam * beam_index); + + simd_float distance_candidate = distances_reg; + simd_uint indices_candidate = current_indices; + + for (uint32_t p = 0; p < N; p++) { + // Use cmplt_min_max_fast for comparison, min values, + // min indices, and max indices. + simd_float min_d_new, max_d_unused; + simd_uint min_idx_new, max_idx_new; + cmplt_min_max_fast( + distance_candidate, + indices_candidate, + min_distances_i[j][p], + min_indices_i[j][p], + min_d_new, + min_idx_new, + max_d_unused, + max_idx_new); + + // The max distance uses distances_reg (the original + // input), NOT distance_candidate. This is a deliberate + // approximation that breaks the data dependency chain. + simd_float max_d_new = simd8float32_max( + min_distances_i[j][p], distances_reg); + + distance_candidate = max_d_new; + indices_candidate = max_idx_new; + + min_distances_i[j][p] = min_d_new; + min_indices_i[j][p] = min_idx_new; + } + } + + current_indices = current_indices + indices_delta; + } + + // fix the indices + for (uint32_t j = 0; j < NBUCKETS / 8; j++) { + const simd_uint offset(n_per_beam * beam_index + j * 8); + for (uint32_t p = 0; p < N; p++) { + min_indices_i[j][p] = 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++) { + uint32_t min_indices_scalar[8]; + float min_distances_scalar[8]; + + min_indices_i[j][p].storeu(min_indices_scalar); + min_distances_i[j][p].storeu(min_distances_scalar); + + for (size_t j8 = 0; j8 < 8; j8++) { + const auto value = min_distances_scalar[j8]; + const auto index = + static_cast(min_indices_scalar[j8]); + if (C::cmp2(bh_val[0], value, bh_ids[0], index)) { + heap_replace_top(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(k, bh_val, bh_ids, value, index); + } + } + } +} + +} // namespace faiss diff --git a/faiss/impl/residual_quantizer_encode_steps.cpp b/faiss/impl/residual_quantizer_encode_steps.cpp index 3fcac277c..0fa19a2ae 100644 --- a/faiss/impl/residual_quantizer_encode_steps.cpp +++ b/faiss/impl/residual_quantizer_encode_steps.cpp @@ -11,12 +11,12 @@ #include #include #include -#include #include #include #include -#include +#include +#include extern "C" { @@ -39,190 +39,6 @@ int sgemm_( namespace faiss { -/******************************************************************** - * Basic routines - ********************************************************************/ - -namespace { - -template -void accum_and_store_tab( - const size_t m_offset, - const float* const __restrict codebook_cross_norms, - const uint64_t* const __restrict codebook_offsets, - const int32_t* const __restrict codes_i, - const size_t b, - const size_t ldc, - const size_t K, - float* const __restrict output) { - // load pointers into registers - const float* cbs[M]; - for (size_t ij = 0; ij < M; ij++) { - const size_t code = static_cast(codes_i[b * m_offset + ij]); - cbs[ij] = &codebook_cross_norms[(codebook_offsets[ij] + code) * ldc]; - } - - // do accumulation in registers using SIMD. - // It is possible that compiler may be smart enough so that - // this manual SIMD unrolling might be unneeded. -#if defined(__AVX2__) || defined(__aarch64__) - const size_t K8 = (K / (8 * NK)) * (8 * NK); - - // process in chunks of size (8 * NK) floats - for (size_t kk = 0; kk < K8; kk += 8 * NK) { - simd8float32 regs[NK]; - for (size_t ik = 0; ik < NK; ik++) { - regs[ik].loadu(cbs[0] + kk + ik * 8); - } - - for (size_t ij = 1; ij < M; ij++) { - for (size_t ik = 0; ik < NK; ik++) { - regs[ik] += simd8float32(cbs[ij] + kk + ik * 8); - } - } - - // write the result - for (size_t ik = 0; ik < NK; ik++) { - regs[ik].storeu(output + kk + ik * 8); - } - } -#else - const size_t K8 = 0; -#endif - - // process leftovers - for (size_t kk = K8; kk < K; kk++) { - float reg = cbs[0][kk]; - for (size_t ij = 1; ij < M; ij++) { - reg += cbs[ij][kk]; - } - output[kk] = reg; - } -} - -template -void accum_and_add_tab( - const size_t m_offset, - const float* const __restrict codebook_cross_norms, - const uint64_t* const __restrict codebook_offsets, - const int32_t* const __restrict codes_i, - const size_t b, - const size_t ldc, - const size_t K, - float* const __restrict output) { - // load pointers into registers - const float* cbs[M]; - for (size_t ij = 0; ij < M; ij++) { - const size_t code = static_cast(codes_i[b * m_offset + ij]); - cbs[ij] = &codebook_cross_norms[(codebook_offsets[ij] + code) * ldc]; - } - - // do accumulation in registers using SIMD. - // It is possible that compiler may be smart enough so that - // this manual SIMD unrolling might be unneeded. -#if defined(__AVX2__) || defined(__aarch64__) - const size_t K8 = (K / (8 * NK)) * (8 * NK); - - // process in chunks of size (8 * NK) floats - for (size_t kk = 0; kk < K8; kk += 8 * NK) { - simd8float32 regs[NK]; - for (size_t ik = 0; ik < NK; ik++) { - regs[ik].loadu(cbs[0] + kk + ik * 8); - } - - for (size_t ij = 1; ij < M; ij++) { - for (size_t ik = 0; ik < NK; ik++) { - regs[ik] += simd8float32(cbs[ij] + kk + ik * 8); - } - } - - // write the result - for (size_t ik = 0; ik < NK; ik++) { - simd8float32 existing(output + kk + ik * 8); - existing += regs[ik]; - existing.storeu(output + kk + ik * 8); - } - } -#else - const size_t K8 = 0; -#endif - - // process leftovers - for (size_t kk = K8; kk < K; kk++) { - float reg = cbs[0][kk]; - for (size_t ij = 1; ij < M; ij++) { - reg += cbs[ij][kk]; - } - output[kk] += reg; - } -} - -template -void accum_and_finalize_tab( - const float* const __restrict codebook_cross_norms, - const uint64_t* const __restrict codebook_offsets, - const int32_t* const __restrict codes_i, - const size_t b, - const size_t ldc, - const size_t K, - const float* const __restrict distances_i, - const float* const __restrict cd_common, - float* const __restrict output) { - // load pointers into registers - const float* cbs[M]; - for (size_t ij = 0; ij < M; ij++) { - const size_t code = static_cast(codes_i[b * M + ij]); - cbs[ij] = &codebook_cross_norms[(codebook_offsets[ij] + code) * ldc]; - } - - // do accumulation in registers using SIMD. - // It is possible that compiler may be smart enough so that - // this manual SIMD unrolling might be unneeded. -#if defined(__AVX2__) || defined(__aarch64__) - const size_t K8 = (K / (8 * NK)) * (8 * NK); - - // process in chunks of size (8 * NK) floats - for (size_t kk = 0; kk < K8; kk += 8 * NK) { - simd8float32 regs[NK]; - for (size_t ik = 0; ik < NK; ik++) { - regs[ik].loadu(cbs[0] + kk + ik * 8); - } - - for (size_t ij = 1; ij < M; ij++) { - for (size_t ik = 0; ik < NK; ik++) { - regs[ik] += simd8float32(cbs[ij] + kk + ik * 8); - } - } - - simd8float32 two(2.0f); - for (size_t ik = 0; ik < NK; ik++) { - // cent_distances[b * K + k] = distances_i[b] + cd_common[k] - // + 2 * dp[k]; - - simd8float32 common_v(cd_common + kk + ik * 8); - common_v = fmadd(two, regs[ik], common_v); - - common_v += simd8float32(distances_i[b]); - common_v.storeu(output + b * K + kk + ik * 8); - } - } -#else - const size_t K8 = 0; -#endif - - // process leftovers - for (size_t kk = K8; kk < K; kk++) { - float reg = cbs[0][kk]; - for (size_t ij = 1; ij < M; ij++) { - reg += cbs[ij][kk]; - } - - output[b * K + kk] = distances_i[b] + cd_common[kk] + 2 * reg; - } -} - -} // anonymous namespace - /******************************************************************** * Single encoding step ********************************************************************/ @@ -331,22 +147,28 @@ void beam_search_encode_step( } std::vector perm(new_beam_size, -1); -#define HANDLE_APPROX(NB, BD) \ - case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B##NB##_D##BD: \ - HeapWithBuckets::bs_addn( \ - beam_size, \ - K, \ - cent_distances_i, \ - new_beam_size, \ - new_distances_i, \ - perm.data()); \ - break; - + auto approx = [&]() { + HeapWithBuckets::bs_addn( + beam_size, + K, + cent_distances_i, + new_beam_size, + new_distances_i, + perm.data()); + }; switch (approx_topk_mode) { - HANDLE_APPROX(8, 3) - HANDLE_APPROX(8, 2) - HANDLE_APPROX(16, 2) - HANDLE_APPROX(32, 2) + case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B8_D3: + approx.template operator()<8, 3>(); + break; + case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B8_D2: + approx.template operator()<8, 2>(); + break; + case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B16_D2: + approx.template operator()<16, 2>(); + break; + case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B32_D2: + approx.template operator()<32, 2>(); + break; default: heap_addn( new_beam_size, @@ -358,8 +180,6 @@ void beam_search_encode_step( } heap_reorder(new_beam_size, new_distances_i, perm.data()); -#undef HANDLE_APPROX - for (int j = 0; j < new_beam_size; j++) { int js = perm[j] / K; int ls = perm[j] % K; @@ -380,6 +200,166 @@ void beam_search_encode_step( } // exposed in the faiss namespace + +namespace { + +// Baseline (scalar) implementation for computing cent_distances. +// Accumulates codebook cross-norms via fvec_add into a temporary buffer. +// Its primary flaw is that it writes too much to the temporary buffer dp. +// This code is kept because it is easy to understand what the optimized +// SIMD version (compute_cent_distances_simd) optimizes exactly. +void compute_cent_distances_baseline( + size_t K, + size_t beam_size, + const float* codebook_cross_norms, + size_t ldc, + const uint64_t* codebook_offsets, + size_t m, + const int32_t* codes_i, + const float* distances_i, + const float* cd_common, + float* cent_distances) { + for (size_t b = 0; b < beam_size; b++) { + std::vector dp(K); + for (size_t m1 = 0; m1 < m; m1++) { + size_t c = codes_i[b * m + m1]; + const float* cb = + &codebook_cross_norms[(codebook_offsets[m1] + c) * ldc]; + fvec_add(K, cb, dp.data(), dp.data()); + } + for (size_t k = 0; k < K; k++) { + cent_distances[b * K + k] = + distances_i[b] + cd_common[k] + 2 * dp[k]; + } + } +} + +// SIMD-optimized implementation for computing cent_distances. +// Uses accum_and_finalize_tab / accum_and_store_tab / accum_and_add_tab +// to accumulate codebook cross-norms in SIMD registers. +template +void compute_cent_distances_simd( + size_t K, + size_t beam_size, + const float* codebook_cross_norms, + size_t ldc, + const uint64_t* codebook_offsets, + size_t m, + const int32_t* codes_i, + const float* distances_i, + const float* cd_common, + float* cent_distances) { + auto do_finalize = [&]() { + for (size_t b = 0; b < beam_size; b++) { + accum_and_finalize_tab( + codebook_cross_norms, + codebook_offsets, + codes_i, + b, + ldc, + K, + distances_i, + cd_common, + cent_distances); + } + }; + + switch (m) { + case 0: + for (size_t b = 0; b < beam_size; b++) { + for (size_t k = 0; k < K; k++) { + cent_distances[b * K + k] = distances_i[b] + cd_common[k]; + } + } + break; + case 1: + do_finalize.template operator()<1>(); + break; + case 2: + do_finalize.template operator()<2>(); + break; + case 3: + do_finalize.template operator()<3>(); + break; + case 4: + do_finalize.template operator()<4>(); + break; + case 5: + do_finalize.template operator()<5>(); + break; + case 6: + do_finalize.template operator()<6>(); + break; + case 7: + do_finalize.template operator()<7>(); + break; + default: { + // m >= 8: accumulate in chunks of 8 into a temporary buffer. + std::vector dp(K); + + for (size_t b = 0; b < beam_size; b++) { + accum_and_store_tab<8, 4, SL>( + m, + codebook_cross_norms, + codebook_offsets, + codes_i, + b, + ldc, + K, + dp.data()); + + for (size_t im = 8; im < ((m + 7) / 8) * 8; im += 8) { + size_t m_left = std::min(m - im, size_t(8)); + auto do_add = [&]() { + accum_and_add_tab( + m, + codebook_cross_norms, + codebook_offsets + im, + codes_i + im, + b, + ldc, + K, + dp.data()); + }; + switch (m_left) { + case 1: + do_add.template operator()<1>(); + break; + case 2: + do_add.template operator()<2>(); + break; + case 3: + do_add.template operator()<3>(); + break; + case 4: + do_add.template operator()<4>(); + break; + case 5: + do_add.template operator()<5>(); + break; + case 6: + do_add.template operator()<6>(); + break; + case 7: + do_add.template operator()<7>(); + break; + case 8: + do_add.template operator()<8>(); + break; + } + } + + for (size_t k = 0; k < K; k++) { + cent_distances[b * K + k] = + distances_i[b] + cd_common[k] + 2 * dp[k]; + } + } + } + } +} + +} // anonymous namespace + void beam_search_encode_step_tab( size_t K, size_t n, @@ -413,145 +393,33 @@ void beam_search_encode_step_tab( cd_common[k] = cent_norms_i[k] - 2 * query_cp_i[k]; } - bool use_baseline_implementation = false; - - // This is the baseline implementation. Its primary flaw - // that it writes way too many info to the temporary buffer - // called dp. - // - // This baseline code is kept intentionally because it is easy to - // understand what an optimized version optimizes exactly. - // - if (use_baseline_implementation) { - for (size_t b = 0; b < beam_size; b++) { - std::vector dp(K); - - for (size_t m1 = 0; m1 < m; m1++) { - size_t c = codes_i[b * m + m1]; - const float* cb = - &codebook_cross_norms - [(codebook_offsets[m1] + c) * ldc]; - fvec_add(K, cb, dp.data(), dp.data()); - } - - for (size_t k = 0; k < K; k++) { - cent_distances[b * K + k] = - distances_i[b] + cd_common[k] + 2 * dp[k]; - } + with_simd_level_256bit([&]() { + if constexpr (SL == SIMDLevel::NONE) { + compute_cent_distances_baseline( + K, + beam_size, + codebook_cross_norms, + ldc, + codebook_offsets, + m, + codes_i, + distances_i, + cd_common.data(), + cent_distances.data()); + } else { + compute_cent_distances_simd( + K, + beam_size, + codebook_cross_norms, + ldc, + codebook_offsets, + m, + codes_i, + distances_i, + cd_common.data(), + cent_distances.data()); } - - } else { - // An optimized implementation that avoids using a temporary buffer - // and does the accumulation in registers. - - // Compute a sum of NK AQ codes. -#define ACCUM_AND_FINALIZE_TAB(NK) \ - case NK: \ - for (size_t b = 0; b < beam_size; b++) { \ - accum_and_finalize_tab( \ - codebook_cross_norms, \ - codebook_offsets, \ - codes_i, \ - b, \ - ldc, \ - K, \ - distances_i, \ - cd_common.data(), \ - cent_distances.data()); \ - } \ - break; - - // this version contains many switch-case scenarios, but - // they won't affect branch predictor. - switch (m) { - case 0: - // trivial case - for (size_t b = 0; b < beam_size; b++) { - for (size_t k = 0; k < K; k++) { - cent_distances[b * K + k] = - distances_i[b] + cd_common[k]; - } - } - break; - - ACCUM_AND_FINALIZE_TAB(1) - ACCUM_AND_FINALIZE_TAB(2) - ACCUM_AND_FINALIZE_TAB(3) - ACCUM_AND_FINALIZE_TAB(4) - ACCUM_AND_FINALIZE_TAB(5) - ACCUM_AND_FINALIZE_TAB(6) - ACCUM_AND_FINALIZE_TAB(7) - - default: { - // m >= 8 case. - - // A temporary buffer has to be used due to the lack of - // registers. But we'll try to accumulate up to 8 AQ codes - // in registers and issue a single write operation to the - // buffer, while the baseline does no accumulation. So, the - // number of write operations to the temporary buffer is - // reduced 8x. - - // allocate a temporary buffer - std::vector dp(K); - - for (size_t b = 0; b < beam_size; b++) { - // Initialize it. Compute a sum of first 8 AQ codes - // because m >= 8 . - accum_and_store_tab<8, 4>( - m, - codebook_cross_norms, - codebook_offsets, - codes_i, - b, - ldc, - K, - dp.data()); - -#define ACCUM_AND_ADD_TAB(NK) \ - case NK: \ - accum_and_add_tab( \ - m, \ - codebook_cross_norms, \ - codebook_offsets + im, \ - codes_i + im, \ - b, \ - ldc, \ - K, \ - dp.data()); \ - break; - - // accumulate up to 8 additional AQ codes into - // a temporary buffer - for (size_t im = 8; im < ((m + 7) / 8) * 8; im += 8) { - size_t m_left = m - im; - if (m_left > 8) { - m_left = 8; - } - - switch (m_left) { - ACCUM_AND_ADD_TAB(1) - ACCUM_AND_ADD_TAB(2) - ACCUM_AND_ADD_TAB(3) - ACCUM_AND_ADD_TAB(4) - ACCUM_AND_ADD_TAB(5) - ACCUM_AND_ADD_TAB(6) - ACCUM_AND_ADD_TAB(7) - ACCUM_AND_ADD_TAB(8) - } - } - - // done. finalize the result - for (size_t k = 0; k < K; k++) { - cent_distances[b * K + k] = - distances_i[b] + cd_common[k] + 2 * dp[k]; - } - } - } - } - - // the optimized implementation ends here - } + }); using C = CMax; int32_t* new_codes_i = new_codes + i * (m + 1) * new_beam_size; float* new_distances_i = new_distances + i * new_beam_size; @@ -564,22 +432,28 @@ void beam_search_encode_step_tab( } std::vector perm(new_beam_size, -1); -#define HANDLE_APPROX(NB, BD) \ - case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B##NB##_D##BD: \ - HeapWithBuckets::bs_addn( \ - beam_size, \ - K, \ - cent_distances_i, \ - new_beam_size, \ - new_distances_i, \ - perm.data()); \ - break; - + auto approx = [&]() { + HeapWithBuckets::bs_addn( + beam_size, + K, + cent_distances_i, + new_beam_size, + new_distances_i, + perm.data()); + }; switch (approx_topk_mode) { - HANDLE_APPROX(8, 3) - HANDLE_APPROX(8, 2) - HANDLE_APPROX(16, 2) - HANDLE_APPROX(32, 2) + case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B8_D3: + approx.template operator()<8, 3>(); + break; + case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B8_D2: + approx.template operator()<8, 2>(); + break; + case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B16_D2: + approx.template operator()<16, 2>(); + break; + case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B32_D2: + approx.template operator()<32, 2>(); + break; default: heap_addn( new_beam_size, @@ -593,8 +467,6 @@ void beam_search_encode_step_tab( heap_reorder(new_beam_size, new_distances_i, perm.data()); -#undef HANDLE_APPROX - for (int j = 0; j < new_beam_size; j++) { int js = perm[j] / K; int ls = perm[j] % K; diff --git a/faiss/impl/residual_quantizer_encode_steps.h b/faiss/impl/residual_quantizer_encode_steps.h index 3d17efd94..1d98aae5f 100644 --- a/faiss/impl/residual_quantizer_encode_steps.h +++ b/faiss/impl/residual_quantizer_encode_steps.h @@ -11,7 +11,7 @@ #include #include -#include +#include namespace faiss { diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index 30099a296..033dc8d07 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -691,7 +691,7 @@ void gpu_sync_all_devices() %template(IndexBinaryIDMap2) faiss::IndexIDMap2Template; -%include +%include #ifdef FAISS_ENABLE_SVS %ignore faiss::to_svs_storage_kind; diff --git a/faiss/utils/approx_topk/approx_topk.h b/faiss/utils/approx_topk/approx_topk.h deleted file mode 100644 index b7bfb7ec4..000000000 --- a/faiss/utils/approx_topk/approx_topk.h +++ /dev/null @@ -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 - -// the list of available modes is in the following file -#include - -#ifdef __AVX2__ -#include -#else -#include -#endif diff --git a/faiss/utils/approx_topk/avx2-inl.h b/faiss/utils/approx_topk/avx2-inl.h deleted file mode 100644 index 3378edf6a..000000000 --- a/faiss/utils/approx_topk/avx2-inl.h +++ /dev/null @@ -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 - -#include - -#include -#include - -namespace faiss { - -template -struct HeapWithBuckets { - // this case was not implemented yet. -}; - -template -struct HeapWithBuckets, 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; - - // 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::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( - 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(k, bh_val, bh_ids, value, index); - } - } - } - } -}; - -} // namespace faiss diff --git a/faiss/utils/approx_topk/mode.h b/faiss/utils/approx_topk/mode.h deleted file mode 100644 index ad9d3ebda..000000000 --- a/faiss/utils/approx_topk/mode.h +++ /dev/null @@ -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, -}; diff --git a/faiss/utils/hamming.h b/faiss/utils/hamming.h index e1a442229..64dfbf746 100644 --- a/faiss/utils/hamming.h +++ b/faiss/utils/hamming.h @@ -34,7 +34,7 @@ // Low-level Hamming distance computations and hamdis_t. #include -#include +#include namespace faiss { diff --git a/tests/test_approx_topk.cpp b/tests/test_approx_topk.cpp index ff4ee1f2e..ebfe01f8a 100644 --- a/tests/test_approx_topk.cpp +++ b/tests/test_approx_topk.cpp @@ -15,7 +15,7 @@ #include #include -#include +#include #include #include