mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
Add FastScanCodeScanner dispatch boundary with per-SIMD TUs (#4897)
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4897 Add `FastScanCodeScanner`, a virtual base that bundles handler + kernel behind the SIMD dispatch boundary. In DD mode, `SINGLE_SIMD_LEVEL = NONE` so the existing fast scan code path uses emulated SIMD types. The new scanner provides per-SIMD translation units (AVX2, AVX512, ARM_NEON) compiled with the correct ISA flags, and a factory function (`make_fast_scan_knn_scanner`) that uses `DISPATCH_SIMDLevel` to select the right TU at runtime. This follows the proven `THE_LEVEL_TO_DISPATCH` pattern from the scalar quantizer per-SIMD TUs (`sq-dispatch.h`). Each per-SIMD TU includes `dispatching.h` which provides: - `ScannerMixIn<Handler>`: wraps a concrete handler and calls accumulation kernels (both search_1 multi-BB and QBS paths) - Factory specialization `make_fast_scan_scanner_impl<SL>()` with combinatorial dispatch over `is_max × with_id_map × handler_type` (SingleResultHandler for k=1, HeapHandler for k≤20, ReservoirHandler for k>20) New files: - `impl/fast_scan/dispatching.h` — dispatch template header - `impl/fast_scan/impl-avx2.cpp` — AVX2 per-SIMD TU - `impl/fast_scan/impl-avx512.cpp` — AVX512 per-SIMD TU - `impl/fast_scan/impl-neon.cpp` — ARM NEON TU (with ARM_SVE forwarding) Modified files: - `impl/fast_scan/pq4_fast_scan.h` — FastScanCodeScanner base + factory decl - `impl/fast_scan/pq4_fast_scan.cpp` — NONE specialization + dispatch wrapper - `xplat.bzl` / `CMakeLists.txt` — register SIMD files and header Note: RaBitQ handler is not wired through FastScanCodeScanner in this diff. That comes in later diffs when callers are switched. Reviewed By: mdouze Differential Revision: D95950483 fbshipit-source-id: 03841671f59c0d99de7d0c7408a64e7a8383a8e7
This commit is contained in:
Gergely Szilvasy
committed by
meta-codesync[bot]
meta-codesync[bot]
parent
5b83ec6856
commit
4a0fba1fe2
@@ -9,16 +9,19 @@
|
||||
# Architecture-specific: only include files for the current build architecture
|
||||
# =============================================================================
|
||||
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
|
||||
utils/simd_impl/distances_avx2.cpp
|
||||
)
|
||||
set(FAISS_SIMD_AVX512_SRC
|
||||
impl/fast_scan/impl-avx512.cpp
|
||||
impl/pq_code_distance/pq_code_distance-avx512.cpp
|
||||
impl/scalar_quantizer/sq-avx512.cpp
|
||||
utils/simd_impl/distances_avx512.cpp
|
||||
)
|
||||
set(FAISS_SIMD_NEON_SRC
|
||||
impl/fast_scan/impl-neon.cpp
|
||||
impl/scalar_quantizer/sq-neon.cpp
|
||||
utils/simd_impl/distances_aarch64.cpp
|
||||
)
|
||||
@@ -117,7 +120,7 @@ set(FAISS_SRC
|
||||
impl/kmeans1d.cpp
|
||||
impl/lattice_Zn.cpp
|
||||
impl/mapped_io.cpp
|
||||
impl/fast_scan/pq4_fast_scan.cpp
|
||||
impl/fast_scan/fast_scan.cpp
|
||||
impl/fast_scan/pq4_fast_scan_search_1.cpp
|
||||
impl/fast_scan/pq4_fast_scan_search_qbs.cpp
|
||||
impl/residual_quantizer_encode_steps.cpp
|
||||
@@ -262,7 +265,9 @@ set(FAISS_HEADERS
|
||||
impl/kmeans1d.h
|
||||
impl/lattice_Zn.h
|
||||
impl/platform_macros.h
|
||||
impl/fast_scan/pq4_fast_scan.h
|
||||
impl/fast_scan/accumulate_loops.h
|
||||
impl/fast_scan/dispatching.h
|
||||
impl/fast_scan/fast_scan.h
|
||||
impl/fast_scan/decompose_qbs.h
|
||||
impl/fast_scan/kernels_simd256.h
|
||||
impl/fast_scan/kernels_simd512.h
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
#include <faiss/impl/LocalSearchQuantizer.h>
|
||||
#include <faiss/impl/ResidualQuantizer.h>
|
||||
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/utils/quantize_lut.h>
|
||||
#include <faiss/utils/utils.h>
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@
|
||||
#include <faiss/impl/IDSelector.h>
|
||||
#include <faiss/impl/RaBitQUtils.h>
|
||||
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/simd_result_handlers.h>
|
||||
#include <faiss/utils/hamming.h>
|
||||
#include <faiss/utils/quantize_lut.h>
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
#include <faiss/impl/AuxIndexStructures.h>
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/impl/simd_dispatch.h>
|
||||
#include <faiss/invlists/BlockInvertedLists.h>
|
||||
#include <faiss/utils/distances.h>
|
||||
|
||||
@@ -19,7 +19,7 @@
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/RaBitQUtils.h>
|
||||
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/simd_result_handlers.h>
|
||||
#include <faiss/invlists/BlockInvertedLists.h>
|
||||
#include <faiss/utils/hamming.h>
|
||||
|
||||
@@ -22,7 +22,7 @@
|
||||
|
||||
#include <faiss/invlists/BlockInvertedLists.h>
|
||||
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/simd_result_handlers.h>
|
||||
|
||||
namespace faiss {
|
||||
|
||||
@@ -17,7 +17,7 @@
|
||||
#include <faiss/impl/RaBitQUtils.h>
|
||||
#include <faiss/impl/RaBitQuantizerMultiBit.h>
|
||||
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/simd_result_handlers.h>
|
||||
#include <faiss/invlists/BlockInvertedLists.h>
|
||||
#include <faiss/utils/distances.h>
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
#include <memory>
|
||||
|
||||
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/utils/utils.h>
|
||||
|
||||
namespace faiss {
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
#include <faiss/impl/RaBitQUtils.h>
|
||||
#include <faiss/impl/RaBitQuantizerMultiBit.h>
|
||||
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/utils/utils.h>
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
*/
|
||||
|
||||
#include <faiss/impl/CodePackerRaBitQ.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
|
||||
#include <cstring>
|
||||
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
/*
|
||||
* 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
|
||||
|
||||
/**
|
||||
* @file accumulate_loops.h
|
||||
* @brief Shared accumulation loop helpers for fast-scan search paths.
|
||||
*
|
||||
* Contains:
|
||||
* - accumulate_fixed_blocks / pq4_accumulate_loop_fixed_scaler
|
||||
* (search_1 multi-BB path, bbs > 32)
|
||||
* - accumulate_q_4step_256 / pq4_accumulate_loop_qbs_fixed_scaler_256
|
||||
* (QBS path, bbs == 32, 256-bit kernel only)
|
||||
*
|
||||
* The QBS helpers use pq4_kernel_qbs_256 exclusively (not decompose_qbs.h)
|
||||
* because decompose_qbs.h includes kernels_simd512.h which uses 512-bit
|
||||
* types that are empty primary templates when SINGLE_SIMD_LEVEL=NONE
|
||||
* (DD mode). SL-parameterizing the 512-bit kernels is future work.
|
||||
*
|
||||
* All functions live in `namespace faiss` (not anonymous) so they can be
|
||||
* shared by both the per-SIMD TU dispatcher (dispatching.h) and the old
|
||||
* free-function search paths (pq4_fast_scan_search_1.cpp).
|
||||
*
|
||||
* The QBS helpers here always use pq4_kernel_qbs_256 (never 512-bit).
|
||||
* This is required for the per-SIMD DD TUs where SINGLE_SIMD_LEVEL=NONE
|
||||
* leaves 512-bit types empty. The old pq4_fast_scan_search_qbs.cpp
|
||||
* continues to use decompose_qbs.h which includes both 256 and 512 paths.
|
||||
*/
|
||||
|
||||
#include <cassert>
|
||||
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/fast_scan/LookupTableScaler.h>
|
||||
#include <faiss/impl/fast_scan/kernels_simd256.h>
|
||||
#include <faiss/impl/fast_scan/simd_result_handlers.h>
|
||||
|
||||
namespace faiss {
|
||||
|
||||
using namespace simd_result_handlers;
|
||||
|
||||
/***************************************************************
|
||||
* Search_1 path helpers (multi-BB kernel, bbs > 32)
|
||||
***************************************************************/
|
||||
|
||||
template <int NQ, int BB, class ResultHandler, class Scaler>
|
||||
void accumulate_fixed_blocks(
|
||||
size_t nb,
|
||||
int nsq,
|
||||
const uint8_t* codes,
|
||||
const uint8_t* LUT,
|
||||
ResultHandler& res,
|
||||
const Scaler& scaler,
|
||||
size_t block_stride) {
|
||||
constexpr int bbs = 32 * BB;
|
||||
for (size_t j0 = 0; j0 < nb; j0 += bbs) {
|
||||
FixedStorageHandler<NQ, 2 * BB> res2;
|
||||
kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2, scaler);
|
||||
res.set_block_origin(0, j0);
|
||||
res2.to_other_handler(res);
|
||||
codes += block_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <class ResultHandler, class Scaler>
|
||||
void pq4_accumulate_loop_fixed_scaler(
|
||||
int nq,
|
||||
size_t nb,
|
||||
int bbs,
|
||||
int nsq,
|
||||
const uint8_t* codes,
|
||||
const uint8_t* LUT,
|
||||
ResultHandler& res,
|
||||
const Scaler& scaler,
|
||||
size_t block_stride) {
|
||||
FAISS_THROW_IF_NOT(is_aligned_pointer(codes));
|
||||
FAISS_THROW_IF_NOT(is_aligned_pointer(LUT));
|
||||
FAISS_THROW_IF_NOT(bbs % 32 == 0);
|
||||
FAISS_THROW_IF_NOT(nb % bbs == 0);
|
||||
|
||||
#define FAISS_ACCLOOP_DISPATCH(NQ, BB) \
|
||||
case NQ * 1000 + BB: \
|
||||
accumulate_fixed_blocks<NQ, BB>( \
|
||||
nb, nsq, codes, LUT, res, scaler, block_stride); \
|
||||
break
|
||||
|
||||
switch (nq * 1000 + bbs / 32) {
|
||||
FAISS_ACCLOOP_DISPATCH(1, 1);
|
||||
FAISS_ACCLOOP_DISPATCH(1, 2);
|
||||
FAISS_ACCLOOP_DISPATCH(1, 3);
|
||||
FAISS_ACCLOOP_DISPATCH(1, 4);
|
||||
FAISS_ACCLOOP_DISPATCH(1, 5);
|
||||
FAISS_ACCLOOP_DISPATCH(2, 1);
|
||||
FAISS_ACCLOOP_DISPATCH(2, 2);
|
||||
FAISS_ACCLOOP_DISPATCH(3, 1);
|
||||
FAISS_ACCLOOP_DISPATCH(4, 1);
|
||||
default:
|
||||
FAISS_THROW_FMT("nq=%d bbs=%d not instantiated", nq, bbs);
|
||||
}
|
||||
#undef FAISS_ACCLOOP_DISPATCH
|
||||
}
|
||||
|
||||
/***************************************************************
|
||||
* QBS path helpers (bbs == 32, 256-bit kernel only)
|
||||
***************************************************************/
|
||||
|
||||
template <int QBS, class ResultHandler, class Scaler>
|
||||
void accumulate_q_4step_256(
|
||||
size_t ntotal2,
|
||||
int nsq,
|
||||
const uint8_t* codes,
|
||||
const uint8_t* LUT0,
|
||||
ResultHandler& res,
|
||||
const Scaler& scaler,
|
||||
size_t block_stride) {
|
||||
constexpr int Q1 = QBS & 15;
|
||||
constexpr int Q2 = (QBS >> 4) & 15;
|
||||
constexpr int Q3 = (QBS >> 8) & 15;
|
||||
constexpr int Q4 = (QBS >> 12) & 15;
|
||||
constexpr int SQ = Q1 + Q2 + Q3 + Q4;
|
||||
|
||||
for (size_t j0 = 0; j0 < ntotal2; j0 += 32) {
|
||||
FixedStorageHandler<SQ, 2> res2;
|
||||
const uint8_t* LUT = LUT0;
|
||||
pq4_kernel_qbs_256<Q1>(nsq, codes, LUT, res2, scaler);
|
||||
LUT += Q1 * nsq * 16;
|
||||
if (Q2 > 0) {
|
||||
res2.set_block_origin(Q1, 0);
|
||||
pq4_kernel_qbs_256<Q2>(nsq, codes, LUT, res2, scaler);
|
||||
LUT += Q2 * nsq * 16;
|
||||
}
|
||||
if (Q3 > 0) {
|
||||
res2.set_block_origin(Q1 + Q2, 0);
|
||||
pq4_kernel_qbs_256<Q3>(nsq, codes, LUT, res2, scaler);
|
||||
LUT += Q3 * nsq * 16;
|
||||
}
|
||||
if (Q4 > 0) {
|
||||
res2.set_block_origin(Q1 + Q2 + Q3, 0);
|
||||
pq4_kernel_qbs_256<Q4>(nsq, codes, LUT, res2, scaler);
|
||||
}
|
||||
res.set_block_origin(0, j0);
|
||||
res2.to_other_handler(res);
|
||||
codes += block_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <class ResultHandler, class Scaler>
|
||||
void pq4_accumulate_loop_qbs_fixed_scaler_256(
|
||||
int qbs,
|
||||
size_t ntotal2,
|
||||
int nsq,
|
||||
const uint8_t* codes,
|
||||
const uint8_t* LUT0,
|
||||
ResultHandler& res,
|
||||
const Scaler& scaler,
|
||||
size_t block_stride) {
|
||||
assert(nsq % 2 == 0);
|
||||
assert(is_aligned_pointer(codes));
|
||||
assert(is_aligned_pointer(LUT0));
|
||||
|
||||
switch (qbs) {
|
||||
#define FAISS_QBS256_DISPATCH(QBS) \
|
||||
case QBS: \
|
||||
accumulate_q_4step_256<QBS>( \
|
||||
ntotal2, nsq, codes, LUT0, res, scaler, block_stride); \
|
||||
return;
|
||||
FAISS_QBS256_DISPATCH(0x3333); // 12
|
||||
FAISS_QBS256_DISPATCH(0x2333); // 11
|
||||
FAISS_QBS256_DISPATCH(0x2233); // 10
|
||||
FAISS_QBS256_DISPATCH(0x333); // 9
|
||||
FAISS_QBS256_DISPATCH(0x2223); // 9
|
||||
FAISS_QBS256_DISPATCH(0x233); // 8
|
||||
FAISS_QBS256_DISPATCH(0x1223); // 8
|
||||
FAISS_QBS256_DISPATCH(0x223); // 7
|
||||
FAISS_QBS256_DISPATCH(0x34); // 7
|
||||
FAISS_QBS256_DISPATCH(0x133); // 7
|
||||
FAISS_QBS256_DISPATCH(0x6); // 6
|
||||
FAISS_QBS256_DISPATCH(0x33); // 6
|
||||
FAISS_QBS256_DISPATCH(0x123); // 6
|
||||
FAISS_QBS256_DISPATCH(0x222); // 6
|
||||
FAISS_QBS256_DISPATCH(0x23); // 5
|
||||
FAISS_QBS256_DISPATCH(0x5); // 5
|
||||
FAISS_QBS256_DISPATCH(0x13); // 4
|
||||
FAISS_QBS256_DISPATCH(0x22); // 4
|
||||
FAISS_QBS256_DISPATCH(0x4); // 4
|
||||
FAISS_QBS256_DISPATCH(0x3); // 3
|
||||
FAISS_QBS256_DISPATCH(0x21); // 3
|
||||
FAISS_QBS256_DISPATCH(0x2); // 2
|
||||
FAISS_QBS256_DISPATCH(0x1); // 1
|
||||
#undef FAISS_QBS256_DISPATCH
|
||||
}
|
||||
|
||||
// Default: qbs not known at compile time
|
||||
for (size_t j0 = 0; j0 < ntotal2; j0 += 32) {
|
||||
const uint8_t* LUT = LUT0;
|
||||
int qi = qbs;
|
||||
int i0 = 0;
|
||||
while (qi) {
|
||||
int nq = qi & 15;
|
||||
qi >>= 4;
|
||||
res.set_block_origin(i0, j0);
|
||||
#define FAISS_NQ256_DISPATCH(NQ) \
|
||||
case NQ: \
|
||||
pq4_kernel_qbs_256<NQ>(nsq, codes, LUT, res, scaler); \
|
||||
break
|
||||
switch (nq) {
|
||||
FAISS_NQ256_DISPATCH(1);
|
||||
FAISS_NQ256_DISPATCH(2);
|
||||
FAISS_NQ256_DISPATCH(3);
|
||||
FAISS_NQ256_DISPATCH(4);
|
||||
#undef FAISS_NQ256_DISPATCH
|
||||
default:
|
||||
FAISS_THROW_FMT("accumulate nq=%d not instantiated", nq);
|
||||
}
|
||||
i0 += nq;
|
||||
LUT += nq * nsq * 16;
|
||||
}
|
||||
codes += block_stride;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace faiss
|
||||
@@ -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.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
/**
|
||||
* @file dispatching.h
|
||||
* @brief Per-SIMD TU dispatch template for fast scan.
|
||||
*
|
||||
* This header is included once per SIMD TU with THE_LEVEL_TO_DISPATCH
|
||||
* set to the desired SIMDLevel. It provides:
|
||||
* - ScannerMixIn: wraps a handler + calls kernel at the TU's SIMD level
|
||||
* - make_fast_scan_scanner_impl<SL>: factory specialization
|
||||
*
|
||||
* Usage (in a per-SIMD .cpp file):
|
||||
* #define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX2
|
||||
* #include <faiss/impl/fast_scan/dispatching.h>
|
||||
*
|
||||
* Kernel helpers come from accumulate_loops.h (search_1 multi-BB path
|
||||
* and QBS 256-bit path). The QBS helpers use pq4_kernel_qbs_256 only
|
||||
* because decompose_qbs.h pulls in 512-bit types that fail with
|
||||
* SINGLE_SIMD_LEVEL=NONE in DD mode.
|
||||
*/
|
||||
|
||||
#ifndef THE_LEVEL_TO_DISPATCH
|
||||
#error "Define THE_LEVEL_TO_DISPATCH before including this header"
|
||||
#endif
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include <faiss/impl/fast_scan/accumulate_loops.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
|
||||
namespace faiss {
|
||||
|
||||
using namespace simd_result_handlers;
|
||||
|
||||
/***************************************************************
|
||||
* ScannerMixIn: wraps a concrete handler + calls accumulation
|
||||
* kernels. Lives behind the virtual FastScanCodeScanner interface
|
||||
* so callers don't need to know the handler type.
|
||||
***************************************************************/
|
||||
|
||||
template <class Handler>
|
||||
struct ScannerMixIn : FastScanCodeScanner {
|
||||
Handler handler_;
|
||||
|
||||
template <typename... Args>
|
||||
explicit ScannerMixIn(Args&&... args)
|
||||
: handler_(std::forward<Args>(args)...) {}
|
||||
|
||||
SIMDResultHandlerToFloat* handler() override {
|
||||
return &handler_;
|
||||
}
|
||||
|
||||
void accumulate_loop(
|
||||
int nq,
|
||||
size_t nb,
|
||||
int bbs,
|
||||
int nsq,
|
||||
const uint8_t* codes,
|
||||
const uint8_t* LUT,
|
||||
int pq2x4_scale,
|
||||
size_t block_stride) override {
|
||||
if (pq2x4_scale) {
|
||||
NormTableScaler<> scaler(pq2x4_scale);
|
||||
pq4_accumulate_loop_fixed_scaler(
|
||||
nq,
|
||||
nb,
|
||||
bbs,
|
||||
nsq,
|
||||
codes,
|
||||
LUT,
|
||||
handler_,
|
||||
scaler,
|
||||
block_stride);
|
||||
} else {
|
||||
DummyScaler<> dummy;
|
||||
pq4_accumulate_loop_fixed_scaler(
|
||||
nq,
|
||||
nb,
|
||||
bbs,
|
||||
nsq,
|
||||
codes,
|
||||
LUT,
|
||||
handler_,
|
||||
dummy,
|
||||
block_stride);
|
||||
}
|
||||
}
|
||||
|
||||
void accumulate_loop_qbs(
|
||||
int qbs,
|
||||
size_t nb,
|
||||
int nsq,
|
||||
const uint8_t* codes,
|
||||
const uint8_t* LUT,
|
||||
int pq2x4_scale,
|
||||
size_t block_stride) override {
|
||||
if (pq2x4_scale) {
|
||||
NormTableScaler<> scaler(pq2x4_scale);
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_256(
|
||||
qbs, nb, nsq, codes, LUT, handler_, scaler, block_stride);
|
||||
} else {
|
||||
DummyScaler<> dummy;
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_256(
|
||||
qbs, nb, nsq, codes, LUT, handler_, dummy, block_stride);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/***************************************************************
|
||||
* Factory specialization for this SIMD level.
|
||||
*
|
||||
* Combinatorial dispatch: is_max × with_id_map × handler type
|
||||
* k == 1: SingleResultHandler
|
||||
* impl even: HeapHandler
|
||||
* impl odd: ReservoirHandler (capacity = 2*k)
|
||||
***************************************************************/
|
||||
|
||||
template <>
|
||||
std::unique_ptr<FastScanCodeScanner> make_fast_scan_scanner_impl<
|
||||
THE_LEVEL_TO_DISPATCH>(
|
||||
bool is_max,
|
||||
int impl,
|
||||
size_t nq,
|
||||
size_t ntotal,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* ids,
|
||||
const IDSelector* sel,
|
||||
bool with_id_map) {
|
||||
// Helper lambda: given comparator C and with_id_map W, select handler
|
||||
auto make = [&]<class C, bool W>() -> std::unique_ptr<FastScanCodeScanner> {
|
||||
if (k == 1) {
|
||||
using H = SingleResultHandler<C, W>;
|
||||
return std::make_unique<ScannerMixIn<H>>(
|
||||
nq, ntotal, distances, ids, sel);
|
||||
} else if (impl % 2 == 0) {
|
||||
using H = HeapHandler<C, W>;
|
||||
return std::make_unique<ScannerMixIn<H>>(
|
||||
nq, ntotal, k, distances, ids, sel);
|
||||
} else {
|
||||
using H = ReservoirHandler<C, W>;
|
||||
return std::make_unique<ScannerMixIn<H>>(
|
||||
nq, ntotal, size_t(k), size_t(2 * k), distances, ids, sel);
|
||||
}
|
||||
};
|
||||
|
||||
if (is_max) {
|
||||
if (with_id_map) {
|
||||
return make.template operator()<CMax<uint16_t, int64_t>, true>();
|
||||
} else {
|
||||
return make.template operator()<CMax<uint16_t, int>, false>();
|
||||
}
|
||||
} else {
|
||||
if (with_id_map) {
|
||||
return make.template operator()<CMin<uint16_t, int64_t>, true>();
|
||||
} else {
|
||||
return make.template operator()<CMin<uint16_t, int>, false>();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace faiss
|
||||
@@ -6,8 +6,9 @@
|
||||
*/
|
||||
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/simd_result_handlers.h>
|
||||
#include <faiss/impl/simd_dispatch.h>
|
||||
|
||||
#include <array>
|
||||
|
||||
@@ -351,3 +352,42 @@ int pq4_pack_LUT_qbs_q_map(
|
||||
}
|
||||
|
||||
} // namespace faiss
|
||||
|
||||
/***************************************************************
|
||||
* FastScanCodeScanner: NONE specialization + dispatch wrapper.
|
||||
*
|
||||
* The NONE specialization provides the scalar fallback.
|
||||
* Per-SIMD specializations (AVX2, AVX512, ARM_NEON) are in
|
||||
* impl-avx2.cpp, impl-avx512.cpp, impl-neon.cpp respectively.
|
||||
***************************************************************/
|
||||
|
||||
#define THE_LEVEL_TO_DISPATCH SIMDLevel::NONE
|
||||
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
|
||||
#undef THE_LEVEL_TO_DISPATCH
|
||||
|
||||
namespace faiss {
|
||||
|
||||
std::unique_ptr<FastScanCodeScanner> make_fast_scan_knn_scanner(
|
||||
bool is_max,
|
||||
int impl,
|
||||
size_t nq,
|
||||
size_t ntotal,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* ids,
|
||||
const IDSelector* sel,
|
||||
bool with_id_map) {
|
||||
DISPATCH_SIMDLevel(
|
||||
make_fast_scan_scanner_impl,
|
||||
is_max,
|
||||
impl,
|
||||
nq,
|
||||
ntotal,
|
||||
k,
|
||||
distances,
|
||||
ids,
|
||||
sel,
|
||||
with_id_map);
|
||||
}
|
||||
|
||||
} // namespace faiss
|
||||
@@ -9,8 +9,10 @@
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <memory>
|
||||
|
||||
#include <faiss/impl/CodePacker.h>
|
||||
#include <faiss/utils/simd_levels.h>
|
||||
|
||||
/** PQ4 SIMD packing and accumulation functions
|
||||
*
|
||||
@@ -25,6 +27,8 @@
|
||||
namespace faiss {
|
||||
|
||||
struct SIMDResultHandler;
|
||||
struct SIMDResultHandlerToFloat;
|
||||
struct IDSelector;
|
||||
|
||||
/** Pack codes for consumption by the SIMD kernels.
|
||||
* The unused bytes are set to 0.
|
||||
@@ -220,4 +224,68 @@ void accumulate_to_mem(
|
||||
const uint8_t* LUT,
|
||||
uint16_t* accu);
|
||||
|
||||
/***************************************************************
|
||||
* FastScanCodeScanner: virtual base that bundles handler + kernel
|
||||
* behind the SIMD dispatch boundary. Per-SIMD TUs instantiate this
|
||||
* with the correct SIMDLevel so that handler and kernel share the
|
||||
* same SIMD types.
|
||||
***************************************************************/
|
||||
|
||||
struct FastScanCodeScanner {
|
||||
virtual ~FastScanCodeScanner() = default;
|
||||
|
||||
/// Access the underlying result handler (for begin/end/normalizer calls)
|
||||
virtual SIMDResultHandlerToFloat* handler() = 0;
|
||||
|
||||
/// Run the search_1 accumulation loop (bbs > 32, multi-BB kernel)
|
||||
virtual void accumulate_loop(
|
||||
int nq,
|
||||
size_t nb,
|
||||
int bbs,
|
||||
int nsq,
|
||||
const uint8_t* codes,
|
||||
const uint8_t* LUT,
|
||||
int pq2x4_scale,
|
||||
size_t block_stride) = 0;
|
||||
|
||||
/// Run the QBS accumulation loop (bbs == 32)
|
||||
virtual void accumulate_loop_qbs(
|
||||
int qbs,
|
||||
size_t nb,
|
||||
int nsq,
|
||||
const uint8_t* codes,
|
||||
const uint8_t* LUT,
|
||||
int pq2x4_scale,
|
||||
size_t block_stride) = 0;
|
||||
};
|
||||
|
||||
/// Per-SIMD factory: explicitly specialized in each per-SIMD TU
|
||||
/// (impl-avx2.cpp, impl-avx512.cpp, impl-neon.cpp, fast_scan.cpp for NONE).
|
||||
/// Not called directly — use make_fast_scan_knn_scanner() instead.
|
||||
template <SIMDLevel SL>
|
||||
std::unique_ptr<FastScanCodeScanner> make_fast_scan_scanner_impl(
|
||||
bool is_max,
|
||||
int impl,
|
||||
size_t nq,
|
||||
size_t ntotal,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* ids,
|
||||
const IDSelector* sel,
|
||||
bool with_id_map);
|
||||
|
||||
/// Runtime dispatch wrapper: selects the best available SIMD level
|
||||
/// (via DISPATCH_SIMDLevel) and delegates to the corresponding
|
||||
/// make_fast_scan_scanner_impl<SL> specialization.
|
||||
std::unique_ptr<FastScanCodeScanner> make_fast_scan_knn_scanner(
|
||||
bool is_max,
|
||||
int impl,
|
||||
size_t nq,
|
||||
size_t ntotal,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* ids,
|
||||
const IDSelector* sel,
|
||||
bool with_id_map = false);
|
||||
|
||||
} // namespace faiss
|
||||
@@ -0,0 +1,13 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
#ifdef COMPILE_SIMD_AVX2
|
||||
|
||||
#define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX2
|
||||
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
|
||||
|
||||
#endif // COMPILE_SIMD_AVX2
|
||||
@@ -0,0 +1,13 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
#ifdef COMPILE_SIMD_AVX512
|
||||
|
||||
#define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX512
|
||||
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
|
||||
|
||||
#endif // COMPILE_SIMD_AVX512
|
||||
@@ -0,0 +1,39 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
#ifdef COMPILE_SIMD_ARM_NEON
|
||||
|
||||
#define THE_LEVEL_TO_DISPATCH SIMDLevel::ARM_NEON
|
||||
#include <faiss/impl/fast_scan/dispatching.h> // IWYU pragma: keep
|
||||
|
||||
// ARM_SVE: forward to ARM_NEON implementation until a dedicated SVE
|
||||
// specialization is written (same pattern as scalar_quantizer/sq-neon.cpp).
|
||||
#ifdef COMPILE_SIMD_ARM_SVE
|
||||
|
||||
namespace faiss {
|
||||
|
||||
template <>
|
||||
std::unique_ptr<FastScanCodeScanner> make_fast_scan_scanner_impl<
|
||||
SIMDLevel::ARM_SVE>(
|
||||
bool is_max,
|
||||
int impl,
|
||||
size_t nq,
|
||||
size_t ntotal,
|
||||
int64_t k,
|
||||
float* distances,
|
||||
int64_t* ids,
|
||||
const IDSelector* sel,
|
||||
bool with_id_map) {
|
||||
return make_fast_scan_scanner_impl<SIMDLevel::ARM_NEON>(
|
||||
is_max, impl, nq, ntotal, k, distances, ids, sel, with_id_map);
|
||||
}
|
||||
|
||||
} // namespace faiss
|
||||
|
||||
#endif // COMPILE_SIMD_ARM_SVE
|
||||
|
||||
#endif // COMPILE_SIMD_ARM_NEON
|
||||
@@ -5,11 +5,9 @@
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/fast_scan/LookupTableScaler.h>
|
||||
#include <faiss/impl/fast_scan/kernels_simd256.h>
|
||||
#include <faiss/impl/fast_scan/accumulate_loops.h>
|
||||
#include <faiss/impl/fast_scan/simd_result_handlers.h>
|
||||
|
||||
namespace faiss {
|
||||
@@ -22,63 +20,6 @@ using namespace simd_result_handlers;
|
||||
|
||||
namespace {
|
||||
|
||||
template <int NQ, int BB, class ResultHandler, class Scaler>
|
||||
void accumulate_fixed_blocks(
|
||||
size_t nb,
|
||||
int nsq,
|
||||
const uint8_t* codes,
|
||||
const uint8_t* LUT,
|
||||
ResultHandler& res,
|
||||
const Scaler& scaler,
|
||||
size_t block_stride) {
|
||||
constexpr int bbs = 32 * BB;
|
||||
for (size_t j0 = 0; j0 < nb; j0 += bbs) {
|
||||
FixedStorageHandler<NQ, 2 * BB> res2;
|
||||
kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2, scaler);
|
||||
res.set_block_origin(0, j0);
|
||||
res2.to_other_handler(res);
|
||||
codes += block_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <class ResultHandler, class Scaler>
|
||||
void pq4_accumulate_loop_fixed_scaler(
|
||||
int nq,
|
||||
size_t nb,
|
||||
int bbs,
|
||||
int nsq,
|
||||
const uint8_t* codes,
|
||||
const uint8_t* LUT,
|
||||
ResultHandler& res,
|
||||
const Scaler& scaler,
|
||||
size_t block_stride) {
|
||||
FAISS_THROW_IF_NOT(is_aligned_pointer(codes));
|
||||
FAISS_THROW_IF_NOT(is_aligned_pointer(LUT));
|
||||
FAISS_THROW_IF_NOT(bbs % 32 == 0);
|
||||
FAISS_THROW_IF_NOT(nb % bbs == 0);
|
||||
|
||||
#define DISPATCH(NQ, BB) \
|
||||
case NQ * 1000 + BB: \
|
||||
accumulate_fixed_blocks<NQ, BB>( \
|
||||
nb, nsq, codes, LUT, res, scaler, block_stride); \
|
||||
break
|
||||
|
||||
switch (nq * 1000 + bbs / 32) {
|
||||
DISPATCH(1, 1);
|
||||
DISPATCH(1, 2);
|
||||
DISPATCH(1, 3);
|
||||
DISPATCH(1, 4);
|
||||
DISPATCH(1, 5);
|
||||
DISPATCH(2, 1);
|
||||
DISPATCH(2, 2);
|
||||
DISPATCH(3, 1);
|
||||
DISPATCH(4, 1);
|
||||
default:
|
||||
FAISS_THROW_FMT("nq=%d bbs=%d not instantiated", nq, bbs);
|
||||
}
|
||||
#undef DISPATCH
|
||||
}
|
||||
|
||||
template <class ResultHandler>
|
||||
void pq4_accumulate_loop_fixed_handler(
|
||||
int nq,
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/fast_scan/LookupTableScaler.h>
|
||||
|
||||
@@ -13,7 +13,7 @@
|
||||
|
||||
#include <faiss/IndexPQFastScan.h>
|
||||
#include <faiss/impl/ProductQuantizer.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
|
||||
namespace {
|
||||
|
||||
|
||||
@@ -12,7 +12,7 @@
|
||||
#include <vector>
|
||||
|
||||
#include <faiss/IndexIVF.h>
|
||||
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/utils/AlignedTable.h>
|
||||
#include <faiss/utils/random.h>
|
||||
|
||||
|
||||
Reference in New Issue
Block a user