facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-11 22:50:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Convert approx_topk and residual_quantizer_encode_steps to SIMDLevel dynamic dispatch (#4923)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4923

Convert HeapWithBuckets and the accum_and_*_tab functions from compile-time #ifdef __AVX2__ / __aarch64__ guards to SIMDLevel-templatized dynamic dispatch, consistent with the rest of FAISS's DD infrastructure (with_simd_level, simdlib, per-ISA .cpp files).

Migration plan: P2234192040

Key changes:
- New approx_topk.h: merges mode.h + heap_with_buckets.h. HeapWithBuckets<C,NB,N> dispatches via with_simd_level_256bit; SIMDLevel::NONE uses generic.h scalar fallback. SIMD primary template declares bs_addn without defining it.
- New simdlib256-inl.h: out-of-line definition of HeapWithBucketsCMaxFloat::bs_addn using simdlib types (simd8float32_tpl/simd8uint32_tpl). Only included by per-ISA .cpp files.
- New rq_beam_search_tab.h / rq_beam_search_tab-inl.h: declarations and definitions of accum_and_store/add/finalize_tab templatized on SIMDLevel, extracted from residual_quantizer_encode_steps.cpp.
- New avx2.cpp / neon.cpp: explicit template instantiations per ISA.
- residual_quantizer_encode_steps.cpp: HANDLE_APPROX / ACCUM_AND_*_TAB macros replaced with templatized lambdas + with_simd_level_256bit dispatch. Accumulation logic decomposed into compute_cent_distances_baseline and compute_cent_distances_simd<SL>.
- Deleted: avx2-inl.h, mode.h, heap_with_buckets.h (merged into approx_topk.h).
- Moved: approx_topk directory from utils/ to impl/.

Reviewed By: algoriddle

Differential Revision: D96360821

fbshipit-source-id: a1facb34b8af065e0a1430897cc05fbe763d3e26
This commit is contained in:
Matthijs Douze
2026-03-14 15:12:14 -07:00
committed by meta-codesync[bot]
parent 58a57e6bca
commit b1feeb79f8
19 changed files with 1042 additions and 691 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+7 -4
View File
@@ -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
faiss/IndexBinaryFlat.h
+1 -1
View File
@@ -14,8 +14,8 @@
#include <faiss/IndexBinary.h>
#include <faiss/impl/approx_topk/approx_topk.h>
#include <faiss/impl/maybe_owned_vector.h>
#include <faiss/utils/approx_topk/mode.h>
namespace faiss {
faiss/impl/LocalSearchQuantizer.cpp
+1 -1
View File
@@ -22,7 +22,7 @@
#include <faiss/utils/distances.h>
#include <faiss/utils/utils.h>
#include <faiss/utils/approx_topk/approx_topk.h>
#include <faiss/impl/approx_topk/approx_topk.h>
// this is needed for prefetching
faiss/impl/ResidualQuantizer.h
+1 -1
View File
@@ -13,7 +13,7 @@
#include <faiss/Clustering.h>
#include <faiss/impl/AdditiveQuantizer.h>
#include <faiss/utils/approx_topk/mode.h>
#include <faiss/impl/approx_topk/approx_topk.h>
namespace faiss {
faiss/impl/approx_topk/approx_topk.h
+215
View File
@@ -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<C, NBUCKETS, N> 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 <cstdint>
#include <faiss/impl/approx_topk/generic.h>
#include <faiss/impl/platform_macros.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/utils/Heap.h>
#include <faiss/utils/simd_levels.h>
// -----------------------------------------------------------------------
// 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 <uint32_t NBUCKETS, uint32_t N, SIMDLevel SL>
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 <uint32_t NBUCKETS, uint32_t N>
struct HeapWithBucketsCMaxFloat<NBUCKETS, N, SIMDLevel::NONE> {
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<NBUCKETS, N>::bs_addn(
beam_size, n_per_beam, distances, k, bh_val, bh_ids);
}
};
// Primary template — not implemented for arbitrary comparators.
template <typename C, uint32_t NBUCKETS, uint32_t N>
struct HeapWithBuckets {
static_assert(
sizeof(C) == 0,
"HeapWithBuckets: unsupported comparator type");
};
// Partial specialization for CMax<float, int> that dispatches
// via with_simd_level_256bit.
template <uint32_t NBUCKETS, uint32_t N>
struct HeapWithBuckets<CMax<float, int>, 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([&]<SIMDLevel SL>() {
HeapWithBucketsCMaxFloat<NBUCKETS, N, SL>::bs_addn(
beam_size, n_per_beam, distances, k, bh_val, bh_ids);
});
}
};
} // namespace faiss
faiss/impl/approx_topk/avx2.cpp
+68
View File
@@ -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 <faiss/impl/approx_topk/rq_beam_search_tab-inl.h>
#include <faiss/impl/approx_topk/simdlib256-inl.h>
#include <faiss/impl/simdlib/simdlib_avx2.h>
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<M, 4, SIMDLevel::AVX2>( \
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<M, 4, SIMDLevel::AVX2>( \
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<M, 4, SIMDLevel::AVX2>( \
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
faiss/utils/approx_topk/generic.h → faiss/impl/approx_topk/generic.h
+15 -8
View File
@@ -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 <typename C, uint32_t NBUCKETS, uint32_t N>
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 <uint32_t NBUCKETS, uint32_t N>
struct HeapWithBuckets<CMax<float, int>, NBUCKETS, N> {
struct HeapWithBucketsGenericCMaxFloat {
static void addn(
// number of elements
const uint32_t n,
@@ -135,4 +131,15 @@ struct HeapWithBuckets<CMax<float, int>, NBUCKETS, N> {
}
};
// Legacy name kept for backward compatibility (used when
// approx_topk.h is not included).
template <typename C, uint32_t NBUCKETS, uint32_t N>
struct HeapWithBucketsGeneric {
// not implemented for arbitrary C
};
template <uint32_t NBUCKETS, uint32_t N>
struct HeapWithBucketsGeneric<CMax<float, int>, NBUCKETS, N>
: HeapWithBucketsGenericCMaxFloat<NBUCKETS, N> {};
} // namespace faiss
faiss/impl/approx_topk/neon.cpp
+68
View File
@@ -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 <faiss/impl/approx_topk/rq_beam_search_tab-inl.h>
#include <faiss/impl/approx_topk/simdlib256-inl.h>
#include <faiss/impl/simdlib/simdlib_neon.h>
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<M, 4, SIMDLevel::ARM_NEON>( \
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<M, 4, SIMDLevel::ARM_NEON>( \
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<M, 4, SIMDLevel::ARM_NEON>( \
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
faiss/impl/approx_topk/rq_beam_search_tab-inl.h
+169
View File
@@ -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 <cstddef>
#include <cstdint>
#include <faiss/impl/approx_topk/rq_beam_search_tab.h>
#include <faiss/impl/simdlib/simdlib.h>
namespace faiss {
template <size_t M, size_t NK, SIMDLevel SL>
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<SL>;
const float* cbs[M];
for (size_t ij = 0; ij < M; ij++) {
const size_t code = static_cast<size_t>(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 <size_t M, size_t NK, SIMDLevel SL>
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<SL>;
const float* cbs[M];
for (size_t ij = 0; ij < M; ij++) {
const size_t code = static_cast<size_t>(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 <size_t M, size_t NK, SIMDLevel SL>
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<SL>;
const float* cbs[M];
for (size_t ij = 0; ij < M; ij++) {
const size_t code = static_cast<size_t>(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
faiss/impl/approx_topk/rq_beam_search_tab.h
+117
View File
@@ -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 <cstddef>
#include <cstdint>
#include <faiss/utils/simd_levels.h>
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 <size_t M, size_t NK, SIMDLevel SL>
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 <size_t M, size_t NK, SIMDLevel SL>
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 <size_t M, size_t NK, SIMDLevel SL>
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
faiss/impl/approx_topk/simdlib256-inl.h
+146
View File
@@ -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 <cstdint>
#include <limits>
#include <faiss/impl/approx_topk/approx_topk.h>
#include <faiss/impl/simdlib/simdlib.h>
#include <faiss/utils/Heap.h>
#include <faiss/utils/simd_levels.h>
namespace faiss {
// Element-wise max of two simd8float32 vectors, implemented via
// cmplt_min_max_fast (which computes both min and max).
template <SIMDLevel SL>
inline simd8float32_tpl<SL> simd8float32_max(
simd8float32_tpl<SL> a,
simd8float32_tpl<SL> b) {
simd8float32_tpl<SL> min_val, max_val;
simd8uint32_tpl<SL> dummy(0u), dmin, dmax;
cmplt_min_max_fast(a, dummy, b, dummy, min_val, dmin, max_val, dmax);
return max_val;
}
template <uint32_t NBUCKETS, uint32_t N, SIMDLevel SL>
void HeapWithBucketsCMaxFloat<NBUCKETS, N, SL>::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<float, int>;
using simd_float = simd8float32_tpl<SL>;
using simd_uint = simd8uint32_tpl<SL>;
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<float>::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<SL>(
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<int32_t>(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
faiss/impl/residual_quantizer_encode_steps.cpp
+230 -358
View File
@@ -11,12 +11,12 @@
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/ResidualQuantizer.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/impl/simdlib/simdlib_dispatch.h>
#include <faiss/utils/Heap.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/utils.h>
#include <faiss/utils/approx_topk/approx_topk.h>
#include <faiss/impl/approx_topk/approx_topk.h>
#include <faiss/impl/approx_topk/rq_beam_search_tab.h>
extern "C" {
@@ -39,190 +39,6 @@ int sgemm_(
namespace faiss {
/********************************************************************
* Basic routines
********************************************************************/
namespace {
template <size_t M, size_t NK>
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<size_t>(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 <size_t M, size_t NK>
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<size_t>(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 <size_t M, size_t NK>
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<size_t>(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<int> perm(new_beam_size, -1);
#define HANDLE_APPROX(NB, BD) \
case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B##NB##_D##BD: \
HeapWithBuckets<C, NB, BD>::bs_addn( \
beam_size, \
K, \
cent_distances_i, \
new_beam_size, \
new_distances_i, \
perm.data()); \
break;
auto approx = [&]<uint32_t NB, uint32_t ND>() {
HeapWithBuckets<C, NB, ND>::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<C>(
new_beam_size,
@@ -358,8 +180,6 @@ void beam_search_encode_step(
}
heap_reorder<C>(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<float> 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 <SIMDLevel SL>
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 = [&]<size_t NK>() {
for (size_t b = 0; b < beam_size; b++) {
accum_and_finalize_tab<NK, 4, SL>(
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<float> 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 = [&]<size_t NK2>() {
accum_and_add_tab<NK2, 4, SL>(
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<float> 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([&]<SIMDLevel SL>() {
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<SL>(
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<NK, 4>( \
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<float> 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<NK, 4>( \
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<float, int>;
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<int> perm(new_beam_size, -1);
#define HANDLE_APPROX(NB, BD) \
case ApproxTopK_mode_t::APPROX_TOPK_BUCKETS_B##NB##_D##BD: \
HeapWithBuckets<C, NB, BD>::bs_addn( \
beam_size, \
K, \
cent_distances_i, \
new_beam_size, \
new_distances_i, \
perm.data()); \
break;
auto approx = [&]<uint32_t NB, uint32_t ND>() {
HeapWithBuckets<C, NB, ND>::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<C>(
new_beam_size,
@@ -593,8 +467,6 @@ void beam_search_encode_step_tab(
heap_reorder<C>(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;
faiss/impl/residual_quantizer_encode_steps.h
+1 -1
View File
@@ -11,7 +11,7 @@
#include <vector>
#include <faiss/Index.h>
#include <faiss/utils/approx_topk/mode.h>
#include <faiss/impl/approx_topk/approx_topk.h>
namespace faiss {
faiss/python/swigfaiss.swig
+1 -1
View File
@@ -691,7 +691,7 @@ void gpu_sync_all_devices()
%template(IndexBinaryIDMap2) faiss::IndexIDMap2Template<faiss::IndexBinary>;
%include <faiss/utils/approx_topk/mode.h>
%include <faiss/impl/approx_topk/approx_topk.h>
#ifdef FAISS_ENABLE_SVS
%ignore faiss::to_svs_storage_kind;
faiss/utils/approx_topk/approx_topk.h
-84
View File
@@ -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
faiss/utils/approx_topk/avx2-inl.h
-196
View File
@@ -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
faiss/utils/approx_topk/mode.h
-34
View File
@@ -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,
};
faiss/utils/hamming.h
+1 -1
View File
@@ -34,7 +34,7 @@
// Low-level Hamming distance computations and hamdis_t.
#include <faiss/utils/hamming_distance/hamdis-inl.h>
#include <faiss/utils/approx_topk/mode.h>
#include <faiss/impl/approx_topk/approx_topk.h>
namespace faiss {
tests/test_approx_topk.cpp
+1 -1
View File
@@ -15,7 +15,7 @@
#include <unordered_set>
#include <vector>
#include <faiss/utils/approx_topk/approx_topk.h>
#include <faiss/impl/approx_topk/approx_topk.h>
#include <faiss/impl/FaissException.h>
#include <faiss/utils/Heap.h>